feat: add AUTHORIZE hook to override credential check during login pr…

…ocess (#339)

changes/339.feature

@@ -0,0 +1 @@
+Add hooking point for AUTHORIZE with FIRST_COMPLETED requirement.

src/ai/backend/gateway/auth.py

@@ -23,6 +23,7 @@
 from ai.backend.common.logging import Logger, BraceStyleAdapter
 from ai.backend.common.plugin.hook import (
     ALL_COMPLETED,
+    FIRST_COMPLETED,
     PASSED,
 )
 from .exceptions import (
@@ -537,9 +538,31 @@ async def authorize(request: web.Request, params: Any) -> web.Response:
         raise InvalidAPIParameters('Unsupported authorization type')
     log.info('AUTH.AUTHORIZE(d:{0[domain]}, u:{0[username]}, passwd:****, type:{0[type]})', params)
     dbpool = request.app['dbpool']
-    user = await check_credential(
-        dbpool,
-        params['domain'], params['username'], params['password'])
+
+    # [Hooking point for AUTHORIZE with the FIRST_COMPLETED requirement]
+    # The hook handlers should accept the whole ``params`` dict, and optional
+    # ``dbpool`` parameter (if the hook needs to query to database).
+    # They should return a corresponding Backend.AI user object after performing
+    # their own authentication steps, like LDAP authentication, etc.
+    params['dbpool'] = dbpool  # hack to execute db commands in the hook (TODO: more general logic)
+    hook_result = await request.app['hook_plugin_ctx'].dispatch(
+        'AUTHORIZE',
+        (params, dbpool),
+        return_when=FIRST_COMPLETED,
+    )
+    if hook_result.status != PASSED:
+        # Did not pass AUTHORIZED hook
+        reason = hook_result.reason
+        raise RejectedByHook(extra_msg=reason)
+    elif 'user' in ChainMap(*hook_result.result):
+        # Passed one of AUTHORIZED hook
+        user = ChainMap(*hook_result.result)['user']
+    else:
+        # No AUTHORIZE hook is defined (proceed with normal login)
+        user = await check_credential(
+            dbpool,
+            params['domain'], params['username'], params['password']
+        )
     if user is None:
         raise AuthorizationFailed('User credential mismatch.')
     if user.get('status') == UserStatus.BEFORE_VERIFICATION: