chore: token policies persisted in DB

diracx-client/tests/test_auth.py

@@ -17,7 +17,6 @@
     "jti": "f0706e0a-af1e-4538-9f1f-7b9620783cba",
     "exp": int((datetime.now(tz=timezone.utc) + timedelta(days=1)).timestamp()),
     "legacy_exchange": False,
-    "dirac_policies": {},
 }
 
 TOKEN_RESPONSE_DICT = {

diracx-core/src/diracx/core/models/auth.py

@@ -59,7 +59,6 @@ class OpenIDConfiguration(TypedDict):
 class TokenPayload(BaseModel):
     jti: str
     exp: UTCDatetime
-    dirac_policies: dict
 
 
 class TokenResponse(BaseModel):
@@ -77,10 +76,13 @@ class AccessTokenPayload(TokenPayload):
     dirac_properties: list[str]
     preferred_username: str
     dirac_group: str
+    dirac_policies: dict
 
 
 class RefreshTokenPayload(TokenPayload):
     legacy_exchange: bool
+    # TODO: this should be removed later (https://github.com/DIRACGrid/diracx/issues/524#issuecomment-4395561176)
+    dirac_policies: dict | None = None
 
 
 class SupportInfo(TypedDict):

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

@@ -5,6 +5,7 @@
 import secrets
 from datetime import UTC, datetime
 from itertools import pairwise
+from typing import Any
 
 from dateutil.relativedelta import relativedelta
 from dateutil.rrule import MONTHLY, rrule
@@ -324,20 +325,15 @@ async def update_authorization_flow_status(
         )
 
     async def insert_refresh_token(
-        self,
-        jti: UUID,
-        subject: str,
-        scope: str,
+        self, jti: UUID, subject: str, scope: str, policies: dict[str, Any]
     ) -> None:
         """Insert a refresh token in the DB.
 
         As well as user attributes required to generate access tokens.
         """
         # Insert values into the DB
         stmt = insert(RefreshTokens).values(
-            jti=str(jti),
-            sub=subject,
-            scope=scope,
+            jti=str(jti), sub=subject, scope=scope, policies=policies
         )
         await self.conn.execute(stmt)
 

diracx-db/src/diracx/db/sql/auth/schema.py

@@ -119,5 +119,8 @@ class RefreshTokens(Base):
 
     # User attributes bound to the refresh token
     sub: Mapped[str] = mapped_column("Sub", String(256), index=True)
+    policies: Mapped[dict[str, Any] | None] = mapped_column(
+        "Policies", nullable=True, default=None
+    )
 
     __table_args__ = (Index("index_status_sub", status, sub),)

diracx-db/tests/auth/test_refresh_token.py

@@ -25,6 +25,7 @@ async def test_insert(auth_db: AuthDB):
             jti1,
             "subject",
             "vo:lhcb property:NormalUser",
+            {"PolicySpecific": "OpenAccessForTest"},
         )
 
     # Insert a second refresh token
@@ -34,6 +35,7 @@ async def test_insert(auth_db: AuthDB):
             jti2,
             "subject",
             "vo:lhcb property:NormalUser",
+            {"PolicySpecific": "OpenAccessForTest"},
         )
 
     # Make sure they don't have the same JWT ID
@@ -46,6 +48,7 @@ async def test_get(auth_db: AuthDB):
     refresh_token_details = {
         "sub": "12345",
         "scope": "vo:lhcb property:NormalUser",
+        "policies": {"PolicySpecific": "OpenAccessForTest"},
     }
 
     # Insert refresh token details
@@ -55,6 +58,7 @@ async def test_get(auth_db: AuthDB):
             jti,
             refresh_token_details["sub"],
             refresh_token_details["scope"],
+            refresh_token_details["policies"],
         )
 
     # Enrich the dict with the generated refresh token attributes
@@ -63,6 +67,7 @@ async def test_get(auth_db: AuthDB):
         "Scope": refresh_token_details["scope"],
         "JTI": jti,
         "Status": RefreshTokenStatus.CREATED,
+        "Policies": refresh_token_details["policies"],
     }
 
     # Get refresh token details
@@ -90,9 +95,7 @@ async def test_get_user_refresh_tokens(auth_db: AuthDB):
     async with auth_db as auth_db:
         for sub in subjects:
             await auth_db.insert_refresh_token(
-                uuid7(),
-                sub,
-                "scope",
+                uuid7(), sub, "scope", {"PolicySpecific": "OpenAccessForTest"}
             )
 
     # Get the refresh tokens of each user
