feat: Better handling of refresh tokens for pilots

Author: Robin-Van-de-Merghel

diracx-db/src/diracx/db/sql/pilot_agents/db.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
 from datetime import datetime, timezone
+from typing import Sequence
 
-from sqlalchemy import DateTime, insert, select, update
+from sqlalchemy import insert, select, update
 from sqlalchemy.exc import IntegrityError, NoResultFound
 
 from diracx.core.exceptions import (
@@ -26,7 +27,7 @@ async def add_pilot_references(
         vo: str,
         grid_type: str = "DIRAC",
         pilot_stamps: dict | None = None,
-    ) -> None:
+    ) -> Sequence:  # Return a list of primary keys
 
         if pilot_stamps is None:
             pilot_stamps = {}
@@ -47,10 +48,18 @@ async def add_pilot_references(
             for ref in pilot_ref
         ]
 
-        # Insert multiple rows in a single execute call
-        stmt = insert(PilotAgents).values(values)
-        await self.conn.execute(stmt)
-        return
+        # Insert multiple rows in a single execute call and use 'returning' to get primary keys
+        stmt = (
+            insert(PilotAgents).values(values).returning(PilotAgents.pilot_id)
+        )  # Assuming 'id' is the primary key
+        result = await self.conn.execute(stmt)
+
+        # Use .scalars() and .all() to get the primary keys directly in a list
+        primary_keys = (
+            result.scalars().all()
+        )  # This returns a flat list of primary keys
+
+        return primary_keys
 
     async def increment_pilot_secret_use(
         self,
@@ -100,32 +109,6 @@ async def verify_pilot_secret(
         # Increment the count
         await self.increment_pilot_secret_use(pilot_id=pilot_id)
 
-    async def register_new_pilot(
-        self,
-        vo: str,
-        pilot_job_reference: str,
-        pilot_stamp: str,
-        grid_type: str = "DIRAC",
-        submission_time: DateTime | None = None,  # ?
-        last_update_time: DateTime | None = None,  # = now?
-    ) -> int | None:
-        stmt = insert(PilotAgents).values(
-            vo=vo,
-            submission_time=submission_time,
-            last_update_time=last_update_time,
-            pilot_job_reference=pilot_job_reference,
-            grid_type=grid_type,
-            pilot_stamp=pilot_stamp,
-        )
-
-        # Execute the request
-        res = await self.conn.execute(stmt)
-
-        new_pilot_id = res.inserted_primary_key
-
-        # Returns the new pilot ID
-        return int(new_pilot_id[0]) if new_pilot_id else None
-
     async def add_pilot_credentials(self, pilot_id: int, pilot_hashed_secret: str):
 
         stmt = insert(PilotRegistrations).values(

diracx-db/tests/pilot_agents/test_pilot_agents_db.py

@@ -38,10 +38,9 @@ async def test_insert_and_select_single(pilot_agents_db: PilotAgentsDB):
 
     async with pilot_agents_db as pilot_agents_db:
         pilot_reference = "pilot-reference-test"
-        await pilot_agents_db.register_new_pilot(
-            vo="pilot-vo",
-            pilot_job_reference=pilot_reference,
-            pilot_stamp="pilot-stamp",
+        await pilot_agents_db.add_pilot_references(
+            vo="lhcb",
+            pilot_ref=[pilot_reference],
             grid_type="grid-type",
         )
 
@@ -51,26 +50,26 @@ async def test_insert_and_select_single(pilot_agents_db: PilotAgentsDB):
             await pilot_agents_db.get_pilot_by_reference("I am a fake ref")
 
         # Set values
-        assert res["VO"] == "pilot-vo"
+        assert res["VO"] == "lhcb"
         assert res["PilotJobReference"] == pilot_reference
-        assert res["PilotStamp"] == "pilot-stamp"
         assert res["GridType"] == "grid-type"
 
-        # Default values
-        assert res["BenchMark"] == 0.0
-        assert res["Status"] == "Unknown"
-
 
 async def test_create_pilot_and_verify_secret(pilot_agents_db: PilotAgentsDB):
 
     async with pilot_agents_db as pilot_agents_db:
         pilot_reference = "pilot-reference-test"
-        pilot_id = await pilot_agents_db.register_new_pilot(
-            vo="pilot-vo",
-            pilot_job_reference=pilot_reference,
-            pilot_stamp="pilot-stamp",
+        pilot_ids = await pilot_agents_db.add_pilot_references(
+            vo="lhcb",
+            pilot_ref=[pilot_reference],
+            grid_type="grid-type",
         )
 
+        assert len(pilot_ids) == 1
+
+        # Only one element
+        pilot_id = pilot_ids[0]
+
         secret = "AW0nd3rfulS3cr3t"
         pilot_hashed_secret = hash(secret)
 

diracx-logic/src/diracx/logic/auth/token.py

@@ -73,9 +73,7 @@ async def get_oidc_token(
             oidc_token_info,
             scope,
             legacy_exchange,
-        ) = await get_oidc_token_info_from_refresh_flow(
-            refresh_token, auth_db, settings
-        )
+        ) = await get_token_info_from_refresh_flow(refresh_token, auth_db, settings)
     else:
         raise NotImplementedError(f"Grant type not implemented {grant_type}")
 
@@ -155,7 +153,7 @@ async def get_oidc_token_info_from_authorization_flow(
     return (oidc_token_info, scope)
 
 
-async def get_oidc_token_info_from_refresh_flow(
+async def get_token_info_from_refresh_flow(
     refresh_token: str, auth_db: AuthDB, settings: AuthSettings
 ) -> tuple[dict, str, bool]:
     """Get OIDC token information from the refresh token DB and check few parameters before returning it."""
@@ -310,7 +308,7 @@ async def exchange_token(
         )
 
     else:
-        preferred_username = oidc_token_info["pilot_reference"]
+        preferred_username = oidc_token_info["preferred_username"]
 
     # Merge the VO with the subject to get a unique DIRAC sub
     sub = f"{vo}:{sub}"
@@ -361,20 +359,32 @@ async def generate_pilot_tokens(
     config: Config,
     settings: AuthSettings,
     available_properties: set[SecurityProperty],
+    refresh_token: str | None = None,
 ) -> tuple[AccessTokenPayload, RefreshTokenPayload]:
 
-    pilot = await get_pilot_informations_by_reference(
-        pilot_db=pilot_db, pilot_job_reference=pilot_job_reference
-    )
+    scope = None
+    pilot_info = None
 
-    pilot_info = {
-        "pilot_reference": pilot["PilotJobReference"],
-        "sub": pilot["PilotJobReference"],
-    }
+    if refresh_token is not None:
+        pilot_info, scope, _ = await get_token_info_from_refresh_flow(
+            refresh_token=refresh_token, auth_db=auth_db, settings=settings
+        )
+    else:
+
+        pilot = await get_pilot_informations_by_reference(
+            pilot_db=pilot_db, pilot_job_reference=pilot_job_reference
+        )
+
+        pilot_info = {
+            "preferred_username": pilot["PilotJobReference"],
+            "sub": pilot["PilotJobReference"],
+        }
+
+        scope = generate_pilot_scope(pilot)
 
     access_token, refresh_token = await exchange_token(
         auth_db=auth_db,
-        scope=generate_pilot_scope(pilot),
+        scope=scope,
         oidc_token_info=pilot_info,
         config=config,
         settings=settings,

diracx-routers/src/diracx/routers/auth/pilot.py

@@ -1,10 +1,21 @@
 from __future__ import annotations
 
-from fastapi import HTTPException, status
+from typing import Annotated
 
-from diracx.core.exceptions import AuthorizationError, PilotNotFoundError
+from fastapi import (
+    Depends,
+    HTTPException,
+    status,
+)
+
+from diracx.core.exceptions import (
+    AuthorizationError,
+    InvalidCredentialsError,
+    PilotNotFoundError,
+)
 from diracx.logic.auth.pilot import try_login
 from diracx.logic.auth.token import create_token, generate_pilot_tokens
+from diracx.routers.pilots.access_policies import RegisteredPilotAccessPolicyCallable
 
 from ..dependencies import (
     AuthDB,
@@ -14,6 +25,7 @@
     PilotAgentsDB,
 )
 from ..fastapi_classes import DiracxRouter
+from ..utils.users import AuthorizedUserInfo, verify_dirac_access_token
 
 router = DiracxRouter(require_auth=False)
 
@@ -63,3 +75,42 @@ async def pilot_login(
         "access_token": create_token(access_token, settings),
         "refresh_token": create_token(refresh_token, settings),
     }
+
+
+@router.post("/pilot-refresh-token")
+async def refresh_pilot_tokens(
+    pilot_db: PilotAgentsDB,
+    auth_db: AuthDB,
+    config: Config,
+    settings: AuthSettings,
+    available_properties: AvailableSecurityProperties,
+    check_permissions: RegisteredPilotAccessPolicyCallable,
+    refresh_token: str,
+    pilot_info: Annotated[AuthorizedUserInfo, Depends(verify_dirac_access_token)],
+):
+    """Endpoint where a pilot can exchange a refresh token against a token."""
+    await check_permissions()
+
+    try:
+        new_access_token, new_refresh_token = await generate_pilot_tokens(
+            pilot_db=pilot_db,
+            auth_db=auth_db,
+            pilot_job_reference=pilot_info.preferred_username,
+            config=config,
+            settings=settings,
+            available_properties=available_properties,
+            refresh_token=refresh_token,
+        )
+    except InvalidCredentialsError as e:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED, detail=str(e)
+        ) from e
+    except ValueError as e:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)
+        ) from e
+
+    return {
+        "access_token": create_token(new_access_token, settings),
+        "refresh_token": create_token(new_refresh_token, settings),
+    }

diracx-routers/src/diracx/routers/pilots/access_policies.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 from collections.abc import Callable
-from enum import StrEnum, auto
 from typing import Annotated
 
 from fastapi import Depends, HTTPException, status
@@ -13,18 +12,6 @@
 from diracx.routers.utils.users import AuthorizedUserInfo
 
 
-class ActionType(StrEnum):
-    #: Create a job or a sandbox
-    CREATE = auto()
-    #: Check job status, download a sandbox
-    READ = auto()
-    #: delete, kill, remove, set status, etc of a job
-    #: delete or assign a sandbox
-    MANAGE = auto()
-    #: Search
-    QUERY = auto()
-
-
 class RegisteredPilotAccessPolicy(BaseAccessPolicy):
 
     @staticmethod

diracx-routers/src/diracx/routers/pilots/debug.py

@@ -16,8 +16,8 @@
 @router.get("/info")
 async def get_pilot_info(
     check_permissions: RegisteredPilotAccessPolicyCallable,
-    user_info: Annotated[AuthorizedUserInfo, Depends(verify_dirac_access_token)],
+    pilot_info: Annotated[AuthorizedUserInfo, Depends(verify_dirac_access_token)],
 ):
     await check_permissions()
 
-    return user_info
+    return pilot_info

diracx-routers/tests/auth/test_pilot_auth.py

@@ -43,10 +43,17 @@ async def test_create_pilot_and_verify_secret(test_client):
 
     async with db as pilot_agents_db:
         # Register a pilot
-        pilot_id = await pilot_agents_db.register_new_pilot(
-            vo=pilot_vo, pilot_job_reference=pilot_reference, pilot_stamp="pilot-stamp"
+        pilot_ids = await pilot_agents_db.add_pilot_references(
+            vo=pilot_vo,
+            pilot_ref=[pilot_reference],
+            grid_type="grid-type",
         )
 
+        assert len(pilot_ids) == 1
+
+        # Only one element
+        pilot_id = pilot_ids[0]
+
         # Add credentials to this pilot
         await pilot_agents_db.add_pilot_credentials(
             pilot_id=pilot_id, pilot_hashed_secret=pilot_hashed_secret
@@ -116,3 +123,45 @@ async def test_create_pilot_and_verify_secret(test_client):
 
     assert r.status_code == 401
     assert r.json()["detail"] == "bad pilot_id / pilot_secret"
+
+    # ----------------- Exchange for new tokens -----------------
+    request_data = {"refresh_token": refresh_token}
+    r = test_client.post(
+        "/api/auth/pilot-refresh-token",
+        params=request_data,
+        headers={"Authorization": f"Bearer {access_token}"},
+    )
+
+    assert r.status_code == 200
+
+    new_access_token = r.json()["access_token"]
+    new_refresh_token = r.json()["refresh_token"]
+
+    # ----------------- Get info with new token -----------------
+    r = test_client.get(
+        "/api/pilots/info", headers={"Authorization": f"Bearer {new_access_token}"}
+    )
+
+    assert r.status_code == 200
+
+    # ----------------- Exchange token with old token -----------------
+    request_data = {"refresh_token": refresh_token}
+    r = test_client.post(
+        "/api/auth/pilot-refresh-token",
+        params=request_data,
+        headers={"Authorization": f"Bearer {access_token}"},
+    )
+
+    assert r.status_code == 401, r.json()
+
+    # ----------------- Exchange token with new token -----------------
+    request_data = {"refresh_token": new_refresh_token}
+    r = test_client.post(
+        "/api/auth/pilot-refresh-token",
+        params=request_data,
+        headers={"Authorization": f"Bearer {new_access_token}"},
+    )
+
+    # RFC6749
+    # https://datatracker.ietf.org/doc/html/rfc6749#section-10.4
+    assert r.status_code == 401, r.json()