diff --git a/diracx-core/src/diracx/core/settings.py b/diracx-core/src/diracx/core/settings.py index 6e97dcfbe..50fe61b72 100644 --- a/diracx-core/src/diracx/core/settings.py +++ b/diracx-core/src/diracx/core/settings.py @@ -32,6 +32,7 @@ SecretStr, TypeAdapter, UrlConstraints, + model_validator, ) from pydantic_settings import BaseSettings, SettingsConfigDict @@ -151,6 +152,21 @@ def create(cls) -> Self: class AuthSettings(ServiceSettingsBase): """Settings for the authentication service.""" + @model_validator(mode="after") + def check_retention_greater_than_expiration(self) -> Self: + """Ensure retention times are bigger than expiration times to avoid deleting valid flows.""" + if self.completed_flow_retention_minutes <= ( + self.device_flow_expiration_seconds / 60 + ) or self.completed_flow_retention_minutes <= ( + self.authorization_flow_expiration_seconds / 60 + ): + raise ValueError( + f"completed_flow_retention_minutes ({self.completed_flow_retention_minutes} minutes) must be bigger" + f" than device_flow_expiration_seconds ({self.device_flow_expiration_seconds / 60} minutes) and" + f" authorization_flow_expiration_seconds: ({self.authorization_flow_expiration_seconds / 60} minutes)" + ) + return self + model_config = SettingsConfigDict( env_prefix="DIRACX_SERVICE_AUTH_", use_attribute_docstrings=True ) @@ -183,6 +199,13 @@ class AuthSettings(ServiceSettingsBase): before it must be exchanged for tokens. Default: 5 minutes. """ + completed_flow_retention_minutes: int = 60 + """Retention time in minutes for completed flow. + + The maximum retention time of flow after being completed + and before they are deleted. Default: 60 minutes. + """ + state_key: FernetKey """Encryption key used to encrypt/decrypt the state parameter passed to the IAM. @@ -226,6 +249,13 @@ class AuthSettings(ServiceSettingsBase): through a new authentication flow. Default: 60 minutes. """ + revoked_refresh_token_retention_minutes: int = 43200 + """Retention time in minutes for revoked refresh tokens. + + The maximum retention time of refresh tokens after being + revoked and before they are deleted. Default: 43200 minutes (30 days). + """ + available_properties: set[SecurityProperty] = Field( default_factory=SecurityProperty.available_properties ) diff --git a/diracx-db/src/diracx/db/sql/auth/db.py b/diracx-db/src/diracx/db/sql/auth/db.py index bcbe814dc..caf1b0226 100644 --- a/diracx-db/src/diracx/db/sql/auth/db.py +++ b/diracx-db/src/diracx/db/sql/auth/db.py @@ -6,7 +6,7 @@ from itertools import pairwise from dateutil.rrule import MONTHLY, rrule -from sqlalchemy import insert, select, text, update +from sqlalchemy import delete, insert, select, text, update from sqlalchemy.exc import IntegrityError, NoResultFound from sqlalchemy.ext.asyncio import AsyncConnection from uuid_utils import UUID, uuid7 @@ -339,3 +339,61 @@ async def revoke_user_refresh_tokens(self, subject): .where(RefreshTokens.sub == subject) .values(status=RefreshTokenStatus.REVOKED) ) + + async def clean_expired_refresh_tokens(self, max_validity: int) -> int: + """Delete expired refresh tokens. + + max_validity: Maximum validity time in minutes for refresh tokens. + """ + expired_date = str( + uuid7_from_datetime(substract_date(minutes=max_validity), randomize=False) + ) + stmt_expired = delete(RefreshTokens).where( + RefreshTokens.status == RefreshTokenStatus.CREATED, + RefreshTokens.jti < expired_date, + ) + res_expired = await self.conn.execute(stmt_expired) + + return res_expired.rowcount + + async def clean_revoked_refresh_tokens(self, max_retention: int) -> int: + """Delete old revoked refresh tokens. + + max_retention: Maximum retention time in minutes for revoked refresh tokens. + """ + revoked_date = str( + uuid7_from_datetime(substract_date(minutes=max_retention), randomize=False) + ) + stmt_revoked = delete(RefreshTokens).where( + RefreshTokens.status == RefreshTokenStatus.REVOKED, + RefreshTokens.jti < revoked_date, + ) + res_revoked = await self.conn.execute(stmt_revoked) + + return res_revoked.rowcount + + async def clean_expired_authorization_flows(self, max_retention: int) -> int: + """Delete old authorization flows. + + max_retention: Maximum retention time in minutes for expired authorization flows. + Must be bigger than authorization_flow_expiration_seconds. + """ + stmt_auth = delete(AuthorizationFlows).where( + AuthorizationFlows.creation_time < substract_date(minutes=max_retention), + ) + res_auth = await self.conn.execute(stmt_auth) + + return res_auth.rowcount + + async def clean_expired_device_flows(self, max_retention: int) -> int: + """Delete old device flows. + + max_retention: Maximum retention time in minutes for expired device flows. + Must be bigger than device_flow_expiration_seconds. + """ + stmt_device = delete(DeviceFlows).where( + DeviceFlows.creation_time < substract_date(minutes=max_retention), + ) + res_device = await self.conn.execute(stmt_device) + + return res_device.rowcount diff --git a/diracx-db/tests/auth/test_authorization_flow.py b/diracx-db/tests/auth/test_authorization_flow.py index 56a3d332c..1fb41731e 100644 --- a/diracx-db/tests/auth/test_authorization_flow.py +++ b/diracx-db/tests/auth/test_authorization_flow.py @@ -5,6 +5,7 @@ from diracx.core.exceptions import AuthorizationError from diracx.db.sql.auth.db import AuthDB +from diracx.db.sql.auth.schema import FlowStatus MAX_VALIDITY = 2 EXPIRED = 0 @@ -74,3 +75,46 @@ async def test_insert(auth_db: AuthDB): ) assert uuid1 != uuid2 + + +async def test_clean_authorization_flows(auth_db: AuthDB): + # Insert two authorization flows + async with auth_db as auth_db: + _ = await auth_db.insert_authorization_flow( + "client_id", "scope", "code_challenge", "S256", "redirect_uri" + ) + uuid2 = await auth_db.insert_authorization_flow( + "client_id2", "scope2", "code_challenge2", "S256", "redirect_uri2" + ) + uuid3 = await auth_db.insert_authorization_flow( + "client_id3", "scope3", "code_challenge3", "S256", "redirect_uri3" + ) + uuid4 = await auth_db.insert_authorization_flow( + "client_id4", "scope4", "code_challenge4", "S256", "redirect_uri4" + ) + + id_token = {"sub": "myIdToken"} + + async with auth_db as auth_db: + _, _ = await auth_db.authorization_flow_insert_id_token(uuid2, id_token, 1) + code3, _ = await auth_db.authorization_flow_insert_id_token(uuid3, id_token, 1) + code4, _ = await auth_db.authorization_flow_insert_id_token(uuid4, id_token, 1) + + async with auth_db as auth_db: + await auth_db.update_authorization_flow_status(code3, FlowStatus.DONE) + await auth_db.update_authorization_flow_status(code4, FlowStatus.ERROR) + + # Check the number of deleted authorization flow (should be 0) + async with auth_db as auth_db: + deleted_auth = await auth_db.clean_expired_authorization_flows(max_retention=30) + assert deleted_auth == 0 + + # Check the number of deleted authorization flow (should be 4: 1 PENDING, 1 READY, 1 DONE, 1 ERROR) + async with auth_db as auth_db: + deleted_auth = await auth_db.clean_expired_authorization_flows(max_retention=0) + assert deleted_auth == 4 + + # Check the number of deleted authorization flow (should be 0 because there is nothing left to delete) + async with auth_db as auth_db: + deleted_auth = await auth_db.clean_expired_authorization_flows(max_retention=0) + assert deleted_auth == 0 diff --git a/diracx-db/tests/auth/test_device_flow.py b/diracx-db/tests/auth/test_device_flow.py index 112e77898..168888f7e 100644 --- a/diracx-db/tests/auth/test_device_flow.py +++ b/diracx-db/tests/auth/test_device_flow.py @@ -8,7 +8,7 @@ from diracx.core.exceptions import AuthorizationError from diracx.db.sql.auth.db import AuthDB -from diracx.db.sql.auth.schema import USER_CODE_LENGTH +from diracx.db.sql.auth.schema import USER_CODE_LENGTH, FlowStatus from diracx.db.sql.utils.functions import substract_date MAX_VALIDITY = 2 @@ -139,3 +139,35 @@ async def test_device_flow_insert_id_token(auth_db: AuthDB): async with auth_db as auth_db: res = await auth_db.get_device_flow(device_code) assert res["IDToken"] == id_token + + +async def test_clean_device_flows(auth_db: AuthDB): + # Insert two device flows + async with auth_db as auth_db: + _, _ = await auth_db.insert_device_flow("client_id1", "scope1") + user_code2, _ = await auth_db.insert_device_flow("client_id2", "scope2") + _, device_code3 = await auth_db.insert_device_flow("client_id3", "scope3") + _, device_code4 = await auth_db.insert_device_flow("client_id4", "scope4") + + async with auth_db as auth_db: + await auth_db.device_flow_validate_user_code(user_code2, MAX_VALIDITY) + await auth_db.device_flow_insert_id_token( + user_code2, {"sub": "myIdToken"}, MAX_VALIDITY + ) + await auth_db.update_device_flow_status(device_code3, FlowStatus.DONE) + await auth_db.update_device_flow_status(device_code4, FlowStatus.ERROR) + + # Check the number of deleted device flows (should be 0) + async with auth_db as auth_db: + deleted_device = await auth_db.clean_expired_device_flows(max_retention=30) + assert deleted_device == 0 + + # Check the number of deleted device flows (should be 4: 1 PENDING, 1 READY, 1 DONE, 1 ERROR) + async with auth_db as auth_db: + deleted_device = await auth_db.clean_expired_device_flows(max_retention=0) + assert deleted_device == 4 + + # Check the number of deleted device flow (should be 0 because there is nothing left to delete) + async with auth_db as auth_db: + deleted_device = await auth_db.clean_expired_device_flows(max_retention=0) + assert deleted_device == 0 diff --git a/diracx-db/tests/auth/test_refresh_token.py b/diracx-db/tests/auth/test_refresh_token.py index 28d6dfe9d..39fc69b37 100644 --- a/diracx-db/tests/auth/test_refresh_token.py +++ b/diracx-db/tests/auth/test_refresh_token.py @@ -257,3 +257,46 @@ async def test_get_refresh_tokens(auth_db: AuthDB): # Check the number of retrieved refresh tokens (should be 3 refresh tokens) assert len(refresh_tokens) == 2 + + +async def test_clean_refresh_tokens(auth_db: AuthDB): + # Insert two refresh tokens + jtis = [] + async with auth_db as auth_db: + for _ in range(2): + jti = uuid7() + await auth_db.insert_refresh_token( + jti, + "subject", + "scope", + ) + jtis.append(jti) + + # Revoke one of the refresh token + async with auth_db as auth_db: + await auth_db.revoke_refresh_token(jtis[0]) + + # Check the number of deleted refresh tokens (should be 0) + async with auth_db as auth_db: + deleted_expired = await auth_db.clean_expired_refresh_tokens(max_validity=10) + assert deleted_expired == 0 + + async with auth_db as auth_db: + deleted_revoked = await auth_db.clean_revoked_refresh_tokens(max_retention=30) + assert deleted_revoked == 0 + + # Check the number of deleted refresh tokens (should be 1 of each) + async with auth_db as auth_db: + deleted_expired = await auth_db.clean_expired_refresh_tokens(max_validity=0) + assert deleted_expired == 1 + + async with auth_db as auth_db: + deleted_revoked = await auth_db.clean_revoked_refresh_tokens(max_retention=0) + assert deleted_revoked == 1 + + # Get all refresh tokens (Admin) + async with auth_db as auth_db: + refresh_tokens = await auth_db.get_user_refresh_tokens() + + # Check the number of retrieved refresh tokens (should be 0) + assert len(refresh_tokens) == 0 diff --git a/diracx-logic/src/diracx/logic/__main__.py b/diracx-logic/src/diracx/logic/__main__.py index 1fe642f4d..d7b7674a5 100644 --- a/diracx-logic/src/diracx/logic/__main__.py +++ b/diracx-logic/src/diracx/logic/__main__.py @@ -90,6 +90,23 @@ async def delete_jwk(args): save_jwks(path, jwks) +async def cleanup_authdb(args): + """Delete expired tokens and flows from the AuthDB.""" + logger.info("Deleting expired tokens and flows") + import os + + from diracx.core.settings import AuthSettings + from diracx.db.sql import AuthDB + from diracx.logic.auth.management import cleanup_expired_data + + settings = AuthSettings() + db_url = os.environ["DIRACX_DB_URL_AUTHDB"] + db = AuthDB(db_url) + async with db.engine_context(): + async with db: + await cleanup_expired_data(db, settings) + + def parse_args(): parser = argparse.ArgumentParser() subparsers = parser.add_subparsers(dest="command", required=True) @@ -120,8 +137,13 @@ def parse_args(): ) delete_jwk_parser.set_defaults(func=delete_jwk) + cleanup_authdb_parser = subparsers.add_parser( + "cleanup-authdb", help="Delete expired tokens and flows from the AuthDB" + ) + cleanup_authdb_parser.set_defaults(func=cleanup_authdb) + args = parser.parse_args() - logger.setLevel(logging.INFO) + logging.basicConfig(level=logging.INFO) asyncio.run(args.func(args)) diff --git a/diracx-logic/src/diracx/logic/auth/management.py b/diracx-logic/src/diracx/logic/auth/management.py index fb6a1963f..57027ad4f 100644 --- a/diracx-logic/src/diracx/logic/auth/management.py +++ b/diracx-logic/src/diracx/logic/auth/management.py @@ -2,6 +2,8 @@ from __future__ import annotations +import logging + from uuid_utils import UUID from diracx.core.exceptions import InvalidCredentialsError @@ -10,6 +12,8 @@ from diracx.db.sql import AuthDB from diracx.logic.auth.utils import verify_dirac_refresh_token +logger = logging.getLogger(__name__) + async def get_refresh_tokens( auth_db: AuthDB, @@ -57,3 +61,26 @@ async def revoke_refresh_token_by_refresh_token( # Decode and verify the refresh token jti, _, _ = await verify_dirac_refresh_token(token, settings) return await revoke_refresh_token_by_jti(auth_db=auth_db, subject=subject, jti=jti) + + +async def cleanup_expired_data(auth_db: AuthDB, settings: AuthSettings) -> None: + """Remove expired data from the auth database.""" + expired_tokens = await auth_db.clean_expired_refresh_tokens( + max_validity=settings.refresh_token_expire_minutes, + ) + logger.info("Deleted %d expired refresh tokens", expired_tokens) + + revoked_tokens = await auth_db.clean_revoked_refresh_tokens( + max_retention=settings.revoked_refresh_token_retention_minutes, + ) + logger.info("Deleted %d revoked refresh tokens", revoked_tokens) + + auth = await auth_db.clean_expired_authorization_flows( + max_retention=settings.completed_flow_retention_minutes, + ) + logger.info("Deleted %d expired authorization flows", auth) + + device = await auth_db.clean_expired_device_flows( + max_retention=settings.completed_flow_retention_minutes, + ) + logger.info("Deleted %d expired device flows", device) diff --git a/docs/admin/reference/env-variables.md b/docs/admin/reference/env-variables.md index 1770414e0..64c5625db 100644 --- a/docs/admin/reference/env-variables.md +++ b/docs/admin/reference/env-variables.md @@ -42,6 +42,15 @@ Expiration time in seconds for authorization code flow. The time window during which the authorization code remains valid before it must be exchanged for tokens. Default: 5 minutes. +### `DIRACX_SERVICE_AUTH_COMPLETED_FLOW_RETENTION_MINUTES` + +*Optional*, default value: `60` + +Retention time in minutes for completed flow. + +The maximum retention time of flow after being completed +and before they are deleted. Default: 60 minutes. + ### `DIRACX_SERVICE_AUTH_STATE_KEY` **Required** @@ -96,6 +105,15 @@ Expiration time in minutes for refresh tokens. The maximum lifetime of refresh tokens before they must be re-issued through a new authentication flow. Default: 60 minutes. +### `DIRACX_SERVICE_AUTH_REVOKED_REFRESH_TOKEN_RETENTION_MINUTES` + +*Optional*, default value: `43200` + +Retention time in minutes for revoked refresh tokens. + +The maximum retention time of refresh tokens after being +revoked and before they are deleted. Default: 43200 minutes (30 days). + ### `DIRACX_SERVICE_AUTH_AVAILABLE_PROPERTIES` *Optional*