feat: Autodeletion of secrets after full use or expiration.

Author: Robin-Van-de-Merghel

diracx-core/src/diracx/core/exceptions.py

@@ -157,11 +157,6 @@ class SecretHasExpiredError(GenericError):
     tail = "has expired"
 
 
-class OverusedSecretError(GenericError):
-    head = "Secret"
-    tail = "too much used"
-
-
 class SecretAlreadyExistsError(GenericError):
     head = "Secret"
     tail = "already exists"

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

@@ -5,12 +5,12 @@
 
 from sqlalchemy import DateTime, bindparam, insert, update
 from sqlalchemy.exc import IntegrityError, OperationalError
+from sqlalchemy.sql import delete
 
 from diracx.core.exceptions import (
     BadPilotCredentialsError,
     BadPilotVOError,
     CredentialsNotFoundError,
-    OverusedSecretError,
     PilotAlreadyExistsError,
     PilotNotFoundError,
     SecretAlreadyExistsError,
@@ -47,6 +47,7 @@ async def increment_pilot_local_secret_and_last_time_use(
         res = await self.conn.execute(stmt)
 
         if res.rowcount == 0:
+            await self.conn.rollback()
             raise PilotNotFoundError(
                 data={
                     "pilot_stamp": pilot_stamp,
@@ -70,7 +71,13 @@ async def increment_global_secret_use(
         res = await self.conn.execute(stmt)
 
         if res.rowcount == 0:
+            await self.conn.rollback()
             raise SecretNotFoundError(data={"secret_id": str(secret_id)})
+        if res.rowcount != 1:
+            await self.conn.rollback()
+            raise DBInBadStateError(
+                detail="This should not happen. Pilot should have a secret, but is not found."
+            )
 
     async def verify_pilot_secret(
         self, pilot_stamp: str, pilot_hashed_secret: str
@@ -90,6 +97,7 @@ async def verify_pilot_secret(
 
         # 3. Compare the secret_id
         if not secret["SecretID"] == pilot_credentials["PilotSecretID"]:
+
             raise BadPilotCredentialsError(
                 data={
                     "pilot_stamp": pilot_stamp,
@@ -104,6 +112,16 @@ async def verify_pilot_secret(
         # Convert the timezone, TODO: Change with #454: https://github.com/DIRACGrid/diracx/pull/454
         expiration = secret["SecretExpirationDate"].replace(tzinfo=timezone.utc)
         if expiration < now:
+
+            try:
+                await self.delete_secrets_bulk([secret["SecretID"]])
+            except SecretNotFoundError as e:
+                await self.conn.rollback()
+
+                raise DBInBadStateError(
+                    detail="This should not happen. Pilot should have a secret, but not found."
+                ) from e
+
             raise SecretHasExpiredError(
                 data={
                     "pilot_hashed_secret": pilot_hashed_secret,
@@ -112,42 +130,40 @@ async def verify_pilot_secret(
                 }
             )
 
-        # 5. Verify the secret counter
-        # 5.1 Only check if the SecretGlobalUseCountMax is defined
-        # If not defined, there is an infinite use.
-        if secret["SecretGlobalUseCountMax"]:
-            # 5.2 Finite use, we check if we can still login
-            if secret["SecretGlobalUseCount"] + 1 > secret["SecretGlobalUseCountMax"]:
-                raise OverusedSecretError(
-                    data={
-                        "pilot_stamp" "pilot_hashed_secret": pilot_hashed_secret,
-                        "secret_global_use_count": secret["SecretGlobalUseCount"],
-                        "secret_global_use_count_max": secret[
-                            "SecretGlobalUseCountMax"
-                        ],
-                    }
-                )
-
-        # 6. Now the pilot is authorized, increment the counters (globally and locally).
+        # 5. Now the pilot is authorized, increment the counters (globally and locally).
         try:
-            # 6.1 Increment the local count
+            # 5.1 Increment the local count
             await self.increment_pilot_local_secret_and_last_time_use(
                 pilot_secret_id=pilot_credentials["PilotSecretID"],
                 pilot_stamp=pilot_credentials["PilotStamp"],
             )
 
-            # 6.2 Increment the global count
+            # 5.2 Increment the global count
             await self.increment_global_secret_use(
                 secret_id=pilot_credentials["PilotSecretID"]
             )
         except Exception as e:  # Generic, to catch it.
             # Should NOT happen
             # Wrapped in a try/catch to still catch in case of an error in the counters
             # Caught and raised here to avoid raising a 4XX error
+            await self.conn.rollback()
+
             raise DBInBadStateError(
                 detail="This should not happen. Pilot has credentials, but has a corrupted secret."
             ) from e
 
+        # 6. Delete all secrets if its count attained the secret_global_use_count_max
+        if secret["SecretGlobalUseCountMax"]:
+            if secret["SecretGlobalUseCount"] + 1 == secret["SecretGlobalUseCountMax"]:
+                try:
+                    await self.delete_secrets_bulk([secret["SecretID"]])
+                except SecretNotFoundError as e:
+                    # Should NOT happen
+                    await self.conn.rollback()
+                    raise DBInBadStateError(
+                        detail="This should not happen. Pilot has credentials, but has corrupted secret."
+                    ) from e
+
     async def add_pilots_bulk(
         self,
         pilot_stamps: list[str],
@@ -217,6 +233,20 @@ async def add_pilots_credentials_bulk(
         # Used later to add an expiration date to the credentials
         return secrets
 
+    async def delete_secrets_bulk(self, secret_ids: list[int]):
+        """Bulk delete secrets."""
+        stmt = delete(PilotSecrets).where(PilotSecrets.secret_id.in_(secret_ids))
+
+        res = await self.conn.execute(stmt)
+
+        if res.rowcount != len(secret_ids):
+            await self.conn.rollback()
+
+            raise SecretNotFoundError(data={"secrets": str(secret_ids)})
+
+        # To avoid raise condition
+        await self.conn.commit()
+
     async def insert_unique_secrets_bulk(
         self,
         hashed_secrets: list[str],

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

@@ -90,7 +90,10 @@ class PilotToSecretMapping(PilotAgentsDBBase):
 
     # Primary key is (PilotSecretID, PilotStamp) pair
     pilot_secret_id = Column(
-        "PilotSecretID", Integer, ForeignKey("PilotSecrets.SecretID"), primary_key=True
+        "PilotSecretID",
+        Integer,
+        ForeignKey("PilotSecrets.SecretID", ondelete="CASCADE"),
+        primary_key=True,
     )
     pilot_stamp = Column("PilotStamp", String(32), primary_key=True)
     # Different from global use: only counts how many a specific pilot used a specific secret

diracx-db/tests/pilot_agents/test_pilot_agents_db.py

@@ -9,7 +9,6 @@
 from diracx.core.exceptions import (
     BadPilotVOError,
     CredentialsNotFoundError,
-    OverusedSecretError,
     PilotNotFoundError,
     SecretHasExpiredError,
     SecretNotFoundError,
@@ -223,7 +222,8 @@ async def test_create_pilot_and_verify_secret_too_much_secret_use(
         )
 
         # Second login, should not work because maxed out at 1 try
-        with pytest.raises(OverusedSecretError):
+        # If the foreign key works, we should have "SecretNotFoundError"
+        with pytest.raises(SecretNotFoundError):
             await pilot_agents_db.verify_pilot_secret(
                 pilot_stamp=pilot_stamp,
                 pilot_hashed_secret=pilot_hashed_secret,

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

@@ -14,7 +14,6 @@
     DiracHttpResponseError,
     ExpiredFlowError,
     InvalidCredentialsError,
-    OverusedSecretError,
     PendingAuthorizationError,
     PilotNotFoundError,
     SecretHasExpiredError,
@@ -305,11 +304,6 @@ async def pilot_login(
             status_code=status.HTTP_401_UNAUTHORIZED,
             detail="bad pilot_secret",
         ) from e
-    except OverusedSecretError as e:
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="secret has been overused",
-        ) from e
     except SecretHasExpiredError as e:
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,

diracx-routers/tests/auth/test_pilot_auth.py

@@ -74,6 +74,19 @@ async def test_create_pilot_and_verify_secret(test_client):
             pilot_secret_expiration_dates=[expiration_date],
         )
 
+    # -----------------  Wrong password  -----------------
+    body = {
+        "pilot_stamp": pilot_stamp,
+        "pilot_secret": "My 1ncr3d1bl3 t0k3n",
+    }
+
+    r = test_client.post("/api/auth/pilot-login", json=body)
+
+    assert r.status_code == 401, r.json()
+    assert r.json()["detail"] == "bad pilot_secret"
+
+    # ----------------- Good password  -----------------
+
     body = {"pilot_stamp": pilot_stamp, "pilot_secret": secret}
 
     r = test_client.post("/api/auth/pilot-login", json=body)
@@ -108,17 +121,6 @@ async def test_create_pilot_and_verify_secret(test_client):
     assert r.status_code == 401, r.json()
     assert r.json()["detail"] == "Invalid JWT"
 
-    # -----------------  Wrong password  -----------------
-    body = {
-        "pilot_stamp": pilot_stamp,
-        "pilot_secret": "My 1ncr3d1bl3 t0k3n",
-    }
-
-    r = test_client.post("/api/auth/pilot-login", json=body)
-
-    assert r.status_code == 401, r.json()
-    assert r.json()["detail"] == "bad pilot_secret"
-
     # -----------------  Wrong ID  -----------------
     body = {"pilot_stamp": "It is a stamp", "pilot_secret": secret}
 
@@ -178,7 +180,45 @@ async def test_create_pilot_and_verify_secret(test_client):
     r = test_client.post("/api/auth/pilot-login", json=body)
 
     assert r.status_code == 401
-    assert r.json()["detail"] == "secret has been overused"
+    assert r.json()["detail"] == "bad credentials"
+
+
+async def test_expired_secret(test_client):
+
+    # see https://github.com/DIRACGrid/diracx/blob/78e00aa57f4191034dbf643c7ed2857a93b53f60/diracx-routers/tests/pilots/test_pilot_logger.py#L37
+    db = test_client.app.dependency_overrides[PilotAgentsDB.transaction].args[0]
+
+    async with db as pilot_agents_db:
+        pilot_stamp = "pilot-stamp"
+        vo = "lhcb"
+        # Register a pilot
+        await pilot_agents_db.add_pilots_bulk(
+            vo=vo,
+            pilot_stamps=[pilot_stamp],
+            grid_type="grid-type",
+        )
+
+        secret = "AW0nd3rfulS3cr3t"
+        pilot_hashed_secret = hash(secret)
+
+        # Add creds
+        secrets_added = await pilot_agents_db.add_pilots_credentials_bulk(
+            pilot_stamps=[pilot_stamp],
+            pilot_hashed_secrets=[pilot_hashed_secret],
+            pilot_secret_use_count_max=1,  # Important later
+            vo=vo,
+        )
+
+        assert len(secrets_added) == 1
+
+        secret_added = secrets_added[0]
+
+        expiration_date = secret_added["SecretCreationDate"] + timedelta(seconds=2)
+
+        await pilot_agents_db.set_secret_expirations_bulk(
+            secret_ids=[secret_added["SecretID"]],
+            pilot_secret_expiration_dates=[expiration_date],
+        )
 
     # ----------------- Secret that expired -----------------
     sleep(2)
@@ -190,6 +230,16 @@ async def test_create_pilot_and_verify_secret(test_client):
     assert r.status_code == 401
     assert r.json()["detail"] == "secret expired"
 
+    # ----------------- Secret that expired, but reused -----------------
+    # Should be deleted by the verify_pilot_secret, because deleted
+
+    body = {"pilot_stamp": pilot_stamp, "pilot_secret": secret}
+
+    r = test_client.post("/api/auth/pilot-login", json=body)
+
+    assert r.status_code == 401
+    assert r.json()["detail"] == "bad credentials"
+
 
 async def test_create_pilots_with_credentials(normal_test_client):
     # Lots of request, to validate that it returns the credentials in the same order as the input references