Merge pull request #337 from minvws/prs-key-id

Added key-id we can send over via the JWE

Author: jaytaph

app/routers/exchange.py

@@ -217,7 +217,7 @@ def exchange_rid(
     if org is None:
         raise OrganizationNotFound(oin)
 
-    pub_key_jwk = key_resolver.resolve(org.id, req.recipientScope)
+    (pub_key_jwk, pub_key_id) = key_resolver.resolve(org.id, req.recipientScope)
     if pub_key_jwk is None:
         logger.warning(
             "no public key found for organization '%s' and scope '%s'",
@@ -232,6 +232,7 @@ def exchange_rid(
         scope=req.recipientScope,
         subject=f"rid:{rid}",
         pub_key=pub_key_jwk,
+        pub_key_id=pub_key_id,
         extra_claims={
             "ridUsage": rid_data["usage"],
         },
@@ -299,7 +300,7 @@ def exchange_pseudonym(
         )
         raise HTTPException(status_code=500, detail="Pseudonym exchange failed")
 
-    pub_key_jwk = key_resolver.resolve(org.id, req.recipientScope)
+    (pub_key_jwk, pub_key_id) = key_resolver.resolve(org.id, req.recipientScope)
     if pub_key_jwk is None:
         logger.warning(
             "no public key found for organization '%s' and scope '%s'",
@@ -313,6 +314,7 @@ def exchange_pseudonym(
         scope=req.recipientScope,
         subject=subject,
         pub_key=pub_key_jwk,
+        pub_key_id=pub_key_id,
     )
 
     return Response(

app/routers/oprf.py

@@ -70,7 +70,7 @@ def post_eval(
     pub_key_jwk = jwk.JWK.from_pem(key_entry.key_data.encode("ascii"))
 
     try:
-        result = oprf_service.eval_blind(req, pub_key_jwk)
+        result = oprf_service.eval_blind(req, pub_key_jwk, key_entry.key_id)
     except ValueError as e:
         log_event(
             logger,

app/services/key_resolver.py

@@ -84,13 +84,15 @@ def resolve_entry(self, org_id: uuid.UUID, scope: str) -> Optional[OrganizationK
         with self.db.get_db_session() as session:
             return session.get_repository(OrganizationKeyRepository).get(org_id, scope)
 
-    def resolve(self, org_id: uuid.UUID, scope: str) -> Optional[jwk.JWK]:
+    def resolve(
+        self, org_id: uuid.UUID, scope: str
+    ) -> tuple[Optional[jwk.JWK], Optional[str]]:
         entry = self.resolve_entry(org_id, scope)
 
         if entry is None:
-            return None
+            return None, None
 
-        return jwk.JWK.from_pem(entry.key_data.encode("ascii"))
+        return jwk.JWK.from_pem(entry.key_data.encode("ascii")), entry.key_id
 
     def create(
         self, org_id: uuid.UUID, scope: list[str], key_id: Optional[str], key_data: str

app/services/oprf/jwe_token.py

@@ -11,6 +11,7 @@ def build(
         scope: str,
         subject: str,
         pub_key: jwk.JWK,
+        pub_key_id: str | None,
         extra_claims: dict[str, Any] = {},
     ) -> str:
         """
@@ -28,7 +29,7 @@ def build(
         }
 
         protected_headers = {
-            "kid": pub_key.thumbprint(),
+            "kid": pub_key_id if pub_key_id else pub_key.thumbprint(),
             "alg": "RSA-OAEP-256",
             "enc": "A256GCM",
             "cty": "application/json",

app/services/oprf/oprf_service.py

@@ -68,7 +68,9 @@ def generate_server_key() -> str:
         """
         return base64.urlsafe_b64encode(pyoprf.keygen()).decode("ascii")
 
-    def eval_blind(self, req: BlindRequest, pub_key: jwk.JWK) -> OprfEvalResult:
+    def eval_blind(
+        self, req: BlindRequest, pub_key: jwk.JWK, pub_key_id: str | None
+    ) -> OprfEvalResult:
         """
         Evaluate a blind and returns a JWE encrypted on the pubkey, plus the
         key versions the blind was evaluated against
@@ -123,6 +125,7 @@ def eval_blind(self, req: BlindRequest, pub_key: jwk.JWK) -> OprfEvalResult:
             scope=req.recipientScope,
             subject=subject,
             pub_key=pub_key,
+            pub_key_id=pub_key_id,
             extra_claims={"extra_versions": extra_versions},
         )
 

tests/test_hsm_key_version_service.py

@@ -234,7 +234,7 @@ def fake_post(url: str, json: dict, **kwargs: object) -> MagicMock:  # type: ign
     )
 
     with patch("app.services.oprf.oprf_service.requests.post", side_effect=fake_post):
-        result = service.eval_blind(req, pub)
+        result = service.eval_blind(req, pub, None)
 
     assert result.key_versions == (2, 7)
 
@@ -328,7 +328,7 @@ def fake_post(url: str, json: dict, **kwargs: object) -> MagicMock:  # type: ign
     )
 
     with patch("app.services.oprf.oprf_service.requests.post", side_effect=fake_post):
-        result = service.eval_blind(req, pub)
+        result = service.eval_blind(req, pub, None)
 
     assert result.key_versions == (3, 5)
 

tests/test_key_resolver.py

@@ -36,8 +36,10 @@ def test_resolver_create_and_resolve_roundtrip(
     assert entry.organization_id == org.id
     assert sorted(entry.scope) == ["lmr", "nvi"]
 
-    key = key_resolver.resolve(org.id, "nvi")
+    (key, key_id) = key_resolver.resolve(org.id, "nvi")
     assert isinstance(key, jwk.JWK)
+    assert key_id is not None
+    assert key_id == "my-key-id"
     assert not key.has_private
 
 

tests/test_logging_events.py

@@ -93,7 +93,7 @@ def test_eval_blind_invalid_input_raises_invalid_blinded_input(
     service = OprfService(server_key=OprfService.generate_server_key())
 
     with pytest.raises(OprfEvaluationError) as exc:
-        service.eval_blind(_blind_request(), pub_key)
+        service.eval_blind(_blind_request(), pub_key, None)
 
     assert exc.value.error_type == "invalid_blinded_input"
 
@@ -110,7 +110,7 @@ def test_eval_blind_without_active_versions_raises_secret_version_destroyed(
     )
 
     with pytest.raises(OprfEvaluationError) as exc:
-        service.eval_blind(_blind_request(), pub_key)
+        service.eval_blind(_blind_request(), pub_key, None)
 
     assert exc.value.error_type == "secret_version_destroyed"
 
@@ -135,6 +135,6 @@ def test_eval_blind_hsm_failure_raises_crypto_evaluation_failure(
         ),
         pytest.raises(OprfEvaluationError) as exc,
     ):
-        service.eval_blind(_blind_request(), pub_key)
+        service.eval_blind(_blind_request(), pub_key, None)
 
     assert exc.value.error_type == "crypto_evaluation_failure"

tests/test_oprf_integration.py

@@ -337,7 +337,7 @@ def test_oprf_eval_failure_emits_failed_event_with_error_type(
     oprf_event_records: List[logging.LogRecord],
 ) -> None:
     class FailingOprfService:
-        def eval_blind(self, req: object, pub_key_jwk: object) -> str:
+        def eval_blind(self, req: object, pub_key_jwk: object, pub_key_id: str|None) -> str:
             raise OprfEvaluationError(
                 "invalid blinded input", error_type="invalid_blinded_input"
             )
@@ -372,7 +372,7 @@ def test_oprf_eval_when_service_rejects_blind_returns_bad_request(
     oprf_context: OprfIntegrationContext,
 ) -> None:
     class FailingOprfService:
-        def eval_blind(self, req: object, pub_key_jwk: object) -> str:
+        def eval_blind(self, req: object, pub_key_jwk: object, pub_key_id: str|None) -> str:
             raise ValueError("invalid blinded input")
 
     app.dependency_overrides[container.get_oprf_service] = lambda: FailingOprfService()