diff --git a/diracx-client/tests/test_auth.py b/diracx-client/tests/test_auth.py index 367295425..2b705450c 100644 --- a/diracx-client/tests/test_auth.py +++ b/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 = { diff --git a/diracx-core/src/diracx/core/models/auth.py b/diracx-core/src/diracx/core/models/auth.py index 7dcc220b4..b33ca1338 100644 --- a/diracx-core/src/diracx/core/models/auth.py +++ b/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): diff --git a/diracx-db/src/diracx/db/sql/auth/db.py b/diracx-db/src/diracx/db/sql/auth/db.py index a35a0ad19..d87283754 100644 --- a/diracx-db/src/diracx/db/sql/auth/db.py +++ b/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,10 +325,7 @@ 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. @@ -335,9 +333,7 @@ async def insert_refresh_token( """ # 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) diff --git a/diracx-db/src/diracx/db/sql/auth/schema.py b/diracx-db/src/diracx/db/sql/auth/schema.py index 0c7554318..54a917d9d 100644 --- a/diracx-db/src/diracx/db/sql/auth/schema.py +++ b/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),) diff --git a/diracx-db/tests/auth/test_refresh_token.py b/diracx-db/tests/auth/test_refresh_token.py index 28d6dfe9d..48eacb85f 100644 --- a/diracx-db/tests/auth/test_refresh_token.py +++ b/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) diff --git a/diracx-logic/src/diracx/logic/auth/token.py b/diracx-logic/src/diracx/logic/auth/token.py index c9b0a7e9c..4ac48c1e6 100644 --- a/diracx-logic/src/diracx/logic/auth/token.py +++ b/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,9 +411,7 @@ 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 @@ -401,9 +419,7 @@ async def insert_refresh_token( # 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 diff --git a/diracx-routers/src/diracx/routers/access_policies.py b/diracx-routers/src/diracx/routers/access_policies.py index c2e498d74..ed82aa489 100644 --- a/diracx-routers/src/diracx/routers/access_policies.py +++ b/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 diff --git a/diracx-routers/src/diracx/routers/auth/token.py b/diracx-routers/src/diracx/routers/auth/token.py index 3b3f7942d..1ff1d2fac 100644 --- a/diracx-routers/src/diracx/routers/auth/token.py +++ b/diracx-routers/src/diracx/routers/auth/token.py @@ -41,35 +41,19 @@ async def mint_token( access_payload: AccessTokenPayload, refresh_payload: RefreshTokenPayload | None, existing_refresh_token: str | None, - all_access_policies: dict[str, BaseAccessPolicy], settings: AuthSettings, ) -> TokenResponse: - """Enrich the token with policy specific content and mint it.""" + """Mint the token.""" if not refresh_payload and not existing_refresh_token: raise HTTPException( status_code=HTTPStatus.INTERNAL_SERVER_ERROR, detail="Refresh token is not set and no refresh token was provided", ) - # Enrich the token with policy specific content - dirac_access_policies = {} - dirac_refresh_policies = {} - for policy_name, policy in all_access_policies.items(): - access_extra, refresh_extra = policy.enrich_tokens( - access_payload, refresh_payload - ) - if access_extra: - dirac_access_policies[policy_name] = access_extra - if refresh_extra: - dirac_refresh_policies[policy_name] = refresh_extra - - # Create the access token - access_payload.dirac_policies = dirac_access_policies access_token = create_token(access_payload, settings) # Create the refresh token if refresh_payload: - refresh_payload.dirac_policies = dirac_refresh_policies refresh_token = create_token(refresh_payload, settings) elif existing_refresh_token: refresh_token = existing_refresh_token @@ -132,6 +116,7 @@ async def get_oidc_token( config, settings, available_properties, + all_access_policies, device_code=device_code, code=code, redirect_uri=redirect_uri, @@ -162,9 +147,7 @@ async def get_oidc_token( status_code=HTTPStatus.FORBIDDEN, detail=str(e), ) from e - return await mint_token( - access_payload, refresh_payload, refresh_token, all_access_policies, settings - ) + return await mint_token(access_payload, refresh_payload, refresh_token, settings) BASE_64_URL_SAFE_PATTERN = ( @@ -209,6 +192,7 @@ async def perform_legacy_exchange( available_properties=available_properties, settings=settings, config=config, + all_access_policies=all_access_policies, expires_minutes=expires_minutes, ) except ValueError as e: @@ -227,6 +211,4 @@ async def perform_legacy_exchange( status_code=HTTPStatus.FORBIDDEN, detail=str(e), ) from e - return await mint_token( - access_payload, refresh_payload, None, all_access_policies, settings - ) + return await mint_token(access_payload, refresh_payload, None, settings) diff --git a/diracx-routers/tests/auth/test_legacy_exchange.py b/diracx-routers/tests/auth/test_legacy_exchange.py index 6f28067e7..e01084b65 100644 --- a/diracx-routers/tests/auth/test_legacy_exchange.py +++ b/diracx-routers/tests/auth/test_legacy_exchange.py @@ -101,12 +101,14 @@ async def test_valid(test_client, legacy_credentials, expires_seconds): ) assert r.status_code == 200 user_info = r.json() + assert user_info["sub"] == "lhcb:b824d4dc-1f9d-4ee8-8df5-c0ae55d46041" assert user_info["vo"] == "lhcb" assert user_info["dirac_group"] == "lhcb_user" assert sorted(user_info["properties"]) == sorted( ["PrivateLimitedDelegation", "NormalUser"] ) + assert user_info["policies"] async def test_refresh_token(test_client, legacy_credentials): diff --git a/diracx-routers/tests/auth/test_standard.py b/diracx-routers/tests/auth/test_standard.py index b70f856a0..342de96ff 100644 --- a/diracx-routers/tests/auth/test_standard.py +++ b/diracx-routers/tests/auth/test_standard.py @@ -1490,3 +1490,54 @@ def fake_iam_server_error(*args, **kwargs): r = test_client.get("/api/auth/device", params={"user_code": user_code}) assert r.status_code == 502 assert r.json()["detail"] == "IAM error" + + +async def test_access_token_contains_policies(test_client, auth_httpx_mock: HTTPXMock): + """Test that dirac_policies exists in the access token after login.""" + # Get access token + access_token = _get_tokens(test_client)["access_token"] + + # Get user infos + r = test_client.get( + "/api/auth/userinfo", headers={"Authorization": f"Bearer {access_token}"} + ) + assert r.status_code == 200 + assert r.json()["policies"] + + +async def test_access_token_contains_policies_after_refresh( + test_client, auth_httpx_mock: HTTPXMock +): + """Test that dirac_policies are persisted in the access token after refreshing the token.""" + # Get tokens + tokens = _get_tokens(test_client) + initial_access_token = tokens["access_token"] + initial_refresh_token = tokens["refresh_token"] + + # Retrieve policies + r = test_client.get( + "/api/auth/userinfo", + headers={"Authorization": f"Bearer {initial_access_token}"}, + ) + assert r.status_code == 200 + initial_policies = r.json()["policies"] + assert initial_policies + + # Refresh the token + r = test_client.post( + "/api/auth/token", + data={ + "grant_type": "refresh_token", + "refresh_token": initial_refresh_token, + "client_id": DIRAC_CLIENT_ID, + }, + ) + assert r.status_code == 200 + new_access_token = r.json()["access_token"] + + # Check if policies are the same + r = test_client.get( + "/api/auth/userinfo", headers={"Authorization": f"Bearer {new_access_token}"} + ) + assert r.status_code == 200 + assert r.json()["policies"] == initial_policies diff --git a/diracx-testing/src/diracx/testing/utils.py b/diracx-testing/src/diracx/testing/utils.py index 9816457f7..e9e9dfc87 100644 --- a/diracx-testing/src/diracx/testing/utils.py +++ b/diracx-testing/src/diracx/testing/utils.py @@ -46,7 +46,7 @@ from uuid_utils import uuid7 from diracx.core.extensions import DiracEntryPoint -from diracx.core.models import AccessTokenPayload, RefreshTokenPayload +from diracx.core.models import AccessTokenPayload from diracx.core.settings import FactorySettings if TYPE_CHECKING: @@ -199,11 +199,8 @@ async def policy( pass @staticmethod - def enrich_tokens( - access_payload: AccessTokenPayload, - refresh_payload: RefreshTokenPayload | None, - ): - return {"PolicySpecific": "OpenAccessForTest"}, {} + def enrich_tokens(access_payload: AccessTokenPayload): + return {"PolicySpecific": "OpenAccessForTest"} enabled_systems = { e.name for e in select_from_extension(group=DiracEntryPoint.SERVICES)