@@ -117,9 +120,7 @@ async def test_revoke(auth_db: AuthDB):
     async with auth_db as auth_db:
         jti = uuid7()
         await auth_db.insert_refresh_token(
-            jti,
-            "subject",
-            "scope",
+            jti, "subject", "scope", {"PolicySpecific": "OpenAccessForTest"}
         )
 
     # Revoke the token
@@ -146,9 +147,7 @@ async def test_revoke_user_refresh_tokens(auth_db: AuthDB):
     async with auth_db as auth_db:
         for sub in subjects:
             await auth_db.insert_refresh_token(
-                uuid7(),
-                sub,
-                "scope",
+                uuid7(), sub, "scope", {"PolicySpecific": "OpenAccessForTest"}
             )
 
     # Revoke the tokens of sub1
@@ -191,9 +190,7 @@ async def test_revoke_and_get_user_refresh_tokens(auth_db: AuthDB):
         for _ in range(nb_tokens):
             jti = uuid7()
             await auth_db.insert_refresh_token(
-                jti,
-                sub,
-                "scope",
+                jti, sub, "scope", {"PolicySpecific": "OpenAccessForTest"}
             )
             jtis.append(jti)
 
@@ -239,9 +236,7 @@ async def test_get_refresh_tokens(auth_db: AuthDB):
     async with auth_db as auth_db:
         for sub in subjects:
             await auth_db.insert_refresh_token(
-                uuid7(),
-                sub,
-                "scope",
+                uuid7(), sub, "scope", {"PolicySpecific": "OpenAccessForTest"}
             )
 
     # Get all refresh tokens (Admin)

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

@@ -6,7 +6,7 @@
 import hashlib
 import re
 from datetime import datetime, timedelta, timezone
-from typing import cast
+from typing import Any, cast
 
 from joserfc import jwt
 from joserfc.jwt import Claims
