Fix(BA-781): Let Auth API handler check all keypairs owned by user

changes/3780.fix.md

@@ -0,0 +1 @@
+Let Auth API handler check all keypairs owned by user

docs/manager/rest-reference/openapi.json

@@ -3,7 +3,7 @@
   "info": {
     "title": "Backend.AI Manager API",
     "description": "Backend.AI Manager REST API specification",
-    "version": "25.2.0",
+    "version": "25.3.0",
     "contact": {
       "name": "Lablup Inc.",
       "url": "https://docs.backend.ai",

src/ai/backend/manager/api/auth.py

@@ -32,6 +32,7 @@
 from ..models.user import (
     INACTIVE_USER_STATUSES,
     UserRole,
+    UserRow,
     UserStatus,
     check_credential,
     compare_to_hashed_password,
@@ -727,33 +728,26 @@ async def authorize(request: web.Request, params: Any) -> web.Response:
         raise AuthorizationFailed("This account needs email verification.")
     if user["status"] in INACTIVE_USER_STATUSES:
         raise AuthorizationFailed("User credential mismatch.")
-    async with root_ctx.db.begin() as conn:
-        query = (
-            sa.select([keypairs.c.access_key, keypairs.c.secret_key])
-            .select_from(keypairs)
-            .where(
-                (keypairs.c.user == user["uuid"]) & (keypairs.c.is_active),
-            )
-            .order_by(sa.desc(keypairs.c.is_admin))
-        )
-        result = await conn.execute(query)
-        keypair = result.first()
-    if keypair is None:
-        raise AuthorizationFailed("No API keypairs found.")
     await check_password_age(root_ctx.db, user, root_ctx.shared_config["auth"])
+    async with root_ctx.db.begin_session() as db_session:
+        user_row = await UserRow.query_user_by_uuid(user["uuid"], db_session)
+        assert user_row is not None
+        main_keypair_row = user_row.get_main_keypair_row()
+    if main_keypair_row is None:
+        raise AuthorizationFailed("No API keypairs found.")
     # [Hooking point for POST_AUTHORIZE]
     # The hook handlers should accept a tuple of the request, user, and keypair objects.
     hook_result = await root_ctx.hook_plugin_ctx.dispatch(
         "POST_AUTHORIZE",
-        (request, params, user, keypair),
+        (request, params, user, main_keypair_row.mapping),
         return_when=FIRST_COMPLETED,
     )
     if hook_result.status != PASSED:
         raise RejectedByHook.from_hook_result(hook_result)
     return web.json_response({
         "data": {
-            "access_key": keypair["access_key"],
-            "secret_key": keypair["secret_key"],
+            "access_key": main_keypair_row.access_key,
+            "secret_key": main_keypair_row.secret_key,
             "role": user["role"],
             "status": user["status"],
         },

src/ai/backend/manager/models/keypair.py

@@ -114,6 +114,27 @@ class KeyPairRow(Base):
 
     user_row = relationship("UserRow", back_populates="keypairs", foreign_keys=keypairs.c.user)
 
+    @property
+    def mapping(self) -> dict[str, Any]:
+        return {
+            "user_id": self.user_id,
+            "access_key": self.access_key,
+            "secret_key": self.secret_key,
+            "is_active": self.is_active,
+            "is_admin": self.is_admin,
+            "created_at": self.created_at,
+            "modified_at": self.modified_at,
+            "last_used": self.last_used,
+            "rate_limit": self.rate_limit,
+            "num_queries": self.num_queries,
+            "ssh_public_key": self.ssh_public_key,
+            "ssh_private_key": self.ssh_private_key,
+            "user": self.user,
+            "resource_policy": self.resource_policy,
+            "dotfiles": self.dotfiles,
+            "bootstrap_script": self.bootstrap_script,
+        }
+
 
 class UserInfo(graphene.ObjectType):
     email = graphene.String()

src/ai/backend/manager/models/user.py

@@ -2,7 +2,7 @@
 
 import enum
 import logging
-from typing import TYPE_CHECKING, Any, Dict, Iterable, Mapping, Optional, Sequence, cast
+from typing import TYPE_CHECKING, Any, Dict, Iterable, Mapping, Optional, Self, Sequence, cast
 from uuid import UUID, uuid4
 
 import aiotools
@@ -18,7 +18,7 @@
 from sqlalchemy.ext.asyncio import AsyncConnection as SAConnection
 from sqlalchemy.ext.asyncio import AsyncEngine as SAEngine
 from sqlalchemy.ext.asyncio import AsyncSession as SASession
-from sqlalchemy.orm import joinedload, load_only, noload, relationship
+from sqlalchemy.orm import joinedload, load_only, noload, relationship, selectinload
 from sqlalchemy.sql.expression import bindparam
 from sqlalchemy.types import VARCHAR, TypeDecorator
 
@@ -48,6 +48,7 @@
 
 if TYPE_CHECKING:
     from .gql import GraphQueryContext
+    from .keypair import KeyPairRow
     from .storage import StorageSessionManager
 
 log = BraceStyleAdapter(logging.getLogger(__spec__.name))
@@ -191,6 +192,41 @@ class UserRow(Base):
         primaryjoin="UserRow.uuid == foreign(VFolderRow.user)",
     )
 
+    @classmethod
+    async def query_user_by_uuid(
+        cls,
+        user_uuid: UUID,
+        db_session: SASession,
+    ) -> Optional[Self]:
+        user_query = (
+            sa.select(UserRow)
+            .where(UserRow.uuid == user_uuid)
+            .options(
+                joinedload(UserRow.main_keypair),
+                selectinload(UserRow.keypairs),
+            )
+        )
+        user_row = await db_session.scalar(user_query)
+        return user_row
+
+    def get_main_keypair_row(self) -> Optional[KeyPairRow]:
+        # `cast()` requires import of KeyPairRow
+        from .keypair import KeyPairRow
+
+        keypair_candidate: Optional[KeyPairRow] = None
+        main_keypair_row = cast(Optional[KeyPairRow], self.main_keypair)
+        if main_keypair_row is None:
+            keypair_rows = cast(list[KeyPairRow], self.keypairs)
+            active_keypairs = [row for row in keypair_rows if row.is_active]
+            for row in active_keypairs:
+                if keypair_candidate is None or not keypair_candidate.is_admin:
+                    keypair_candidate = row
+            if keypair_candidate is not None:
+                self.main_keypair = keypair_candidate
+        else:
+            keypair_candidate = main_keypair_row
+        return keypair_candidate
+
 
 class UserGroup(graphene.ObjectType):
     id = graphene.UUID()