Merge pull request #311 from minvws/key-id

User can now specify a key-id to their public key so we can

Author: jaytaph

app/config.py

@@ -111,6 +111,7 @@ def parse_space_separated(cls, v: Any) -> list[str] | None:
             return v.split() or None
         return v if v is None or isinstance(v, list) else None
 
+
 class Config(BaseModel):
     app: ConfigApp
     database: ConfigDatabase

app/db/entities/organization_key.py

@@ -28,6 +28,7 @@ class OrganizationKey(Base):
         "scope", JSONB, nullable=False, server_default="{}"
     )
     key_data: Mapped[str] = mapped_column("key_data", Text, nullable=False)
+    key_id: Mapped[str | None] = mapped_column("key_id", Text, nullable=True)
 
     organization = relationship("Organization", back_populates="keys")
 
@@ -37,4 +38,5 @@ def to_dict(self) -> dict[str, Any]:
             # We omit organization_id since this is an internal detail.
             "scope": self.scope,
             "key_data": self.key_data,
+            "key_id": self.key_id or "",
         }

app/db/repositories/org_key_repository.py

@@ -46,7 +46,7 @@ def get_by_org(self, org_id: uuid.UUID) -> Optional[Sequence[OrganizationKey]]:
         return self.db_session.execute(query).scalars().all()
 
     def create(
-        self, org_id: uuid.UUID, scope: list[str], key_data: str
+        self, org_id: uuid.UUID, scope: list[str], key_data: str, key_id: Optional[str]
     ) -> OrganizationKey:
         """
         Creates a new key entry.
@@ -55,6 +55,7 @@ def create(
             organization_id=org_id,
             scope=scope,
             key_data=key_data,
+            key_id=key_id,
         )
         self.db_session.add(entry)
         self.db_session.flush()

app/models/requests.py

@@ -1,7 +1,7 @@
 import base64
 import logging
 from datetime import datetime
-from typing import Any, Literal, List
+from typing import Any, Literal, List, Optional
 
 from pydantic import BaseModel, ConfigDict, model_validator, Field, field_validator
 
@@ -16,6 +16,7 @@
 
 class RegisterRequest(BaseModel):
     scope: List[str]
+    key_id: Optional[str]
 
 
 class OrgRequest(BaseModel):

app/routers/key.py

@@ -34,7 +34,7 @@ def post_key(
 
     # Create the key entry
     try:
-        key_resolver.create(org.id, req.scope, mtls_pub_key)
+        key_resolver.create(org.id, req.scope, req.key_id, mtls_pub_key)
     except AlreadyExistsError:
         logger.warning("key already exists for org_id=%s scope=%r", org.id, req.scope)
         raise HTTPException(

app/services/key_resolver.py

@@ -89,15 +89,15 @@ def resolve(self, org_id: uuid.UUID, scope: str) -> Optional[jwk.JWK]:
         return jwk.JWK.from_pem(entry.key_data.encode("ascii"))
 
     def create(
-        self, org_id: uuid.UUID, scope: list[str], key_data: str
+        self, org_id: uuid.UUID, scope: list[str], key_id: Optional[str], key_data: str
     ) -> OrganizationKey:
         scope = _normalize_scope(scope)
         key_data = key_data.strip()
 
         with self.db.get_db_session() as session:
             try:
                 entry = session.get_repository(OrganizationKeyRepository).create(
-                    org_id, scope, key_data
+                    org_id, scope, key_data, key_id
                 )
             except Exception as e:
                 logger.exception(

sql/007-key-id.sql

@@ -0,0 +1,3 @@
+-- Add key ID to be returned in the JWE kid header
+
+ALTER TABLE organization_key ADD COLUMN key_id TEXT;

tests/test_hsm_key_version_integration.py

@@ -99,7 +99,7 @@ def test_new_key_version_is_added_to_jwe(
         oin=OIN, name=f"Org {OIN}", max_key_usage=RidUsage.ReversiblePseudonym
     )
     private_key_pem, public_key_pem = _generate_rsa_keypair()
-    key_resolver.create(org.id, [SCOPE], public_key_pem)
+    key_resolver.create(org.id, [SCOPE], None, public_key_pem)
 
     # Route OPRF evaluation through a (mocked) HSM that reads its active key
     # versions from the same database the endpoint writes to.

tests/test_key_resolver.py

@@ -16,7 +16,7 @@ def test_resolver_create_and_resolve_roundtrip(
     key_resolver: KeyResolver, org_service: OrgService
 ) -> None:
     org = org_service.create(
-        oin="ura:00000099000000001000",
+        oin="oin:00000099000000001000",
         name="test org",
         max_key_usage=RidUsage.ReversiblePseudonym,
     )
@@ -27,7 +27,7 @@ def test_resolver_create_and_resolve_roundtrip(
         scope=["NVI", " lmr "],
         pub_key=TEST_PUBKEY,
     )
-    entry = key_resolver.create(org.id, req.scope, req.pub_key)
+    entry = key_resolver.create(org.id, req.scope, "my-key-id", req.pub_key)
 
     assert entry.organization_id == org.id
     assert sorted(entry.scope) == ["lmr", "nvi"]
@@ -46,7 +46,7 @@ def test_resolver_get_and_delete(
         max_key_usage=RidUsage.ReversiblePseudonym,
     )
 
-    e = key_resolver.create(org.id, ["*"], TEST_PUBKEY)
+    e = key_resolver.create(org.id, ["*"], "my-key-id", TEST_PUBKEY)
 
     items = key_resolver.get_by_org(org.id) or []
     assert len(items) == 1
@@ -60,3 +60,37 @@ def test_resolver_get_and_delete(
 
     items2 = key_resolver.get_by_org(org.id)
     assert items2 == []
+
+
+def test_resolver_create_persists_key_id(
+    key_resolver: KeyResolver, org_service: OrgService
+) -> None:
+    org = org_service.create(
+        oin="oin:00000099000000001000", name="test org", max_key_usage=RidUsage.ReversiblePseudonym
+    )
+
+    entry = key_resolver.create(org.id, ["nvi"], "kid-2024", TEST_PUBKEY)
+    assert entry.key_id == "kid-2024"
+
+    # key_id survives a round-trip through the database
+    stored = key_resolver.get_by_id(entry.id)
+    assert stored is not None
+    assert stored.key_id == "kid-2024"
+    assert stored.to_dict()["key_id"] == "kid-2024"
+
+
+def test_resolver_create_without_key_id(
+    key_resolver: KeyResolver, org_service: OrgService
+) -> None:
+    org = org_service.create(
+        oin="oin:00000099000000001000", name="test org", max_key_usage=RidUsage.ReversiblePseudonym
+    )
+
+    entry = key_resolver.create(org.id, ["nvi"], None, TEST_PUBKEY)
+    assert entry.key_id is None
+
+    stored = key_resolver.get_by_id(entry.id)
+    assert stored is not None
+    assert stored.key_id is None
+    # to_dict() represents a missing key_id as an empty string
+    assert stored.to_dict()["key_id"] == ""

tests/test_oprf_integration.py

@@ -56,7 +56,7 @@ def setup_org_and_key(
         max_key_usage=RidUsage.ReversiblePseudonym,
     )
     private_key_pem, public_key_pem = generate_rsa_keypair()
-    key_resolver.create(org.id, [scope], public_key_pem)
+    key_resolver.create(org.id, [scope], None, public_key_pem)
 
     return private_key_pem