@@ -45,6 +45,7 @@ async def get_oidc_token(
     config: Config,
     settings: AuthSettings,
     available_properties: set[SecurityProperty],
+    all_access_policies: dict[str, Any],
     device_code: str | None = None,
     code: str | None = None,
     redirect_uri: str | None = None,
@@ -55,6 +56,7 @@ async def get_oidc_token(
     legacy_exchange = False
     include_refresh_token = True
     refresh_token_expire_minutes = None
+    policies = None
 
     if grant_type == GrantType.device_code:
         assert device_code is not None
@@ -77,6 +79,7 @@ async def get_oidc_token(
             legacy_exchange,
             refresh_token_expire_minutes,
             include_refresh_token,
+            policies,
         ) = await get_oidc_token_info_from_refresh_flow(
             refresh_token, auth_db, settings
         )
@@ -91,9 +94,11 @@ async def get_oidc_token(
         config,
         settings,
         available_properties,
+        all_access_policies,
         legacy_exchange=legacy_exchange,
         refresh_token_expire_minutes=refresh_token_expire_minutes,
         include_refresh_token=include_refresh_token,
+        policies=policies,
     )
 
 
@@ -163,7 +168,7 @@ async def get_oidc_token_info_from_authorization_flow(
 
 async def get_oidc_token_info_from_refresh_flow(
     refresh_token: str, auth_db: AuthDB, settings: AuthSettings
-) -> tuple[dict, str, bool, float, bool]:
+) -> tuple[dict, str, bool, float, bool, dict]:
     """Get OIDC token information from the refresh token DB and check few parameters before returning it."""
     # Decode the refresh token to get the JWT ID
     jti, exp, legacy_exchange = await verify_dirac_refresh_token(
@@ -216,12 +221,14 @@ async def get_oidc_token_info_from_refresh_flow(
         "sub": sub.split(":", 1)[1],
     }
     scope = refresh_token_attributes["Scope"]
+    policies = refresh_token_attributes["Policies"]
     return (
         oidc_token_info,
         scope,
         legacy_exchange,
         remaining_minutes,
         include_refresh_token,
+        policies,
     )
 
 
@@ -240,6 +247,7 @@ async def perform_legacy_exchange(
     available_properties: set[SecurityProperty],
     settings: AuthSettings,
     config: Config,
+    all_access_policies: dict[str, Any],
     expires_minutes: float | None = None,
 ) -> tuple[AccessTokenPayload, RefreshTokenPayload | None]:
     """Endpoint used by legacy DIRAC to mint tokens for proxy -> token exchange."""
@@ -265,8 +273,10 @@ async def perform_legacy_exchange(
         config,
         settings,
         available_properties,
+        all_access_policies,
         refresh_token_expire_minutes=expires_minutes,
         legacy_exchange=True,
+        policies=None,
     )
 
 
@@ -277,10 +287,12 @@ async def exchange_token(
     config: Config,
     settings: AuthSettings,
     available_properties: set[SecurityProperty],
+    all_access_policies: dict[str, Any],
     *,
     refresh_token_expire_minutes: float | None = None,
     legacy_exchange: bool = False,
     include_refresh_token: bool = True,
+    policies: dict[str, Any] | None = None,
 ) -> tuple[AccessTokenPayload, RefreshTokenPayload | None]:
     """Exchange the OIDC token for a DIRAC generated access token."""
     # Extract dirac attributes from the OIDC scope
@@ -316,6 +328,33 @@ async def exchange_token(
     # Merge the VO with the subject to get a unique DIRAC sub
     sub = f"{vo}:{sub}"
 
+    # Generate access token payload
+    # For now, the access token is only used to access DIRAC services,
+    # therefore, the audience is not set and checked
+    access_jti = uuid7()
+    access_exp = uuid7_to_datetime(access_jti) + timedelta(
+        minutes=settings.access_token_expire_minutes
+    )
+    access_payload = AccessTokenPayload(
+        sub=sub,
+        vo=vo,
+        iss=settings.token_issuer,
+        dirac_properties=list(properties),
+        jti=str(access_jti),
+        preferred_username=preferred_username,
+        dirac_group=dirac_group,
+        exp=access_exp,
+        dirac_policies={},
+    )
+
+    # Enrich the token with policy specific content
+    if policies is not None:
+        access_payload.dirac_policies = policies
+    else:
+        for policy_name, policy in all_access_policies.items():
+            if access_extra := policy.enrich_tokens(access_payload):
+                access_payload.dirac_policies[policy_name] = access_extra
+
     refresh_payload: RefreshTokenPayload | None = None
     if include_refresh_token:
         # Insert the refresh token with user details into the RefreshTokens table
@@ -324,6 +363,7 @@ async def exchange_token(
             auth_db=auth_db,
             subject=sub,
             scope=scope,
+            policies=access_payload.dirac_policies,
         )
 
         # Generate refresh token payload
@@ -338,28 +378,8 @@ async def exchange_token(
             # legacy_exchange is used to indicate that the original refresh token
             # was obtained from the legacy_exchange endpoint
             legacy_exchange=legacy_exchange,
-            dirac_policies={},
         )
 
-    # Generate access token payload
-    # For now, the access token is only used to access DIRAC services,
-    # therefore, the audience is not set and checked
-    access_jti = uuid7()
-    access_exp = uuid7_to_datetime(access_jti) + timedelta(
-        minutes=settings.access_token_expire_minutes
-    )
-    access_payload = AccessTokenPayload(
-        sub=sub,
-        vo=vo,
-        iss=settings.token_issuer,
-        dirac_properties=list(properties),
-        jti=str(access_jti),
-        preferred_username=preferred_username,
-        dirac_group=dirac_group,
-        exp=access_exp,
-        dirac_policies={},
-    )
-
     return access_payload, refresh_payload
 
 
@@ -391,19 +411,15 @@ def _sign_token_payload(claims: dict, settings: AuthSettings) -> str:
 
 
 async def insert_refresh_token(
-    auth_db: AuthDB,
-    subject: str,
-    scope: str,
+    auth_db: AuthDB, subject: str, scope: str, policies: dict[str, Any]
 ) -> UUID:
     """Insert a refresh token into the database and return the JWT ID."""
     # Generate a JWT ID
     jti = uuid7()
 
     # Insert the refresh token into the DB
     await auth_db.insert_refresh_token(
-        jti=jti,
-        subject=subject,
-        scope=scope,
+        jti=jti, subject=subject, scope=scope, policies=policies
     )
     return jti
 

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

@@ -33,7 +33,6 @@
 from diracx.core.extensions import DiracEntryPoint, select_from_extension
 from diracx.core.models import (
     AccessTokenPayload,
-    RefreshTokenPayload,
 )
 from diracx.core.settings import DevelopmentSettings
 from diracx.routers.dependencies import auto_inject
@@ -90,18 +89,15 @@ async def policy(policy_name: str, user_info: AuthorizedUserInfo, /):
         return
 
     @staticmethod
-    def enrich_tokens(
-        access_payload: AccessTokenPayload, refresh_payload: RefreshTokenPayload | None
-    ) -> tuple[dict, dict]:
-        """Add content to access or refresh payload when issuing a token.
+    def enrich_tokens(access_payload: AccessTokenPayload) -> dict:
+        """Add content to access payload when issuing a token.
 
-        Content can be whatever is desired inside the access or refresh payload.
+        Content can be whatever is desired inside the access payload.
 
         :param access_payload: access token payload
-        :param refresh_payload: refresh token payload
-        :returns: extra content for both payload
+        :returns: extra content for the access payload
         """
-        return {}, {}
+        return {}
 
 
 @auto_inject