From c0e41f4a84f7f22641adce43dc373ed647fd78a7 Mon Sep 17 00:00:00 2001 From: Atul Patil Date: Tue, 25 Aug 2026 17:52:13 +0530 Subject: [PATCH] Add single-round-trip mutual attestation This CL lets both endpoints verify each other over one /_attest-connection exchange, rather than the two flights the earlier mutual-attestation design required. The client exports the TLS session EKM, signs its own proof over it, and carries that proof in the request; the server verifies it before producing anything, then answers with its own proof, which the client verifies against the same EKM. Neither peer contributes a nonce. The EKM already derives from both peers' handshake randoms, so it is unique and unforgeable per session and serves as the freshness challenge directly. Dropping the nonces is what collapses the exchange to one round trip, and it removes the per-connection pending-handshake state the two-flight design had to keep between flights. Because both sides now sign the same value, peer_role and mode (plus protocol_version) inside SessionSignaturePayload provide the domain separation the separate nonces used to: a proof cannot be reflected back and accepted as the other side's, and a server-only proof cannot satisfy a mutual check. All three fields sit at their defaults for server-only proofs, so proto3 omits them and the existing signed payload is byte-for-byte unchanged. The binding is constant for the lifetime of the session, so re-attesting it would prove nothing new. Mutual sessions are therefore attested exactly once: the client refuses to revalidate, and the server answers a repeat attestation with 403 plus Connection: close and revokes the session's authorization. A failed exchange does not consume the session, so a rejected client may retry. Re-attestation means opening a new connection. The client shows its evidence before it can verify the server, which is the cost of a single round trip; MUTUAL_ATTESTATION.md documents the trade-off. --- ARCHITECTURE.md | 9 + MUTUAL_ATTESTATION.md | 198 ++++++++++ README.md | 3 +- src/prompt_encryption_sdk/attestation.py | 210 ++++++++++ src/prompt_encryption_sdk/client/client.py | 61 ++- .../client/connection.py | 171 +++++++-- src/prompt_encryption_sdk/client/validator.py | 75 +++- .../proto/attestation.proto | 57 ++- src/prompt_encryption_sdk/server/__init__.py | 42 +- src/prompt_encryption_sdk/server/asgi.py | 32 +- .../server/attestation.py | 265 ++++++++++--- src/prompt_encryption_sdk/server/wsgi.py | 35 +- tests/attestation_protocol_test.py | 237 ++++++++++++ tests/client/client_test.py | 65 ++++ tests/client/connection_test.py | 144 +++++++ tests/mutual_attestation_test.py | 325 ++++++++++++++++ ...mutual_attestation_tls_integration_test.py | 363 ++++++++++++++++++ tests/server/asgi_test.py | 70 +++- tests/server/attestation_test.py | 275 +++++++++++++ tests/server/wsgi_test.py | 45 ++- 20 files changed, 2548 insertions(+), 134 deletions(-) create mode 100644 MUTUAL_ATTESTATION.md create mode 100644 src/prompt_encryption_sdk/attestation.py create mode 100644 tests/attestation_protocol_test.py create mode 100644 tests/mutual_attestation_test.py create mode 100644 tests/mutual_attestation_tls_integration_test.py diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 3c92276..26dd8b3 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -141,6 +141,15 @@ By verifying this entire chain (Hardware $$\rightarrow$$ Token $$\rightarrow$$ I --- +## Mutual Attestation +When the client also runs inside a TEE, the same single `/_attest-connection` round trip can verify both directions: the client's proof travels in the request, the server verifies it before producing its own, and the server's proof travels back in the response. + +That exchange drops the client and server nonces entirely and uses the TLS Exported Keying Material as the challenge — the EKM is already derived from both peers' handshake randoms, so it is unique and unforgeable per session. The consequence is that the binding is constant for the session's lifetime, so a mutually attested session is attested exactly once and is never revalidated; re-attestation means opening a new connection. + +See [Mutual Attestation](MUTUAL_ATTESTATION.md) for the protocol, the role/mode domain separation that replaces the nonces, and the API. + +--- + ## Threat Model Matrix & Security Considerations ### What the SDK Protects Against | Threat | Protected? | Mechanism | diff --git a/MUTUAL_ATTESTATION.md b/MUTUAL_ATTESTATION.md new file mode 100644 index 0000000..a9fbcd4 --- /dev/null +++ b/MUTUAL_ATTESTATION.md @@ -0,0 +1,198 @@ +# Mutual Attestation (Single Round Trip) + +The base SDK proves one direction: the client verifies that the server is a +genuine TEE running approved software. Mutual attestation proves the other +direction too, so a server can refuse to answer anything but a confidential +client. + +It costs **one round trip** — the same `/_attest-connection` request the +server-only flow already makes. + +--- + +## The exchange + +```ascii +Client (TEE) Server (TEE) +| | +| [1] Standard TLS 1.3 handshake | +|<============================================================>| +| | ++--[2] EKM = export_keying_material("EXPORTER-Prompt-Encryption-SDK", 32) ++--[3] Proof_client = Sign(K_client, {SHA256(EKM), SHA256(T_client), +| role=CLIENT, mode=MUTUAL, version=1}) +| | +| [4] POST /_attest-connection { mode=MUTUAL, Proof_client } | +|------------------------------------------------------------>| +| +--[5] EKM derived independently +| +--[6] Verify Proof_client, else 4xx/5xx +| +--[7] Proof_server = Sign(K_server, {..., +| | role=SERVER, mode=MUTUAL}) +| | +| [8] { Proof_server, mode=MUTUAL, mutual_attestation_complete }| +|<------------------------------------------------------------| +| | ++--[9] Verify the echoed mode, then verify Proof_server | +| | +| [10] Verified E2E encrypted application data | +|=============================================================>| +``` + +The server verifies the client's proof **before** producing its own, so a +client that fails policy never receives a fresh server attestation. + +--- + +## Why there are no nonces + +Both peers bind to the TLS **Exported Keying Material** (RFC 5705), exported +with the protocol label and *no context*. + +The EKM is derived from the TLS master secret, which is derived from both +peers' handshake randoms. It is therefore already unique to this session and +unforgeable without the session keys — it *is* the freshness challenge. A +client nonce would add nothing the handshake did not already provide, and a +server nonce would force a second flight to carry it. + +Removing the nonces is what collapses the exchange to one round trip. It also +removes the per-connection handshake state the server would otherwise have to +keep between flights. + +### What replaces the nonces + +Because both peers now sign over the *same* value, two fields inside +`SessionSignaturePayload` do the work the separate nonces used to: + +| Field | Purpose | +| :--- | :--- | +| `peer_role` | `CLIENT` or `SERVER`. Stops a proof from being reflected back and accepted as the other side's proof. | +| `mode` + `protocol_version` | Pins the semantics of the binding, so a server-only proof cannot satisfy a mutual check (and vice versa). | + +Both are absent from the wire for server-only proofs, so the existing signed +payload is byte-for-byte unchanged. + +--- + +## No revalidation: a session is attested exactly once + +The EKM is constant for the lifetime of a TLS session. Re-attesting that +session would re-sign an identical payload and prove nothing new — there is no +fresh challenge left to answer. + +So mutual sessions are **not** revalidated: + +* **Client:** `revalidate_session()` raises, and the lazy revalidation timer is + disabled. Passing `revalidation_timeout` alongside `mutual_attestation` is a + `ValueError`. +* **Server:** a second `/_attest-connection` on a mutually attested session is + refused with `403 Forbidden` and `Connection: close`, and the session's + authorization is revoked so no further application traffic is served on it. + +To re-attest — after a key rotation, or when the attestation token nears +expiry — **open a new connection**. The new TLS session brings a new EKM and +therefore a genuinely fresh challenge. + +A *failed* exchange does not consume the session: a client that is rejected on +policy may retry on the same connection. + +--- + +## The trade-off: the client shows its evidence first + +Collapsing to one round trip means the client cannot verify the server before +sending its own proof. Whoever terminates the TLS session sees the client's GCA +attestation token — its platform and workload measurements — even if that peer +then fails the client's policy check. + +What limits the exposure: + +* The evidence travels **inside** the TLS tunnel, after the client has already + validated the server's certificate in the normal handshake. It is not visible + to network observers. +* The client's proof is signed over **this session's** EKM, so a peer that + harvests it cannot replay it on any other connection. +* `peer_role` is `CLIENT`, so the harvested proof cannot be turned around and + presented as a server proof either. + +What it costs: a peer holding a valid certificate for the endpoint learns the +client's measurements before proving it is a TEE. If your threat model treats +those measurements as secret from a certificate-holding-but-unattested peer, +prefer a two-flight exchange where the server proves itself first — at the cost +of the extra round trip and the per-connection handshake state it requires. + +--- + +## Usage + +### Server + +```python +from prompt_encryption_sdk.proto import attestation_pb2 +from prompt_encryption_sdk.server import run_uvicorn_app # or run_gunicorn_app + +client_policy = attestation_pb2.AttestationPolicy( + hw_model=attestation_pb2.HARDWARE_MODEL_TDX, + workload=attestation_pb2.WorkloadPolicy( + image_hash="sha256:YOUR_EXPECTED_CLIENT_IMAGE_HASH" + ), + gce_instance=attestation_pb2.GceInstancePolicy(project_id="your-project"), +) + +run_uvicorn_app( + app, + client_policy=client_policy, + # Omit to keep accepting server-only clients as well. + require_mutual_attestation=True, + host="0.0.0.0", + port=8443, + ssl_certfile="server.crt", + ssl_keyfile="server.key", +) +``` + +### Client + +The client needs its own TEE identity — the same `TokenManager` the server +uses, rotating its own keys and fetching its own attestation token. + +```python +from prompt_encryption_sdk.client import PromptEncryptionClient +from prompt_encryption_sdk.server import KeyManager, TokenManager + +token_manager = TokenManager(key_manager=KeyManager()) + +with token_manager: # Starts background key and token rotation. + client = PromptEncryptionClient( + server_policy, + mutual_attestation=True, + client_token_manager=token_manager, + ) + with client.session() as session: + response = session.post("https://server:8443/infer", json=payload) +``` + +`prompt_encryption_sdk.server` imports lazily, so a client can use +`KeyManager` and `TokenManager` without installing the server extras. + +--- + +## Failure modes + +| Situation | Result | +| :--- | :--- | +| Server answers a mutual request in server-only mode | Client raises `AttestationHandshakeError` — no silent downgrade | +| Client proof fails policy or signature check | Server errors out; no server proof is produced | +| Client proof is bound to a different TLS session | Signature check fails — the EKM does not transfer | +| Server proof replayed to the server as a client proof | Fails — `peer_role` differs | +| Server-only client reaches a `require_mutual_attestation` server | Rejected | +| Mutual request carries a `nonce` | Rejected — the peers must derive the same binding | +| Attested session re-attests | `403`, connection closed, authorization revoked | + +--- + +## Wire compatibility + +Every new proto field is added, none renumbered. `AttestationMode`'s zero value +is `SERVER_ONLY`, so an existing client's request parses as before and its +signed payload is unchanged. A server built from this SDK keeps serving +server-only clients unless `require_mutual_attestation` is set. diff --git a/README.md b/README.md index 164f330..5f36e99 100644 --- a/README.md +++ b/README.md @@ -9,6 +9,7 @@ This SDK is particularly useful for protecting sensitive data, ensuring that pro - **End-to-End Encryption**: Encrypts data in transit directly into the secure enclave, protecting it from intermediate infrastructure. - **Attested TLS**: Seamlessly ties Google Cloud Confidential Space attestation into the TLS handshake. - **Client & Server Support**: High-level primitives to build attested clients and compliant ASGI/WSGI servers. +- **Mutual Attestation**: When the client also runs in a TEE, both peers can verify each other in a single round trip. See [Mutual Attestation](MUTUAL_ATTESTATION.md). ## Prerequisites @@ -16,7 +17,7 @@ This SDK is particularly useful for protecting sensitive data, ensuring that pro - OpenSSL (for generating certificates if running a local test server) - A Google Cloud Project with Confidential Space enabled (for production server deployments) -For more technical details on how the SDK works, see the [Architecture Documentation](ARCHITECTURE.md). +For more technical details on how the SDK works, see the [Architecture Documentation](ARCHITECTURE.md). For the client-verifies-server-and-server-verifies-client flow, see [Mutual Attestation](MUTUAL_ATTESTATION.md). ## Installation diff --git a/src/prompt_encryption_sdk/attestation.py b/src/prompt_encryption_sdk/attestation.py new file mode 100644 index 0000000..c588e75 --- /dev/null +++ b/src/prompt_encryption_sdk/attestation.py @@ -0,0 +1,210 @@ +# Copyright 2026 The Prompt Encryption SDK Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared proof primitives for post-handshake attested TLS. + +Both endpoints produce the same shape of proof, so the signing logic lives +here rather than in the server package. Mutual attestation is a single round +trip: the client puts its proof in the request, the server verifies it and +puts its own proof in the response. + +Neither side contributes a nonce. The TLS Exported Keying Material (RFC 5705) +is already unique to the session -- it is derived from both peers' handshake +randoms -- so it *is* the freshness challenge. The consequence is that the +binding is constant for the lifetime of the session, which is why a session is +attested exactly once; see `AttestationReplayError`. +""" + +from collections.abc import Sequence +import hashlib +from typing import Protocol + +from prompt_encryption_sdk.proto import attestation_pb2 + + +# Wire version of the mutual exchange. Version 0 is the server-only flow. +MUTUAL_ATTESTATION_PROTOCOL_VERSION = 1 + + +class AttestationReplayError(Exception): + """Raised when a TLS session that already attested attests again. + + Because the signed binding is the session's EKM, a second exchange over the + same session would re-sign the identical payload and prove nothing new. The + connection is dropped rather than re-attested. + """ + + +class _SigningKeyManager(Protocol): + """The subset of KeyManager used to produce channel-bound proofs.""" + + def sign_payload(self, payload: bytes) -> bytes: + ... + + def sign_payload_mldsa(self, payload: bytes) -> bytes: + ... + + +class IdentitySnapshotProvider(Protocol): + """Provides an atomic snapshot of a TEE identity and its signing keys.""" + + key_manager: _SigningKeyManager + + def get_identity_snapshot(self) -> tuple[bytes, bytes, bytes]: + ... + + +def validate_required_verifiers(verifier_types: Sequence[int]) -> None: + """Rejects empty or unsupported attestation verifier requirements.""" + if not verifier_types: + raise ValueError("At least one required_verifier_type must be specified.") + if any( + verifier_type != attestation_pb2.VERIFIER_TYPE_GCA + for verifier_type in verifier_types + ): + raise ValueError( + f"Unsupported verifier types requested: {list(verifier_types)}" + ) + + +def build_signature_payload( + *, + tls_ekm: bytes, + attestation_token: bytes, + peer_role: int = attestation_pb2.ATTESTATION_PEER_ROLE_UNSPECIFIED, + mode: int = attestation_pb2.ATTESTATION_MODE_SERVER_ONLY, +) -> bytes: + """Builds the canonical payload signed by both the classical and PQC keys. + + In mutual mode both peers bind to the *same* EKM, so `peer_role` and `mode` + are the domain separators that keep one side's proof from being replayed as + the other's. + + Args: + tls_ekm: Exported Keying Material for this TLS session. + attestation_token: The raw GCA attestation token bytes. + peer_role: Which endpoint produced the proof. + mode: The negotiated attestation mode. + + Returns: + The serialized SessionSignaturePayload. For server-only proofs every added + field is at its default and therefore absent from the wire, so the bytes + are unchanged from the pre-mutual-attestation SDK. + """ + protocol_version = ( + MUTUAL_ATTESTATION_PROTOCOL_VERSION + if mode == attestation_pb2.ATTESTATION_MODE_MUTUAL + else 0 + ) + return attestation_pb2.SessionSignaturePayload( + ekm_hash=hashlib.sha256(tls_ekm).digest(), + token_hash=hashlib.sha256(attestation_token).digest(), + peer_role=peer_role, + mode=mode, + protocol_version=protocol_version, + ).SerializeToString(deterministic=True) + + +class AttestationProver: + """Produces a TEE identity proof bound to a live TLS session.""" + + def __init__(self, identity: IdentitySnapshotProvider): + self._identity = identity + + def __repr__(self) -> str: + return f"AttestationProver(identity={self._identity!r})" + + def create_proof( + self, + tls_ekm: bytes, + *, + peer_role: int = attestation_pb2.ATTESTATION_PEER_ROLE_UNSPECIFIED, + mode: int = attestation_pb2.ATTESTATION_MODE_SERVER_ONLY, + ) -> attestation_pb2.AttestationProof: + """Creates an ECDSA + ML-DSA proof for the supplied channel binding. + + Args: + tls_ekm: Exported Keying Material for this TLS session. + peer_role: The role this endpoint plays in the exchange. + mode: The negotiated attestation mode. + + Returns: + An AttestationProof carrying the GCA token, both public keys, and both + signatures over the canonical payload. + + Raises: + RuntimeError: If this endpoint has no attestation token yet. + """ + ecdsa_public_key, pqc_public_key, attestation_token = ( + self._identity.get_identity_snapshot() + ) + if not attestation_token: + raise RuntimeError("Attestation identity has no token.") + + payload = build_signature_payload( + tls_ekm=tls_ekm, + attestation_token=attestation_token, + peer_role=peer_role, + mode=mode, + ) + return attestation_pb2.AttestationProof( + evidence=[ + attestation_pb2.AttestationEvidence( + verifier_type=attestation_pb2.VERIFIER_TYPE_GCA, + gca_bundle=attestation_pb2.GcaTrustBundle( + attestation_token=attestation_token.decode("utf-8") + ), + ) + ], + instance_public_key=attestation_pb2.EcdsaP256PublicKey( + key_bytes=ecdsa_public_key + ), + session_signature=self._identity.key_manager.sign_payload(payload), + pqc_public_key=attestation_pb2.MlDsaPublicKey( + serialized_public_keyset=pqc_public_key + ), + pqc_session_signature=self._identity.key_manager.sign_payload_mldsa( + payload + ), + ) + + +def proof_from_response( + response: attestation_pb2.AttestConnectionResponse, +) -> attestation_pb2.AttestationProof: + """Copies the top-level response proof fields into a role-neutral proof.""" + proof = attestation_pb2.AttestationProof( + instance_public_key=response.instance_public_key, + session_signature=response.session_signature, + pqc_public_key=response.pqc_public_key, + pqc_session_signature=response.pqc_session_signature, + ) + proof.evidence.extend(response.evidence) + return proof + + +def response_from_proof( + proof: attestation_pb2.AttestationProof, + **response_fields, +) -> attestation_pb2.AttestConnectionResponse: + """Copies a proof into the AttestConnectionResponse wire shape.""" + response = attestation_pb2.AttestConnectionResponse( + instance_public_key=proof.instance_public_key, + session_signature=proof.session_signature, + pqc_public_key=proof.pqc_public_key, + pqc_session_signature=proof.pqc_session_signature, + **response_fields, + ) + response.evidence.extend(proof.evidence) + return response diff --git a/src/prompt_encryption_sdk/client/client.py b/src/prompt_encryption_sdk/client/client.py index 34932c6..09e6d9e 100644 --- a/src/prompt_encryption_sdk/client/client.py +++ b/src/prompt_encryption_sdk/client/client.py @@ -17,6 +17,7 @@ import dataclasses from typing import Any +from prompt_encryption_sdk import attestation as attestation_protocol from prompt_encryption_sdk.proto import attestation_pb2 from google.protobuf import text_format import requests @@ -33,9 +34,13 @@ def __init__( self, *, policy: attestation_pb2.AttestationPolicy, revalidation_timeout: int | None, + mutual_attestation: bool = False, + attestation_prover: "attestation_protocol.AttestationProver | None" = None, **kwargs ): self._policy = policy + self._mutual_attestation = mutual_attestation + self._attestation_prover = attestation_prover self._revalidation_timeout = revalidation_timeout if self._revalidation_timeout is None: self._revalidation_timeout = constants.DEFAULT_REVALIDATION_TIMEOUT @@ -45,6 +50,7 @@ def __repr__(self) -> str: policy_str = text_format.MessageToString(self._policy, as_one_line=True) return ( f"{self.__class__.__name__}(policy={policy_str}, " + f"mutual_attestation={self._mutual_attestation!r}, " f"revalidation_timeout={self._revalidation_timeout!r})" ) @@ -61,6 +67,8 @@ def init_poolmanager( maxsize=maxsize, block=block, policy=self._policy, + mutual_attestation=self._mutual_attestation, + attestation_prover=self._attestation_prover, revalidation_timeout=self._revalidation_timeout, **pool_kwargs ) @@ -81,21 +89,61 @@ class PromptEncryptionClient: with client.session() as session: response = session.post("https://server/infer", json=data) + # Mutual: this client runs in a TEE too and proves it in the same + # round trip. No revalidation; open a new connection to re-attest. + with token_manager: # server.TokenManager for this client's identity + client = PromptEncryptionClient( + policy, + mutual_attestation=True, + client_token_manager=token_manager, + ) + Attributes: policy: The security policy to enforce. revalidation_timeout: Seconds before a session is considered 'stale' and - requires re-attestation. Default: 3300s (55m). + requires re-attestation. Default: 3300s (55m). Not supported when + mutual_attestation is enabled. + mutual_attestation: Also prove this client's own TEE identity to the + server, in the same round trip. Sessions established this way are + never revalidated; the server drops any session that tries. + client_token_manager: A started TokenManager (or any object exposing + `key_manager` and `get_identity_snapshot()`) holding this confidential + client's keys and attestation token. Required in mutual mode. """ policy: attestation_pb2.AttestationPolicy revalidation_timeout: int | None = None + mutual_attestation: bool = False + client_token_manager: Any | None = None + + def __post_init__(self) -> None: + if not self.mutual_attestation: + return + if self.client_token_manager is None: + raise ValueError( + "client_token_manager is required when mutual_attestation is" + " enabled." + ) + if self.revalidation_timeout is not None: + raise ValueError( + "revalidation_timeout is not supported with mutual_attestation; a" + " mutually attested session is attested exactly once." + ) def session(self) -> requests.Session: """Creates a new requests.Session with the Attested Adapter mounted.""" session = requests.Session() - adapter = AttestedHTTPSAdapter( + adapter_kwargs = dict( policy=self.policy, revalidation_timeout=self.revalidation_timeout ) + if self.mutual_attestation: + adapter_kwargs.update( + mutual_attestation=True, + attestation_prover=attestation_protocol.AttestationProver( + self.client_token_manager + ), + ) + adapter = AttestedHTTPSAdapter(**adapter_kwargs) session.mount("https://", adapter) return session @@ -106,8 +154,15 @@ def post( url: str, json: Any, revalidation_timeout: int | None = None, + mutual_attestation: bool = False, + client_token_manager: Any | None = None, **kwargs ) -> requests.Response: """Convenience method for a single-shot attested request.""" - with cls(policy, revalidation_timeout=revalidation_timeout).session() as sess: + with cls( + policy, + revalidation_timeout=revalidation_timeout, + mutual_attestation=mutual_attestation, + client_token_manager=client_token_manager, + ).session() as sess: return sess.post(url, json=json, **kwargs) diff --git a/src/prompt_encryption_sdk/client/connection.py b/src/prompt_encryption_sdk/client/connection.py index edd1e68..76fbdd6 100644 --- a/src/prompt_encryption_sdk/client/connection.py +++ b/src/prompt_encryption_sdk/client/connection.py @@ -17,6 +17,12 @@ This module defines `AttestedHTTPSConnection`, `AttestedHTTPSConnectionPool`, and `AttestedPoolManager` which extend urllib3's classes to perform attestation handshakes and session revalidation. + +Mutual attestation is a single round trip. The client exports the session EKM, +signs its own proof over it, and sends that proof in the very first request; +the server verifies it and answers with its own proof. Since the binding is +the session EKM rather than a per-exchange nonce, a mutually attested session +is never revalidated -- the server drops any session that tries. """ import datetime @@ -27,6 +33,7 @@ import ssl import time +from prompt_encryption_sdk import attestation as attestation_protocol from prompt_encryption_sdk.ekm import exporter as ekm_exporter from prompt_encryption_sdk.proto import attestation_pb2 from google.protobuf import json_format @@ -56,6 +63,12 @@ def __init__( **kwargs, ): self._policy = kwargs.pop("policy", None) + self._mutual_attestation = kwargs.pop("mutual_attestation", False) + self._attestation_prover = kwargs.pop("attestation_prover", None) + if self._mutual_attestation and self._attestation_prover is None: + raise ValueError( + "attestation_prover is required when mutual_attestation is enabled." + ) # Default revalidation timeout is 55 minutes. revalidation_timeout = kwargs.pop( "revalidation_timeout", datetime.timedelta(seconds=3300) @@ -78,6 +91,7 @@ def __repr__(self) -> str: return ( f"AttestedHTTPSConnection(host={self.host!r}, port={self.port}, " f"is_attested={self.is_attested}, " + f"mutual_attestation={self._mutual_attestation}, " f"revalidation_timeout={self._revalidation_timeout})" ) @@ -129,8 +143,42 @@ def _process_attestation_response( ) from e def _perform_attestation_handshake(self) -> None: - """Executes the attestation protocol (Request -> Response -> Validate).""" + """Executes server-only or mutual post-handshake attestation.""" + try: + if self._mutual_attestation: + self._perform_mutual_attestation() + else: + self._perform_server_attestation() + # Update timestamp on success + self._last_attestation_time = time.time() + + except (socket.error, ssl.SSLError) as e: + raise exceptions.AttestationHandshakeError( + "Network error during attestation" + ) from e + + def _send_attestation_request( + self, request: attestation_pb2.AttestConnectionRequest + ) -> attestation_pb2.AttestConnectionResponse: + """Sends the attestation flight over the established TLS socket.""" + json_body = json_format.MessageToJson(request) + # Note: We call the superclass's `request` method explicitly here to + # avoid a recursive loop during the handshake. + # Using HTTPConnection request method to send raw bytes. + super().request( + "POST", + constants.ATTESTATION_ENDPOINT, + body=json_body, + headers={ + "Content-Type": "application/json", + "Accept": "application/json, application/x-protobuf", + }, + ) + return self._process_attestation_response(self.getresponse()) + + def _perform_server_attestation(self) -> None: + """Runs the server-only exchange, challenging with a fresh client nonce.""" # A. Generate a fresh nonce nonce = secrets.token_bytes(constants.NONCE_LENGTH) @@ -138,40 +186,79 @@ def _perform_attestation_handshake(self) -> None: required_verifier_type=[attestation_pb2.VerifierType.VERIFIER_TYPE_GCA], nonce=nonce, ) - json_body = json_format.MessageToJson(attestation_req) # B. Send Request - try: - # Note: We call the superclass's `request` method explicitly here to - # avoid a recursive loop during the handshake. - # Using HTTPConnection request method to send raw bytes. - super().request( - "POST", - constants.ATTESTATION_ENDPOINT, - body=json_body, - headers={ - "Content-Type": "application/json", - "Accept": "application/json, application/x-protobuf", - }, - ) - response = self.getresponse() - attest_resp = self._process_attestation_response(response) + attest_resp = self._send_attestation_request(attestation_req) - # D. Export Keying Material - tls_ekm = self._ekm_exporter_fn( - self.sock, constants.EKM_LENGTH, constants.EKM_LABEL, context=nonce # pyrefly: ignore[bad-argument-type] - ) + # C. Export Keying Material + tls_ekm = self._ekm_exporter_fn( + self.sock, constants.EKM_LENGTH, constants.EKM_LABEL, context=nonce # pyrefly: ignore[bad-argument-type] + ) - # E. Validate - att_validator = self._attestation_validator_cls(self._policy) - att_validator.validate(attest_resp, tls_ekm) + # D. Validate + self._validate_server_proof(attest_resp, tls_ekm) - # Update timestamp on success - self._last_attestation_time = time.time() + def _perform_mutual_attestation(self) -> None: + """Runs the one-round-trip mutual exchange. - except (socket.error, ssl.SSLError) as e: + The EKM is exported with no context, so both peers derive the same value + straight from the TLS session and neither has to contribute a nonce. + """ + # A. Export Keying Material -- this is the challenge for both peers. + tls_ekm = self._ekm_exporter_fn( + self.sock, constants.EKM_LENGTH, constants.EKM_LABEL # pyrefly: ignore[bad-argument-type] + ) + + # B. Prove our own identity before asking the server for its proof. + client_proof = self._attestation_prover.create_proof( + tls_ekm, + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_CLIENT, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ) + request = attestation_pb2.AttestConnectionRequest( + required_verifier_type=[attestation_pb2.VERIFIER_TYPE_GCA], + protocol_version=( + attestation_protocol.MUTUAL_ATTESTATION_PROTOCOL_VERSION + ), + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + client_attestation=client_proof, + ) + response = self._send_attestation_request(request) + + # C. Fail closed unless the server echoed the mode we asked for. A server + # that answers with a server-only proof must not be silently accepted. + if ( + response.protocol_version + != attestation_protocol.MUTUAL_ATTESTATION_PROTOCOL_VERSION + or response.mode != attestation_pb2.ATTESTATION_MODE_MUTUAL + or not response.mutual_attestation_complete + ): raise exceptions.AttestationHandshakeError( - "Network error during attestation" - ) from e + "Mutual attestation negotiation failed; refusing server-only" + " downgrade." + ) + + # D. Validate the server proof against the same binding. + self._validate_server_proof( + response, + tls_ekm, + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_SERVER, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ) + + def _validate_server_proof( + self, + response: attestation_pb2.AttestConnectionResponse, + tls_ekm: bytes, + **binding, + ) -> None: + """Validates a server proof and releases validator-owned resources.""" + att_validator = self._attestation_validator_cls(self._policy) + try: + att_validator.validate(response, tls_ekm, **binding) + finally: + close = getattr(att_validator, "close", None) + if close is not None: + close() def connect(self) -> None: """Establishes the TLS connection and performs attestation.""" @@ -226,6 +313,11 @@ def request(self, method, url, body=None, headers=None, **kwargs) -> None: def _should_revalidate(self) -> bool: """Checks if the time since last attestation exceeds the timeout.""" + # A mutual session binds to the session EKM, which never changes, so + # re-attesting it would prove nothing. The server drops such a session. + if self._mutual_attestation: + return False + # Initial connection case if self._last_attestation_time == 0.0: return False @@ -238,7 +330,16 @@ def revalidate_session(self) -> None: """Performs session revalidation on the existing connection. Triggers a fresh attestation handshake to verify key rotation and validity. + + Raises: + exceptions.PromptEncryptionError: If the socket is closed, or if this is + a mutually attested session, which cannot be revalidated. """ + if self._mutual_attestation: + raise exceptions.PromptEncryptionError( + "Mutually attested sessions cannot be revalidated. Open a new" + " connection to re-attest." + ) if self.sock is None: raise exceptions.PromptEncryptionError( "Cannot revalidate: Socket is closed." @@ -265,6 +366,8 @@ def __init__( *, port=None, policy=None, + mutual_attestation=False, + attestation_prover=None, revalidation_timeout=datetime.timedelta(seconds=3300), ekm_exporter_fn=ekm_exporter.export_keying_material, attestation_validator_cls=validator.AttestationValidator, @@ -272,11 +375,15 @@ def __init__( ): kwargs.update({ "policy": policy, + "mutual_attestation": mutual_attestation, + "attestation_prover": attestation_prover, "revalidation_timeout": revalidation_timeout, "ekm_exporter_fn": ekm_exporter_fn, "attestation_validator_cls": attestation_validator_cls, }) self._policy = policy + self._mutual_attestation = mutual_attestation + self._attestation_prover = attestation_prover self._revalidation_timeout = revalidation_timeout super().__init__(host, port=port, **kwargs) @@ -295,12 +402,16 @@ def __init__( self, *, policy=None, + mutual_attestation=False, + attestation_prover=None, revalidation_timeout=datetime.timedelta(seconds=3300), ekm_exporter_fn=ekm_exporter.export_keying_material, attestation_validator_cls=validator.AttestationValidator, **kwargs, ): self._policy = policy + self._mutual_attestation = mutual_attestation + self._attestation_prover = attestation_prover self._revalidation_timeout = revalidation_timeout self._ekm_exporter_fn = ekm_exporter_fn self._attestation_validator_cls = attestation_validator_cls @@ -321,6 +432,8 @@ def _new_pool( host, port=port, policy=self._policy, + mutual_attestation=self._mutual_attestation, + attestation_prover=self._attestation_prover, ekm_exporter_fn=self._ekm_exporter_fn, attestation_validator_cls=self._attestation_validator_cls, revalidation_timeout=self._revalidation_timeout, diff --git a/src/prompt_encryption_sdk/client/validator.py b/src/prompt_encryption_sdk/client/validator.py index 9ee17a6..64543bf 100644 --- a/src/prompt_encryption_sdk/client/validator.py +++ b/src/prompt_encryption_sdk/client/validator.py @@ -22,6 +22,7 @@ import types from typing import Any +from prompt_encryption_sdk import attestation as attestation_protocol from prompt_encryption_sdk.client import constants from prompt_encryption_sdk.client import exceptions from prompt_encryption_sdk.proto import attestation_pb2 @@ -151,6 +152,9 @@ def validate( response: attestation_pb2.AttestConnectionResponse, tls_ekm: bytes, expected_nonce: bytes | None = None, + *, + peer_role: int = attestation_pb2.ATTESTATION_PEER_ROLE_UNSPECIFIED, + mode: int = attestation_pb2.ATTESTATION_MODE_SERVER_ONLY, ) -> None: """Validates the AttestConnectionResponse from the server. @@ -158,12 +162,48 @@ def validate( response: The parsed proto response. tls_ekm: The Exported Keying Material from the TLS socket. expected_nonce: The fresh challenge nonce generated by the client. + peer_role: Expected producer role included in the signed payload. + mode: Negotiated attestation mode included in the signed payload. + + Raises: + AttestationVerificationError: If validation fails. + PolicyViolationError: If policy check fails. + """ + self.validate_proof( + attestation_protocol.proof_from_response(response), + tls_ekm, + expected_nonce=expected_nonce, + peer_role=peer_role, + mode=mode, + ) + + def validate_proof( + self, + proof: attestation_pb2.AttestationProof, + tls_ekm: bytes, + expected_nonce: bytes | None = None, + *, + peer_role: int = attestation_pb2.ATTESTATION_PEER_ROLE_UNSPECIFIED, + mode: int = attestation_pb2.ATTESTATION_MODE_SERVER_ONLY, + ) -> None: + """Validates a role-neutral proof produced by either peer. + + The server uses this to check a confidential client's proof, so the checks + below are deliberately symmetric: nothing here assumes the prover is the + server. + + Args: + proof: The proof to validate. + tls_ekm: The Exported Keying Material from the TLS socket. + expected_nonce: The fresh challenge nonce generated by the verifier. + peer_role: Expected producer role included in the signed payload. + mode: Negotiated attestation mode included in the signed payload. Raises: AttestationVerificationError: If validation fails. PolicyViolationError: If policy check fails. """ - if not response.evidence: + if not proof.evidence: raise exceptions.AttestationVerificationError( "No attestation evidence provided." ) @@ -172,7 +212,7 @@ def validate( gca_bundle = next( ( ev.gca_bundle - for ev in response.evidence + for ev in proof.evidence if ev.verifier_type == attestation_pb2.VerifierType.VERIFIER_TYPE_GCA ), @@ -197,12 +237,12 @@ def validate( # 4. Verify Instance Key Binding # Checks that the Instance Public Key hashes are inside the Token's 'eat_nonce' - ecdsa_pub_bytes = response.instance_public_key.key_bytes + ecdsa_pub_bytes = proof.instance_public_key.key_bytes if not ecdsa_pub_bytes: raise exceptions.AttestationVerificationError( "ECDSA public key is missing." ) - pqc_pub_bytes = response.pqc_public_key.serialized_public_keyset + pqc_pub_bytes = proof.pqc_public_key.serialized_public_keyset if not pqc_pub_bytes: raise exceptions.AttestationVerificationError( "PQC public key is missing." @@ -216,37 +256,36 @@ def validate( ) # 5. Verify TLS Session Binding (Signature over EKM) - if not response.session_signature: + if not proof.session_signature: raise exceptions.AttestationVerificationError( "session signature is missing." ) - if not response.pqc_session_signature: + if not proof.pqc_session_signature: raise exceptions.AttestationVerificationError( "PQC session signature is missing." ) # Reconstruct the payload to verify the cryptographic binding of the - # TLS session. - token_hash_bytes = hashlib.sha256( - gca_bundle.attestation_token.encode("utf-8") - ).digest() - ekm_hash_bytes = hashlib.sha256(tls_ekm).digest() - payload_msg = attestation_pb2.SessionSignaturePayload( - ekm_hash=ekm_hash_bytes, token_hash=token_hash_bytes + # TLS session. In mutual mode both peers bind to the same EKM, so + # peer_role and mode are what stop a proof from being reflected back. + payload = attestation_protocol.build_signature_payload( + tls_ekm=tls_ekm, + attestation_token=gca_bundle.attestation_token.encode("utf-8"), + peer_role=peer_role, + mode=mode, ) - payload = payload_msg.SerializeToString() # Verify classical signature self._verify_session_signature( - response.instance_public_key, - signature=response.session_signature, + proof.instance_public_key, + signature=proof.session_signature, payload=payload, ) # Verify PQC signature self._verify_session_signature_mldsa( - response.pqc_public_key, - signature=response.pqc_session_signature, + proof.pqc_public_key, + signature=proof.pqc_session_signature, payload=payload, ) diff --git a/src/prompt_encryption_sdk/proto/attestation.proto b/src/prompt_encryption_sdk/proto/attestation.proto index 1cc8c87..8b4283e 100644 --- a/src/prompt_encryption_sdk/proto/attestation.proto +++ b/src/prompt_encryption_sdk/proto/attestation.proto @@ -15,11 +15,37 @@ enum VerifierType { VERIFIER_TYPE_GCA = 1; } +// AttestationMode is selected by the client on its single attestation flight. +// The zero value is intentionally server-only so requests from existing +// clients keep their current behavior. +enum AttestationMode { + ATTESTATION_MODE_SERVER_ONLY = 0; + ATTESTATION_MODE_MUTUAL = 1; +} + +// AttestationPeerRole prevents a proof produced by one side from being +// reflected back and accepted as a proof for the other side. +enum AttestationPeerRole { + ATTESTATION_PEER_ROLE_UNSPECIFIED = 0; + ATTESTATION_PEER_ROLE_SERVER = 1; + ATTESTATION_PEER_ROLE_CLIENT = 2; +} + message AttestConnectionRequest { // A list of verifiers that the server must be endorsed by. repeated VerifierType required_verifier_type = 1; - // A random nonce for generating unique TLS session material. + // A random nonce for generating unique TLS session material. Server-only + // mode uses it as the EKM context. Mutual mode derives freshness from the + // TLS session itself, so this field must be empty there. bytes nonce = 2; + // Protocol version 0 is the server-only flow. Mutual attestation uses + // version 1 and fails closed if the peer does not echo it. + uint32 protocol_version = 3; + AttestationMode mode = 4; + // The client's own TEE proof. Present only in mutual mode, where it makes + // the exchange a single round trip: the client proves itself in the request + // and the server proves itself in the response. + AttestationProof client_attestation = 5; } // GcaTrustBundle contains the materials needed to verify a GCA-based @@ -55,6 +81,27 @@ message SessionSignaturePayload { bytes ekm_hash = 1; // The raw 32-byte SHA-256 hash of the attestation token. bytes token_hash = 2; + // The remaining fields are absent on the wire for server-only proofs, + // preserving their existing canonical signed payload byte-for-byte. + // + // Which endpoint produced this proof. Both peers bind to the same EKM in + // mutual mode, so the role is what stops a proof from being reflected. + AttestationPeerRole peer_role = 3; + // The negotiated mode, which pins the semantics of the signed binding. + AttestationMode mode = 4; + // The protocol version the proof was produced under. + uint32 protocol_version = 5; +} + +// AttestationProof is role-neutral evidence that either endpoint can produce. +// AttestConnectionResponse keeps its original top-level proof fields so the +// server-only wire format is unchanged. +message AttestationProof { + repeated AttestationEvidence evidence = 1; + EcdsaP256PublicKey instance_public_key = 2; + bytes session_signature = 3; + MlDsaPublicKey pqc_public_key = 4; + bytes pqc_session_signature = 5; } // AttestConnectionResponse is sent by the server to the client. It is sent only @@ -71,6 +118,14 @@ message AttestConnectionResponse { MlDsaPublicKey pqc_public_key = 4; // PQC signature over the SessionSignaturePayload bytes pqc_session_signature = 5; + // Negotiation echo. These fields are default/absent for the server-only + // flow. The client fails closed unless the server echoes the mutual mode it + // asked for, which is what prevents a silent downgrade. + uint32 protocol_version = 6; + AttestationMode mode = 7; + // True once the server has verified this client's proof and authorized the + // TLS session for application traffic. + bool mutual_attestation_complete = 8; } // Enum for the supported TEE hardware CPU models. diff --git a/src/prompt_encryption_sdk/server/__init__.py b/src/prompt_encryption_sdk/server/__init__.py index 01ba7bb..4654747 100644 --- a/src/prompt_encryption_sdk/server/__init__.py +++ b/src/prompt_encryption_sdk/server/__init__.py @@ -12,15 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Server-side components for Prompt Encryption SDK.""" +"""Server-side components for Prompt Encryption SDK. -from prompt_encryption_sdk.server.asgi import PromptEncryptionASGIMiddleware -from prompt_encryption_sdk.server.asgi import run_uvicorn_app -from prompt_encryption_sdk.server.attestation import AttestedTLS -from prompt_encryption_sdk.server.keys import KeyManager -from prompt_encryption_sdk.server.token import TokenManager -from prompt_encryption_sdk.server.wsgi import PromptEncryptionWSGIMiddleware -from prompt_encryption_sdk.server.wsgi import run_gunicorn_app +Imports are lazy so a mutually attesting *client* can use KeyManager and +TokenManager to hold its own TEE identity without installing the ASGI, WSGI, +Gunicorn, or Uvicorn dependencies. +""" + +import importlib __all__ = ( "AttestedTLS", @@ -31,3 +30,30 @@ "run_gunicorn_app", "run_uvicorn_app", ) + +_EXPORTS = { + "AttestedTLS": (".attestation", "AttestedTLS"), + "PromptEncryptionASGIMiddleware": ( + ".asgi", + "PromptEncryptionASGIMiddleware", + ), + "PromptEncryptionWSGIMiddleware": ( + ".wsgi", + "PromptEncryptionWSGIMiddleware", + ), + "KeyManager": (".keys", "KeyManager"), + "TokenManager": (".token", "TokenManager"), + "run_gunicorn_app": (".wsgi", "run_gunicorn_app"), + "run_uvicorn_app": (".asgi", "run_uvicorn_app"), +} + + +def __getattr__(name: str): + """Resolves a public name to its module on first access.""" + try: + module_name, attribute_name = _EXPORTS[name] + except KeyError as e: + raise AttributeError(name) from e + value = getattr(importlib.import_module(module_name, __name__), attribute_name) + globals()[name] = value + return value diff --git a/src/prompt_encryption_sdk/server/asgi.py b/src/prompt_encryption_sdk/server/asgi.py index 6a1e704..98bea54 100644 --- a/src/prompt_encryption_sdk/server/asgi.py +++ b/src/prompt_encryption_sdk/server/asgi.py @@ -18,6 +18,7 @@ import json import logging import weakref +from prompt_encryption_sdk import attestation as attestation_protocol from prompt_encryption_sdk.proto import attestation_pb2 from prompt_encryption_sdk.server import attestation from prompt_encryption_sdk.server import keys @@ -79,6 +80,20 @@ async def __call__(self, scope, receive, send): if scope["path"] == "/_attest-connection": try: await self.handle_attestation(scope, receive, send) + except attestation_protocol.AttestationReplayError as e: + # The session already proved itself against an EKM that cannot change. + # Revoke it and close the connection rather than re-attesting it. + logging.warning("Dropping re-attestation on an attested session: %r", e) + replayed_socket = scope.get("extensions", {}).get("tls_socket") + if replayed_socket is not None: + self._attested_sockets.discard(replayed_socket) + error_response = responses.JSONResponse( + {"error": str(e)}, + status_code=http.HTTPStatus.FORBIDDEN, + # Tells the ASGI server to tear down this TLS session. + headers={"Connection": "close"}, + ) + await error_response(scope, receive, send) except Exception as e: # pylint: disable=broad-except logging.exception("Error during handling attest connection request") status_code = ( @@ -151,7 +166,13 @@ async def handle_attestation(self, scope, receive, send): def run_uvicorn_app( - app, *, key_manager=None, token_manager=None, **kwargs + app, + *, + key_manager=None, + token_manager=None, + client_policy: attestation_pb2.AttestationPolicy | None = None, + require_mutual_attestation: bool = False, + **kwargs, ) -> None: """Runs a uvicorn app with PromptEncryptionASGIMiddleware. @@ -161,6 +182,9 @@ def run_uvicorn_app( one will be created. token_manager: An optional token.TokenManager instance. If not provided, a new one will be created using the key_manager. + client_policy: Policy used to verify confidential clients. Providing it + enables mutual attestation. + require_mutual_attestation: Reject server-only clients. **kwargs: Additional keyword arguments to pass to uvicorn.run. """ if key_manager is None: @@ -171,6 +195,10 @@ def run_uvicorn_app( if "log_config" not in kwargs: kwargs["log_level"] = kwargs.get("log_level", "info") with token_manager: - atls = attestation.AttestedTLS(token_manager) + atls = attestation.AttestedTLS( + token_manager, + client_policy=client_policy, + require_mutual_attestation=require_mutual_attestation, + ) protected_app = PromptEncryptionASGIMiddleware(app, atls) uvicorn.run(protected_app, **kwargs) diff --git a/src/prompt_encryption_sdk/server/attestation.py b/src/prompt_encryption_sdk/server/attestation.py index e8a4b52..461cb06 100644 --- a/src/prompt_encryption_sdk/server/attestation.py +++ b/src/prompt_encryption_sdk/server/attestation.py @@ -14,24 +14,84 @@ """Attested TLS logic for Prompt Encryption SDK.""" -import hashlib +import threading from typing import Any +import weakref from absl import logging +from prompt_encryption_sdk import attestation as attestation_protocol +from prompt_encryption_sdk.client import validator from prompt_encryption_sdk.ekm import exporter from prompt_encryption_sdk.proto import attestation_pb2 from prompt_encryption_sdk.server import token _EKM_LABEL = b"EXPORTER-Prompt-Encryption-SDK" +_EKM_LENGTH = 32 + class AttestedTLS: - """Handles AttestConnection logic.""" + """Handles AttestConnection logic. + + Two modes share one round trip: + + * Server-only (the default) returns the server's proof, bound to an EKM + derived with the client's nonce as context. + * Mutual verifies the client's proof carried in the request and, only if it + passes, returns the server's proof. Both proofs bind to the same + no-context EKM: the TLS session is itself the freshness challenge, so + neither side needs to contribute a nonce. + + Because a mutual proof binds to a value that is constant for the session, + a second attestation on an already-mutually-attested session would prove + nothing new. Such a session is not re-attested; `attest_connection` raises + `AttestationReplayError` and the caller drops the connection. + """ + + def __init__( + self, + token_manager: token.TokenManager, + *, + client_policy: attestation_pb2.AttestationPolicy | None = None, + require_mutual_attestation: bool = False, + attestation_validator_cls=validator.AttestationValidator, + ): + """Initializes the handler. + + Args: + token_manager: Supplies this server's keys and attestation token. + client_policy: Policy a confidential client's evidence must satisfy. + Providing it enables mutual attestation. + require_mutual_attestation: Reject server-only clients. + attestation_validator_cls: Dependency injection for the validator used + to check client proofs. + + Raises: + ValueError: If mutual attestation is required without a client policy. + """ + if require_mutual_attestation and client_policy is None: + raise ValueError( + "client_policy is required when mutual attestation is required." + ) - def __init__(self, token_manager: token.TokenManager): self.token_manager = token_manager + self._prover = attestation_protocol.AttestationProver(token_manager) + self._client_policy = client_policy + self._require_mutual_attestation = require_mutual_attestation + self._attestation_validator_cls = attestation_validator_cls + self._client_validator = None + # Guards both the lazily built validator and the attested-session set, + # which are shared across concurrently served connections. + self._lock = threading.Lock() + # Sessions that already completed mutual attestation. Weak so entries + # disappear with the TLS socket. + self._mutually_attested = weakref.WeakSet() def __repr__(self): - return f"AttestedTLS(token_manager={self.token_manager!r})" + return ( + f"AttestedTLS(token_manager={self.token_manager!r}," + f" client_policy={self._client_policy!r}," + f" require_mutual_attestation={self._require_mutual_attestation!r})" + ) def attest_connection( self, @@ -46,6 +106,9 @@ def attest_connection( the TLS session material and hash of the attestation token with the private key and includes the signature in the response. + In mutual mode it first verifies the client proof carried in the request, + and only produces the server proof if that check passes. + Args: request: The AttestConnectionRequest message. ssl_obj: The SSL object from the TLS connection. @@ -55,12 +118,128 @@ def attest_connection( the server's public key and signed TLS session material. Raises: - ValueError: If no required_verifier_type is specified or if an unsupported - verifier type is requested. + AttestationReplayError: If this TLS session already completed mutual + attestation. + ValueError: If the request is malformed or asks for an unsupported + verifier, mode, or protocol version. RuntimeError: If EKM extraction fails. """ - if not request.required_verifier_type: - raise ValueError("At least one required_verifier_type must be specified.") + attestation_protocol.validate_required_verifiers( + request.required_verifier_type + ) + + with self._lock: + already_attested = ssl_obj in self._mutually_attested + if already_attested: + raise attestation_protocol.AttestationReplayError( + "This TLS session is already mutually attested. Re-attestation over" + " the same session proves no new freshness; establish a new" + " connection instead." + ) + + if request.mode == attestation_pb2.ATTESTATION_MODE_MUTUAL: + return self._attest_mutual(request, ssl_obj=ssl_obj) + + if request.mode != attestation_pb2.ATTESTATION_MODE_SERVER_ONLY: + raise ValueError(f"Unsupported attestation mode: {request.mode}") + if self._require_mutual_attestation: + raise ValueError("This server requires mutual attestation.") + if request.HasField("client_attestation"): + raise ValueError( + "client_attestation is valid only for mutual attestation." + ) + return self._attest_server_only(request, ssl_obj=ssl_obj) + + def _attest_server_only( + self, + request: attestation_pb2.AttestConnectionRequest, + *, + ssl_obj: Any, + ) -> attestation_pb2.AttestConnectionResponse: + """Runs the server-only exchange, using the client nonce as EKM context.""" + tls_ekm = self._export_ekm(ssl_obj, context=request.nonce) + return attestation_protocol.response_from_proof( + self._prover.create_proof(tls_ekm) + ) + + def _attest_mutual( + self, + request: attestation_pb2.AttestConnectionRequest, + *, + ssl_obj: Any, + ) -> attestation_pb2.AttestConnectionResponse: + """Verifies the client's proof, then answers with the server's own.""" + if ( + request.protocol_version + != attestation_protocol.MUTUAL_ATTESTATION_PROTOCOL_VERSION + ): + raise ValueError( + "Unsupported mutual attestation protocol version:" + f" {request.protocol_version}" + ) + if self._client_policy is None: + raise ValueError("Server is not configured to verify client evidence.") + if request.nonce: + # Mutual mode binds to the bare session EKM. Honoring a caller-supplied + # context here would let the peers derive different bindings. + raise ValueError("nonce must be empty for mutual attestation.") + if not request.HasField("client_attestation"): + raise ValueError("Client attestation proof is missing.") + + tls_ekm = self._export_ekm(ssl_obj, context=None) + self._client_verifier().validate_proof( + request.client_attestation, + tls_ekm, + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_CLIENT, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ) + + proof = self._prover.create_proof( + tls_ekm, + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_SERVER, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ) + with self._lock: + self._mutually_attested.add(ssl_obj) + return attestation_protocol.response_from_proof( + proof, + protocol_version=( + attestation_protocol.MUTUAL_ATTESTATION_PROTOCOL_VERSION + ), + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + mutual_attestation_complete=True, + ) + + def _client_verifier(self): + """Returns the shared client-proof validator, building it on first use. + + Construction reaches out for OIDC discovery, so it is deferred until a + mutual request actually arrives and then reused across connections. + """ + if self._client_validator is None: + with self._lock: + if self._client_validator is None: + self._client_validator = self._attestation_validator_cls( + self._client_policy + ) + return self._client_validator + + @staticmethod + def _export_ekm(ssl_obj: Any, *, context: bytes | None) -> bytes: + """Extracts EKM using the standard library, falling back to the C extension. + + Args: + ssl_obj: The SSL object from the TLS connection. + context: RFC 5705 context, or None for no context at all. Both peers + must agree, since an empty context and no context derive different + keying material. + + Returns: + The exported keying material. + + Raises: + RuntimeError: If both extraction paths fail. + """ ekm_bytes = None first_ekm_exception = None # Attempt to extract EKM using the standard library. @@ -70,9 +249,9 @@ def attest_connection( if hasattr(ssl_obj, "export_keying_material"): try: ekm_bytes = ssl_obj.export_keying_material( - _EKM_LABEL, 32, context=request.nonce + _EKM_LABEL, _EKM_LENGTH, context=context ) - except Exception as e: + except Exception as e: # pylint: disable=broad-except # If export_keying_material fails, we will try to extract EKM using # exporter.export_keying_material. logging.exception("Failed to extract EKM using export_keying_material.") @@ -82,61 +261,19 @@ def attest_connection( target_obj = getattr(ssl_obj, "_sslobj", ssl_obj) ekm_bytes = exporter.export_keying_material( sock=target_obj, - length=32, + length=_EKM_LENGTH, label=_EKM_LABEL, - context=request.nonce, + context=context, ) - if ekm_bytes is None: - if first_ekm_exception: - raise RuntimeError( - "EKM extraction failed. The initial attempt using" - " ssl_obj.export_keying_material failed." - ) from first_ekm_exception - else: - raise RuntimeError( - "EKM extraction failed. Both ssl_obj.export_keying_material and" - " the fallback exporter failed to extract keying material." - ) - - ecdsa_pub, pqc_pub, attestation_token = ( - self.token_manager.get_identity_snapshot() - ) - token_hash_bytes = hashlib.sha256(attestation_token).digest() - ekm_hash_bytes = hashlib.sha256(ekm_bytes).digest() - payload = attestation_pb2.SessionSignaturePayload( - ekm_hash=ekm_hash_bytes, token_hash=token_hash_bytes - ) - ecdsa_signature = self.token_manager.key_manager.sign_payload( - payload.SerializeToString() - ) - pqc_signature = self.token_manager.key_manager.sign_payload_mldsa( - payload.SerializeToString() - ) - - for verifier_type in request.required_verifier_type: - if verifier_type != attestation_pb2.VerifierType.VERIFIER_TYPE_GCA: - raise ValueError( - "Unsupported verifier types requested:" - f" {request.required_verifier_type}" - ) - - response = attestation_pb2.AttestConnectionResponse( - evidence=[ - attestation_pb2.AttestationEvidence( - verifier_type=attestation_pb2.VerifierType.VERIFIER_TYPE_GCA, - gca_bundle=attestation_pb2.GcaTrustBundle( - attestation_token=attestation_token.decode("utf-8") - ), - ) - ], - instance_public_key=attestation_pb2.EcdsaP256PublicKey( - key_bytes=ecdsa_pub - ), - session_signature=ecdsa_signature, - pqc_public_key=attestation_pb2.MlDsaPublicKey( - serialized_public_keyset=pqc_pub - ), - pqc_session_signature=pqc_signature, + if ekm_bytes is not None: + return ekm_bytes + if first_ekm_exception: + raise RuntimeError( + "EKM extraction failed. The initial attempt using" + " ssl_obj.export_keying_material failed." + ) from first_ekm_exception + raise RuntimeError( + "EKM extraction failed. Both ssl_obj.export_keying_material and" + " the fallback exporter failed to extract keying material." ) - return response diff --git a/src/prompt_encryption_sdk/server/wsgi.py b/src/prompt_encryption_sdk/server/wsgi.py index d93a2f9..c5a2ab3 100644 --- a/src/prompt_encryption_sdk/server/wsgi.py +++ b/src/prompt_encryption_sdk/server/wsgi.py @@ -19,6 +19,7 @@ from typing import Any, Iterable import weakref +from prompt_encryption_sdk import attestation as attestation_protocol from prompt_encryption_sdk.proto import attestation_pb2 from prompt_encryption_sdk.server import attestation from prompt_encryption_sdk.server import keys @@ -121,6 +122,7 @@ def handle_attestation( Returns: An iterable of bytes representing the response body. """ + ssl_obj = environ.get("prompt_encryption.socket") try: body = self._parse_request(environ) @@ -140,9 +142,7 @@ def handle_attestation( ignore_unknown_fields=True, ) - # Get socket from environ (injected by our patched gunicorn/worker) - ssl_obj = environ.get("prompt_encryption.socket") - + # Socket comes from environ (injected by our patched gunicorn/worker). if not ssl_obj: raise RuntimeError( "TLS Socket not found. Server must be started via" @@ -168,6 +168,24 @@ def handle_attestation( ) return [response_data] + except attestation_protocol.AttestationReplayError as e: + # The session already proved itself against an EKM that cannot change. + # Revoke it and close the connection rather than re-attesting it. + logging.warning("Dropping re-attestation on an attested session: %r", e) + if ssl_obj is not None: + self._attested_sockets.discard(ssl_obj) + error_json = json.dumps({"error": str(e)}).encode("utf-8") + start_response( + "403 Forbidden", + [ + ("Content-Type", "application/json"), + ("Content-Length", str(len(error_json))), + # Tells the WSGI server to tear down this TLS session. + ("Connection", "close"), + ], + ) + return [error_json] + except json_format.ParseError: logging.exception("Error parsing AttestConnectionRequest") error_json = json.dumps({"error": "Invalid request"}).encode("utf-8") @@ -222,6 +240,8 @@ def run_gunicorn_app( workers: int = 1, ssl_certfile: str | None = None, ssl_keyfile: str | None = None, + client_policy: attestation_pb2.AttestationPolicy | None = None, + require_mutual_attestation: bool = False, attested_tls_cls: type[attestation.AttestedTLS] = attestation.AttestedTLS, standalone_app_cls: type[gunicorn.app.base.BaseApplication] = ( _StandaloneApplication @@ -239,6 +259,9 @@ def run_gunicorn_app( workers: Number of workers. ssl_certfile: Path to SSL certificate file. ssl_keyfile: Path to SSL key file. + client_policy: Policy used to verify confidential clients. Providing it + enables mutual attestation. + require_mutual_attestation: Reject server-only clients. attested_tls_cls: Dependency injection for AttestedTLS class. standalone_app_cls: Dependency injection for StandaloneApplication class. **kwargs: Additional keyword arguments to pass to Gunicorn options. @@ -249,7 +272,11 @@ def run_gunicorn_app( token_manager = token.TokenManager(key_manager=key_manager) with token_manager: - atls = attested_tls_cls(token_manager) + atls = attested_tls_cls( + token_manager, + client_policy=client_policy, + require_mutual_attestation=require_mutual_attestation, + ) middleware_app = PromptEncryptionWSGIMiddleware(app, atls) options = { diff --git a/tests/attestation_protocol_test.py b/tests/attestation_protocol_test.py new file mode 100644 index 0000000..9d5af3d --- /dev/null +++ b/tests/attestation_protocol_test.py @@ -0,0 +1,237 @@ +# Copyright 2026 The Prompt Encryption SDK Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the shared attestation proof primitives.""" + +import hashlib +from unittest import mock + +from absl.testing import absltest +from absl.testing import parameterized +from prompt_encryption_sdk import attestation as attestation_protocol +from prompt_encryption_sdk.proto import attestation_pb2 + + +_EKM = b"e" * 32 +_TOKEN = b"attestation-token" + + +class ValidateRequiredVerifiersTest(parameterized.TestCase): + + def test_accepts_gca(self): + attestation_protocol.validate_required_verifiers( + [attestation_pb2.VERIFIER_TYPE_GCA] + ) + + def test_rejects_empty(self): + with self.assertRaisesRegex( + ValueError, "At least one required_verifier_type must be specified." + ): + attestation_protocol.validate_required_verifiers([]) + + def test_rejects_unsupported(self): + with self.assertRaisesRegex(ValueError, "Unsupported verifier types"): + attestation_protocol.validate_required_verifiers( + [attestation_pb2.VERIFIER_TYPE_UNSPECIFIED] + ) + + def test_rejects_unsupported_mixed_with_supported(self): + with self.assertRaisesRegex(ValueError, "Unsupported verifier types"): + attestation_protocol.validate_required_verifiers([ + attestation_pb2.VERIFIER_TYPE_GCA, + attestation_pb2.VERIFIER_TYPE_UNSPECIFIED, + ]) + + +class BuildSignaturePayloadTest(parameterized.TestCase): + + def test_server_only_payload_is_wire_compatible(self): + """The pre-mutual-attestation signed bytes must not change.""" + legacy = attestation_pb2.SessionSignaturePayload( + ekm_hash=hashlib.sha256(_EKM).digest(), + token_hash=hashlib.sha256(_TOKEN).digest(), + ).SerializeToString() + + self.assertEqual( + attestation_protocol.build_signature_payload( + tls_ekm=_EKM, attestation_token=_TOKEN + ), + legacy, + ) + + def test_mutual_payload_carries_role_mode_and_version(self): + payload = attestation_pb2.SessionSignaturePayload.FromString( + attestation_protocol.build_signature_payload( + tls_ekm=_EKM, + attestation_token=_TOKEN, + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_CLIENT, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ) + ) + + self.assertEqual( + payload, + attestation_pb2.SessionSignaturePayload( + ekm_hash=hashlib.sha256(_EKM).digest(), + token_hash=hashlib.sha256(_TOKEN).digest(), + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_CLIENT, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + protocol_version=( + attestation_protocol.MUTUAL_ATTESTATION_PROTOCOL_VERSION + ), + ), + ) + + def test_roles_produce_distinct_payloads_for_the_same_ekm(self): + """Role separation is the only thing preventing proof reflection.""" + server = attestation_protocol.build_signature_payload( + tls_ekm=_EKM, + attestation_token=_TOKEN, + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_SERVER, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ) + client = attestation_protocol.build_signature_payload( + tls_ekm=_EKM, + attestation_token=_TOKEN, + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_CLIENT, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ) + + self.assertNotEqual(server, client) + + def test_modes_produce_distinct_payloads_for_the_same_ekm(self): + mutual = attestation_protocol.build_signature_payload( + tls_ekm=_EKM, + attestation_token=_TOKEN, + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_SERVER, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ) + server_only = attestation_protocol.build_signature_payload( + tls_ekm=_EKM, + attestation_token=_TOKEN, + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_SERVER, + mode=attestation_pb2.ATTESTATION_MODE_SERVER_ONLY, + ) + + self.assertNotEqual(mutual, server_only) + + def test_payload_is_deterministic(self): + self.assertEqual( + attestation_protocol.build_signature_payload( + tls_ekm=_EKM, + attestation_token=_TOKEN, + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_CLIENT, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ), + attestation_protocol.build_signature_payload( + tls_ekm=_EKM, + attestation_token=_TOKEN, + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_CLIENT, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ), + ) + + +class AttestationProverTest(parameterized.TestCase): + + def _identity(self, token=_TOKEN): + identity = mock.MagicMock() + identity.get_identity_snapshot.return_value = ( + b"ecdsa-public", + b"mldsa-public", + token, + ) + identity.key_manager.sign_payload.return_value = b"ecdsa-signature" + identity.key_manager.sign_payload_mldsa.return_value = b"mldsa-signature" + return identity + + def test_create_proof_signs_the_canonical_payload_with_both_keys(self): + identity = self._identity() + + proof = attestation_protocol.AttestationProver(identity).create_proof( + _EKM, + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_CLIENT, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ) + + expected_payload = attestation_protocol.build_signature_payload( + tls_ekm=_EKM, + attestation_token=_TOKEN, + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_CLIENT, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ) + identity.key_manager.sign_payload.assert_called_once_with(expected_payload) + identity.key_manager.sign_payload_mldsa.assert_called_once_with( + expected_payload + ) + self.assertEqual(proof.session_signature, b"ecdsa-signature") + self.assertEqual(proof.pqc_session_signature, b"mldsa-signature") + self.assertEqual(proof.instance_public_key.key_bytes, b"ecdsa-public") + self.assertEqual( + proof.pqc_public_key.serialized_public_keyset, b"mldsa-public" + ) + self.assertEqual( + proof.evidence[0].verifier_type, attestation_pb2.VERIFIER_TYPE_GCA + ) + self.assertEqual( + proof.evidence[0].gca_bundle.attestation_token, _TOKEN.decode("utf-8") + ) + + def test_create_proof_without_token_fails(self): + with self.assertRaisesRegex(RuntimeError, "no token"): + attestation_protocol.AttestationProver( + self._identity(token=b"") + ).create_proof(_EKM) + + +class ProofResponseConversionTest(parameterized.TestCase): + + def _proof(self): + return attestation_pb2.AttestationProof( + evidence=[ + attestation_pb2.AttestationEvidence( + verifier_type=attestation_pb2.VERIFIER_TYPE_GCA, + gca_bundle=attestation_pb2.GcaTrustBundle( + attestation_token="token" + ), + ) + ], + instance_public_key=attestation_pb2.EcdsaP256PublicKey( + key_bytes=b"ecdsa" + ), + session_signature=b"sig", + pqc_public_key=attestation_pb2.MlDsaPublicKey( + serialized_public_keyset=b"mldsa" + ), + pqc_session_signature=b"pqc-sig", + ) + + def test_round_trips_through_the_response_shape(self): + proof = self._proof() + + response = attestation_protocol.response_from_proof( + proof, + protocol_version=1, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + mutual_attestation_complete=True, + ) + + self.assertTrue(response.mutual_attestation_complete) + self.assertEqual( + attestation_protocol.proof_from_response(response), proof + ) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/client/client_test.py b/tests/client/client_test.py index 94705a7..398b889 100644 --- a/tests/client/client_test.py +++ b/tests/client/client_test.py @@ -15,6 +15,7 @@ from unittest import mock from absl.testing import parameterized +from prompt_encryption_sdk import attestation as attestation_protocol from prompt_encryption_sdk.client import client from prompt_encryption_sdk.client import connection import requests @@ -81,6 +82,8 @@ def test_init_poolmanager_initializes_attested_pool(self, mock_pool_manager): maxsize=maxsize, block=block, policy=self.mock_policy, + mutual_attestation=False, + attestation_prover=None, revalidation_timeout=self.timeout, extra=extra_kwarg, ) @@ -208,5 +211,67 @@ def test_post_classmethod(self, mock_session_method): mock_session_method.return_value.__exit__.assert_called_once() +class MutualPromptEncryptionClientTest(parameterized.TestCase): + """Covers the mutual-attestation configuration surface.""" + + def setUp(self): + super().setUp() + self.mock_policy = mock.create_autospec( + client.attestation_pb2.AttestationPolicy + ) + self.token_manager = mock.MagicMock() + + def test_mutual_requires_a_client_token_manager(self): + with self.assertRaisesRegex( + ValueError, "client_token_manager is required" + ): + client.PromptEncryptionClient(self.mock_policy, mutual_attestation=True) + + def test_mutual_rejects_a_revalidation_timeout(self): + """A mutually attested session is attested exactly once.""" + with self.assertRaisesRegex( + ValueError, "revalidation_timeout is not supported" + ): + client.PromptEncryptionClient( + self.mock_policy, + mutual_attestation=True, + client_token_manager=self.token_manager, + revalidation_timeout=600, + ) + + @mock.patch.object(requests, "Session", autospec=True) + @mock.patch.object(client, "AttestedHTTPSAdapter", autospec=True) + def test_session_wires_a_prover_over_the_client_identity( + self, mock_adapter_cls, mock_session_cls + ): + del mock_session_cls + sdk = client.PromptEncryptionClient( + self.mock_policy, + mutual_attestation=True, + client_token_manager=self.token_manager, + ) + + sdk.session() + + adapter_kwargs = mock_adapter_cls.call_args.kwargs + self.assertTrue(adapter_kwargs["mutual_attestation"]) + self.assertIsInstance( + adapter_kwargs["attestation_prover"], + attestation_protocol.AttestationProver, + ) + + @mock.patch.object(requests, "Session", autospec=True) + @mock.patch.object(client, "AttestedHTTPSAdapter", autospec=True) + def test_server_only_session_does_not_mention_mutual_attestation( + self, mock_adapter_cls, mock_session_cls + ): + del mock_session_cls + sdk = client.PromptEncryptionClient(self.mock_policy) + + sdk.session() + + self.assertNotIn("mutual_attestation", mock_adapter_cls.call_args.kwargs) + + if __name__ == "__main__": googletest.main() diff --git a/tests/client/connection_test.py b/tests/client/connection_test.py index 16054ce..1a322e5 100644 --- a/tests/client/connection_test.py +++ b/tests/client/connection_test.py @@ -21,12 +21,14 @@ from absl.testing import absltest from absl.testing import parameterized +from prompt_encryption_sdk import attestation as attestation_protocol from prompt_encryption_sdk.client import connection from prompt_encryption_sdk.client import constants from prompt_encryption_sdk.client import exceptions from prompt_encryption_sdk.client import validator from prompt_encryption_sdk.ekm import exporter as ekm_exporter from prompt_encryption_sdk.proto import attestation_pb2 +from google.protobuf import json_format import urllib3 from urllib3.connection import HTTPSConnection from urllib3.connectionpool import HTTPConnectionPool @@ -369,5 +371,147 @@ def test_new_pool_http_returns_standard_pool(self): self.assertIsInstance(pool, HTTPConnectionPool) +class MutualAttestedHTTPSConnectionTest(parameterized.TestCase): + """Covers the client half of the single-round-trip mutual exchange.""" + + def setUp(self): + super().setUp() + self.prover = mock.MagicMock() + self.client_proof = attestation_pb2.AttestationProof( + session_signature=b"client-sig" + ) + self.prover.create_proof.return_value = self.client_proof + + self.validator_cls = mock.MagicMock() + self.validator_inst = self.validator_cls.return_value + + self.export_ekm = mock.MagicMock(return_value=b"session-ekm") + + self.conn = connection.AttestedHTTPSConnection( + host="example.com", + port=443, + policy=attestation_pb2.AttestationPolicy(), + mutual_attestation=True, + attestation_prover=self.prover, + ekm_exporter_fn=self.export_ekm, + attestation_validator_cls=self.validator_cls, + ) + + def _stub_response(self, response: attestation_pb2.AttestConnectionResponse): + mock_response = mock.create_autospec(urllib3.HTTPResponse, instance=True) + mock_response.status = 200 + mock_response.headers = {"Content-Type": "application/x-protobuf"} + mock_response.data = response.SerializeToString() + self.conn.getresponse = mock.MagicMock(return_value=mock_response) + return mock_response + + def _mutual_response(self, **overrides): + fields = dict( + protocol_version=( + attestation_protocol.MUTUAL_ATTESTATION_PROTOCOL_VERSION + ), + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + mutual_attestation_complete=True, + ) + fields.update(overrides) + return attestation_pb2.AttestConnectionResponse(**fields) + + def test_requires_a_prover_when_mutual_attestation_is_enabled(self): + with self.assertRaisesRegex(ValueError, "attestation_prover is required"): + connection.AttestedHTTPSConnection( + host="example.com", port=443, mutual_attestation=True + ) + + @mock.patch.object(HTTPSConnection, "request", autospec=True) + def test_one_round_trip_carries_the_client_proof(self, mock_super_request): + self._stub_response(self._mutual_response()) + + self.conn._perform_attestation_handshake() + + with self.subTest("EkmExportedWithoutContext"): + self.export_ekm.assert_called_once_with( + self.conn.sock, constants.EKM_LENGTH, constants.EKM_LABEL + ) + with self.subTest("ClientProvesItselfAsTheClientRole"): + self.prover.create_proof.assert_called_once_with( + b"session-ekm", + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_CLIENT, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ) + with self.subTest("ExactlyOneFlight"): + mock_super_request.assert_called_once() + with self.subTest("RequestCarriesTheProofAndNoNonce"): + sent = json_format.Parse( + mock_super_request.call_args.kwargs["body"], + attestation_pb2.AttestConnectionRequest(), + ) + self.assertEqual(sent.client_attestation, self.client_proof) + self.assertEqual(sent.mode, attestation_pb2.ATTESTATION_MODE_MUTUAL) + self.assertEqual(sent.protocol_version, 1) + self.assertEmpty(sent.nonce) + with self.subTest("ServerProofValidatedAgainstTheSameEkm"): + self.validator_inst.validate.assert_called_once_with( + mock.ANY, + b"session-ekm", + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_SERVER, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ) + with self.subTest("ValidatorResourcesReleased"): + self.validator_inst.close.assert_called_once() + + @parameterized.named_parameters( + ("server_only_mode", {"mode": attestation_pb2.ATTESTATION_MODE_SERVER_ONLY}), + ("wrong_protocol_version", {"protocol_version": 2}), + ("not_complete", {"mutual_attestation_complete": False}), + ) + @mock.patch.object(HTTPSConnection, "request", autospec=True) + def test_downgrade_is_refused(self, overrides, mock_super_request): + del mock_super_request + self._stub_response(self._mutual_response(**overrides)) + + with self.assertRaisesRegex( + exceptions.AttestationHandshakeError, "refusing server-only downgrade" + ): + self.conn._perform_attestation_handshake() + + self.validator_inst.validate.assert_not_called() + + @mock.patch.object(HTTPSConnection, "request", autospec=True) + def test_invalid_server_proof_propagates(self, mock_super_request): + del mock_super_request + self._stub_response(self._mutual_response()) + self.validator_inst.validate.side_effect = ( + exceptions.AttestationVerificationError("bad server proof") + ) + + with self.assertRaises(exceptions.AttestationVerificationError): + self.conn._perform_attestation_handshake() + + def test_mutual_sessions_are_never_revalidated(self): + self.conn.is_attested = True + self.conn._last_attestation_time = time.time() - 10_000 + + self.assertFalse(self.conn._should_revalidate()) + + @mock.patch.object(HTTPSConnection, "request", autospec=True) + def test_request_does_not_revalidate_a_mutual_session( + self, mock_super_request + ): + self.conn.is_attested = True + self.conn._last_attestation_time = time.time() - 10_000 + + with mock.patch.object(self.conn, "revalidate_session") as mock_reval: + self.conn.request("GET", "/infer") + + mock_reval.assert_not_called() + mock_super_request.assert_called_once() + + def test_explicit_revalidation_is_refused(self): + with self.assertRaisesRegex( + exceptions.PromptEncryptionError, "cannot be revalidated" + ): + self.conn.revalidate_session() + + if __name__ == "__main__": absltest.main() diff --git a/tests/mutual_attestation_test.py b/tests/mutual_attestation_test.py new file mode 100644 index 0000000..e358da7 --- /dev/null +++ b/tests/mutual_attestation_test.py @@ -0,0 +1,325 @@ +# Copyright 2026 The Prompt Encryption SDK Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Mock-TEE tests for the single-round-trip mutual attestation exchange. + +Real ECDSA P-256 and ML-DSA-65 keys are generated and every signature is +genuinely produced and verified. Only the Google Cloud Attestation token +issuance and its OIDC verification are faked, since those require a real +Confidential Space workload. +""" + +import hashlib +import pathlib + +from absl.testing import absltest +from prompt_encryption_sdk import attestation as attestation_protocol +from prompt_encryption_sdk.client import exceptions +from prompt_encryption_sdk.client import validator +from prompt_encryption_sdk.proto import attestation_pb2 +from prompt_encryption_sdk.server import attestation as server_attestation +from prompt_encryption_sdk.server import keys + + +_EKM_LABEL = b"EXPORTER-Prompt-Encryption-SDK" +_EKM_LENGTH = 32 + + +class _Identity: + """A stand-in TokenManager holding real keys and a fake GCA token.""" + + def __init__(self, key_manager: keys.KeyManager, token: bytes): + self.key_manager = key_manager + self._token = token + + def get_identity_snapshot(self) -> tuple[bytes, bytes, bytes]: + return ( + self.key_manager.get_current_public_key(), + self.key_manager.get_current_pqc_public_key(), + self._token, + ) + + +class _FakeOIDCValidator: + """Returns pre-canned claims instead of verifying a real GCA token.""" + + def __init__(self, claims_by_token): + self._claims_by_token = claims_by_token + + def validate_token(self, token: str): + try: + return self._claims_by_token[token] + except KeyError as e: + raise exceptions.AttestationVerificationError("Unknown token.") from e + + def close(self): + pass + + +class _MockTLSSocket: + """Models both peers deriving identical EKM from one TLS session.""" + + def __init__(self, secret: bytes = b"mock-tls-master-secret"): + self._secret = secret + + def export_keying_material(self, label, length, *, context=None): + seed = self._secret + label + (b"" if context is None else b"\x01" + context) + return hashlib.sha256(seed).digest()[:length] + + +def _claims_for(identity: _Identity): + """Builds GCA claims that bind both public keys via eat_nonce.""" + return { + "eat_nonce": [ + keys.calculate_fingerprint( + identity.key_manager.get_current_public_key() + ), + keys.calculate_fingerprint( + identity.key_manager.get_current_pqc_public_key() + ), + ], + "submods": {"container": {}, "gce": {}}, + } + + +class MutualAttestationTest(absltest.TestCase): + + def setUp(self): + super().setUp() + root = pathlib.Path(self.create_tempdir().full_path) + self.server_identity = self._make_identity(root / "server", b"server-token") + self.client_identity = self._make_identity(root / "client", b"client-token") + self.oidc = _FakeOIDCValidator({ + "server-token": _claims_for(self.server_identity), + "client-token": _claims_for(self.client_identity), + }) + self.policy = attestation_pb2.AttestationPolicy() + self.client_validator = validator.AttestationValidator( + self.policy, oidc_validator=self.oidc + ) + self.server = server_attestation.AttestedTLS( + self.server_identity, + client_policy=self.policy, + attestation_validator_cls=lambda policy: ( + validator.AttestationValidator(policy, oidc_validator=self.oidc) + ), + ) + self.tls_socket = _MockTLSSocket() + self.tls_ekm = self.tls_socket.export_keying_material( + _EKM_LABEL, _EKM_LENGTH + ) + self.client_prover = attestation_protocol.AttestationProver( + self.client_identity + ) + + def _make_identity(self, path: pathlib.Path, token: bytes) -> _Identity: + path.mkdir() + key_manager = keys.KeyManager( + private_key_path=path / "ecdsa-private.pem", + public_key_path=path / "ecdsa-public.pem", + pqc_private_key_path=path / "mldsa-private.bin", + pqc_public_key_path=path / "mldsa-public.bin", + ) + key_manager.generate_key_pair() + return _Identity(key_manager, token) + + def _client_request(self, **overrides): + fields = dict( + required_verifier_type=[attestation_pb2.VERIFIER_TYPE_GCA], + protocol_version=( + attestation_protocol.MUTUAL_ATTESTATION_PROTOCOL_VERSION + ), + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + client_attestation=self.client_prover.create_proof( + self.tls_ekm, + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_CLIENT, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ), + ) + fields.update(overrides) + return attestation_pb2.AttestConnectionRequest(**fields) + + def test_one_round_trip_verifies_both_peers(self): + response = self.server.attest_connection( + self._client_request(), ssl_obj=self.tls_socket + ) + + self.assertTrue(response.mutual_attestation_complete) + # The client verifies the server proof against the same session EKM. + self.client_validator.validate( + response, + self.tls_ekm, + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_SERVER, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ) + + def test_server_proof_cannot_be_reflected_as_a_client_proof(self): + """Both proofs bind to the same EKM, so only the role separates them.""" + response = self.server.attest_connection( + self._client_request(), ssl_obj=self.tls_socket + ) + + with self.assertRaises(exceptions.AttestationVerificationError): + self.client_validator.validate( + response, + self.tls_ekm, + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_CLIENT, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ) + + def test_client_proof_cannot_be_replayed_to_the_server_as_a_server_proof(self): + client_proof = self.client_prover.create_proof( + self.tls_ekm, + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_CLIENT, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ) + + with self.assertRaises(exceptions.AttestationVerificationError): + self.client_validator.validate_proof( + client_proof, + self.tls_ekm, + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_SERVER, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ) + + def test_a_server_only_proof_is_not_accepted_as_a_mutual_one(self): + """Mode is signed, so an old-style proof cannot satisfy mutual mode.""" + server_only_proof = self.client_prover.create_proof(self.tls_ekm) + + with self.assertRaises(exceptions.AttestationVerificationError): + self.client_validator.validate_proof( + server_only_proof, + self.tls_ekm, + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_CLIENT, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ) + + def test_a_proof_from_another_tls_session_is_rejected(self): + """The EKM is the freshness challenge, so it must not transfer.""" + other_session_ekm = _MockTLSSocket( + b"a-different-tls-session" + ).export_keying_material(_EKM_LABEL, _EKM_LENGTH) + stolen_proof = attestation_protocol.AttestationProver( + self.client_identity + ).create_proof( + other_session_ekm, + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_CLIENT, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ) + + with self.assertRaises(exceptions.AttestationVerificationError): + self.server.attest_connection( + self._client_request(client_attestation=stolen_proof), + ssl_obj=self.tls_socket, + ) + + def test_a_client_outside_the_policy_is_rejected(self): + strict_policy = attestation_pb2.AttestationPolicy( + gce_instance=attestation_pb2.GceInstancePolicy(project_id="trusted") + ) + server = server_attestation.AttestedTLS( + self.server_identity, + client_policy=strict_policy, + require_mutual_attestation=True, + attestation_validator_cls=lambda policy: ( + validator.AttestationValidator(policy, oidc_validator=self.oidc) + ), + ) + + with self.assertRaises(exceptions.PolicyViolationError): + server.attest_connection(self._client_request(), ssl_obj=self.tls_socket) + + def test_a_client_with_an_unverifiable_token_is_rejected(self): + unknown = _Identity(self.client_identity.key_manager, b"forged-token") + forged_proof = attestation_protocol.AttestationProver(unknown).create_proof( + self.tls_ekm, + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_CLIENT, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ) + + with self.assertRaises(exceptions.AttestationVerificationError): + self.server.attest_connection( + self._client_request(client_attestation=forged_proof), + ssl_obj=self.tls_socket, + ) + + def test_a_client_key_not_bound_by_its_token_is_rejected(self): + """The GCA token's eat_nonce must cover the keys that signed the proof.""" + unbound_root = pathlib.Path(self.create_tempdir().full_path) / "unbound" + unbound = self._make_identity(unbound_root, b"client-token") + unbound_proof = attestation_protocol.AttestationProver( + unbound + ).create_proof( + self.tls_ekm, + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_CLIENT, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ) + + with self.assertRaisesRegex( + exceptions.AttestationVerificationError, "Instance Key binding failed" + ): + self.server.attest_connection( + self._client_request(client_attestation=unbound_proof), + ssl_obj=self.tls_socket, + ) + + def test_a_tampered_pqc_signature_is_rejected(self): + proof = self.client_prover.create_proof( + self.tls_ekm, + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_CLIENT, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ) + proof.pqc_session_signature = bytes( + [proof.pqc_session_signature[0] ^ 0xFF] + ) + proof.pqc_session_signature[1:] + + with self.assertRaisesRegex( + exceptions.AttestationVerificationError, "PQC session signature" + ): + self.server.attest_connection( + self._client_request(client_attestation=proof), + ssl_obj=self.tls_socket, + ) + + def test_a_tampered_ecdsa_signature_is_rejected(self): + proof = self.client_prover.create_proof( + self.tls_ekm, + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_CLIENT, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ) + proof.session_signature = bytes( + [proof.session_signature[0] ^ 0xFF] + ) + proof.session_signature[1:] + + with self.assertRaisesRegex( + exceptions.AttestationVerificationError, "session signature" + ): + self.server.attest_connection( + self._client_request(client_attestation=proof), + ssl_obj=self.tls_socket, + ) + + def test_the_session_is_attested_exactly_once(self): + self.server.attest_connection( + self._client_request(), ssl_obj=self.tls_socket + ) + + with self.assertRaises(attestation_protocol.AttestationReplayError): + self.server.attest_connection( + self._client_request(), ssl_obj=self.tls_socket + ) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/mutual_attestation_tls_integration_test.py b/tests/mutual_attestation_tls_integration_test.py new file mode 100644 index 0000000..ce875ed --- /dev/null +++ b/tests/mutual_attestation_tls_integration_test.py @@ -0,0 +1,363 @@ +# Copyright 2026 The Prompt Encryption SDK Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""End-to-end mutual attestation over a real TLS session. + +This drives the real client connection class against the real WSGI middleware +across a genuine TLS socket, so the Exported Keying Material both peers bind +to is produced by OpenSSL rather than a stub. Every ECDSA and ML-DSA signature +is really produced and really verified. + +Only Google Cloud Attestation is faked: token issuance and its OIDC +verification require a live Confidential Space workload. +""" + +import http.server +import io +import socket +import ssl +import sys +import threading +import urllib.parse + +from absl.testing import absltest +from prompt_encryption_sdk import attestation as attestation_protocol +from prompt_encryption_sdk.client import connection +from prompt_encryption_sdk.client import exceptions +from prompt_encryption_sdk.client import validator +from prompt_encryption_sdk.proto import attestation_pb2 +from prompt_encryption_sdk.server import attestation as server_attestation +from prompt_encryption_sdk.server import keys +from prompt_encryption_sdk.server import wsgi +from OpenSSL import crypto + + +_HOST = "localhost" +_RSA_KEY_SIZE = 2048 +_CERT_SERIAL_NUMBER = 1000 +_CERT_VALIDITY_SECONDS = 10 * 365 * 24 * 60 * 60 +_APP_BODY = b'{"result": "ok"}' + + +class _Identity: + """A stand-in TokenManager holding real keys and a fake GCA token.""" + + def __init__(self, key_manager: keys.KeyManager, token: bytes): + self.key_manager = key_manager + self._token = token + + def get_identity_snapshot(self) -> tuple[bytes, bytes, bytes]: + return ( + self.key_manager.get_current_public_key(), + self.key_manager.get_current_pqc_public_key(), + self._token, + ) + + +class _FakeOIDCValidator: + """Returns pre-canned claims instead of verifying a real GCA token.""" + + def __init__(self, claims_by_token): + self._claims_by_token = claims_by_token + + def validate_token(self, token: str): + try: + return self._claims_by_token[token] + except KeyError as e: + raise exceptions.AttestationVerificationError("Unknown token.") from e + + def close(self): + pass + + +def _claims_for(identity: _Identity): + """Builds GCA claims that bind both public keys via eat_nonce.""" + return { + "eat_nonce": [ + keys.calculate_fingerprint( + identity.key_manager.get_current_public_key() + ), + keys.calculate_fingerprint( + identity.key_manager.get_current_pqc_public_key() + ), + ], + "submods": {"container": {}, "gce": {}}, + } + + +class _WSGIBridgeHandler(http.server.BaseHTTPRequestHandler): + """Minimal keep-alive WSGI bridge that exposes the raw TLS socket. + + `wsgiref` closes the connection after every request, which would defeat the + whole point of post-handshake attestation, so this hand-rolled handler is + used instead. It mirrors what the Gunicorn worker patch does in production: + put the live SSL socket into the WSGI environ. + """ + + protocol_version = "HTTP/1.1" + + def log_message(self, format, *args): # pylint: disable=redefined-builtin + del format, args # Keep the test output clean. + + def do_GET(self): + self._run_wsgi_app() + + def do_POST(self): + self._run_wsgi_app() + + def _run_wsgi_app(self): + content_length = int(self.headers.get("Content-Length") or 0) + body = self.rfile.read(content_length) if content_length else b"" + environ = { + "REQUEST_METHOD": self.command, + "PATH_INFO": urllib.parse.urlparse(self.path).path, + "QUERY_STRING": urllib.parse.urlparse(self.path).query, + "SERVER_PROTOCOL": self.request_version, + "CONTENT_LENGTH": str(content_length), + "CONTENT_TYPE": self.headers.get("Content-Type", ""), + "wsgi.input": io.BytesIO(body), + "wsgi.errors": sys.stderr, + "wsgi.url_scheme": "https", + # The production Gunicorn worker injects this same key. + "prompt_encryption.socket": self.connection, + } + + captured = [] + + def start_response(status, headers, exc_info=None): + del exc_info + captured.append((status, headers)) + return lambda data: None + + chunks = self.server.wsgi_app(environ, start_response) + response_body = b"".join(chunks) + status, headers = captured[0] + code, _, reason = status.partition(" ") + self.send_response_only(int(code), reason) + for name, value in headers: + # send_header sets close_connection when it sees "Connection: close". + self.send_header(name, value) + self.end_headers() + self.wfile.write(response_body) + + +class _TLSWSGIServer(http.server.ThreadingHTTPServer): + """Wraps every accepted connection in TLS before handling it.""" + + daemon_threads = True + + def __init__(self, address, handler_cls, *, ssl_context, wsgi_app): + self.ssl_context = ssl_context + self.wsgi_app = wsgi_app + super().__init__(address, handler_cls) + + def get_request(self): + raw_sock, addr = super().get_request() + return self.ssl_context.wrap_socket(raw_sock, server_side=True), addr + + def handle_error(self, request, client_address): + del request, client_address # Client-side teardown races are expected. + + +class MutualAttestationOverRealTlsTest(absltest.TestCase): + + def setUp(self): + super().setUp() + self.server_identity = self._make_identity("server", b"server-token") + self.client_identity = self._make_identity("client", b"client-token") + oidc = _FakeOIDCValidator({ + "server-token": _claims_for(self.server_identity), + "client-token": _claims_for(self.client_identity), + }) + self.validator_factory = lambda policy: validator.AttestationValidator( + policy, oidc_validator=oidc + ) + self.policy = attestation_pb2.AttestationPolicy() + + self.attested_tls = server_attestation.AttestedTLS( + self.server_identity, + client_policy=self.policy, + require_mutual_attestation=True, + attestation_validator_cls=self.validator_factory, + ) + self.middleware = wsgi.PromptEncryptionWSGIMiddleware( + self._app, self.attested_tls + ) + self.port = self._start_server(self.middleware) + + def _app(self, unused_environ, start_response): + start_response( + "200 OK", + [ + ("Content-Type", "application/json"), + ("Content-Length", str(len(_APP_BODY))), + ], + ) + return [_APP_BODY] + + def _make_identity(self, name: str, token: bytes) -> _Identity: + temp_dir = self.create_tempdir(name).full_path + key_manager = keys.KeyManager( + private_key_path=f"{temp_dir}/ecdsa-private.pem", + public_key_path=f"{temp_dir}/ecdsa-public.pem", + pqc_private_key_path=f"{temp_dir}/mldsa-private.bin", + pqc_public_key_path=f"{temp_dir}/mldsa-public.bin", + ) + key_manager.generate_key_pair() + return _Identity(key_manager, token) + + def _start_server(self, wsgi_app) -> int: + pkey = crypto.PKey() + pkey.generate_key(crypto.TYPE_RSA, _RSA_KEY_SIZE) + cert = crypto.X509() + cert.get_subject().CN = _HOST + cert.set_serial_number(_CERT_SERIAL_NUMBER) + cert.gmtime_adj_notBefore(0) + cert.gmtime_adj_notAfter(_CERT_VALIDITY_SECONDS) + cert.set_issuer(cert.get_subject()) + cert.set_pubkey(pkey) + cert.sign(pkey, "sha256") + cert_file = self.create_tempfile( + "server.crt", content=crypto.dump_certificate(crypto.FILETYPE_PEM, cert) + ) + key_file = self.create_tempfile( + "server.key", content=crypto.dump_privatekey(crypto.FILETYPE_PEM, pkey) + ) + + ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + ssl_context.load_cert_chain( + certfile=cert_file.full_path, keyfile=key_file.full_path + ) + + server = _TLSWSGIServer( + (_HOST, 0), + _WSGIBridgeHandler, + ssl_context=ssl_context, + wsgi_app=wsgi_app, + ) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + self.addCleanup(thread.join, 5.0) + self.addCleanup(server.server_close) + self.addCleanup(server.shutdown) + return server.server_address[1] + + def _make_connection(self, **overrides): + kwargs = dict( + host=_HOST, + port=self.port, + policy=self.policy, + mutual_attestation=True, + attestation_prover=attestation_protocol.AttestationProver( + self.client_identity + ), + attestation_validator_cls=self.validator_factory, + cert_reqs="CERT_NONE", + assert_hostname=False, + ) + kwargs.update(overrides) + conn = connection.AttestedHTTPSConnection(**kwargs) + self.addCleanup(conn.close) + return conn + + def test_one_round_trip_attests_both_peers_and_authorizes_traffic(self): + conn = self._make_connection() + + conn.connect() + + with self.subTest("BothPeersAttested"): + self.assertTrue(conn.is_attested) + with self.subTest("ApplicationTrafficAuthorized"): + conn.request("GET", "/infer") + response = conn.getresponse() + self.assertEqual(response.status, 200) + self.assertEqual(response.data, _APP_BODY) + + def test_a_client_the_server_does_not_trust_is_rejected(self): + """A real policy mismatch must fail the connection, not downgrade it.""" + self.attested_tls._client_policy = attestation_pb2.AttestationPolicy( + gce_instance=attestation_pb2.GceInstancePolicy(project_id="trusted") + ) + conn = self._make_connection() + + with self.assertRaises(exceptions.AttestationHandshakeError): + conn.connect() + + self.assertFalse(conn.is_attested) + + def test_a_server_only_client_is_rejected_when_mutual_is_required(self): + conn = self._make_connection( + mutual_attestation=False, attestation_prover=None + ) + + with self.assertRaises(exceptions.AttestationHandshakeError): + conn.connect() + + def test_re_attesting_the_same_session_drops_the_connection(self): + conn = self._make_connection() + conn.connect() + conn.request("GET", "/infer") + self.assertEqual(conn.getresponse().status, 200) + + # A client that ignores the no-revalidation rule and re-attests anyway. + with self.assertRaises(exceptions.AttestationHandshakeError) as ctx: + conn._perform_mutual_attestation() + # _process_attestation_response re-wraps the status error, so the server's + # refusal shows up as the chained cause. + self.assertIn("403", str(ctx.exception.__cause__)) + + with self.subTest("SessionNoLongerServesApplicationTraffic"): + self.assertFalse(self._application_request_succeeds(conn)) + + def test_client_side_revalidation_is_refused_before_it_reaches_the_wire(self): + conn = self._make_connection() + conn.connect() + + with self.assertRaisesRegex( + exceptions.PromptEncryptionError, "cannot be revalidated" + ): + conn.revalidate_session() + + with self.subTest("SessionStillUsable"): + self.assertTrue(self._application_request_succeeds(conn)) + + def test_a_fresh_connection_re_attests_successfully(self): + """Re-attestation is done by reconnecting, not by revalidating.""" + first = self._make_connection() + first.connect() + first.close() + + second = self._make_connection() + second.connect() + + self.assertTrue(second.is_attested) + self.assertTrue(self._application_request_succeeds(second)) + + def _application_request_succeeds(self, conn) -> bool: + """Returns whether an application request is served over this session.""" + try: + conn.request("GET", "/infer") + return conn.getresponse().status == 200 + except ( + exceptions.PromptEncryptionError, + OSError, + socket.error, + ssl.SSLError, + http.client.HTTPException, + ): + return False + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/server/asgi_test.py b/tests/server/asgi_test.py index 775ec52..885f37e 100644 --- a/tests/server/asgi_test.py +++ b/tests/server/asgi_test.py @@ -22,6 +22,7 @@ from absl.testing import absltest from absl.testing import parameterized +from prompt_encryption_sdk import attestation as attestation_protocol from prompt_encryption_sdk.proto import attestation_pb2 from prompt_encryption_sdk.server import asgi from prompt_encryption_sdk.server import attestation @@ -178,6 +179,61 @@ async def receive(): asyncio.run(run()) + def test_replayed_attestation_drops_the_connection(self): + """A mutually attested session may not attest again; it is dropped.""" + + async def run(): + self.mw._attested_sockets.add(self.ssl_obj) + self.mock_attested_tls.attest_connection.side_effect = ( + attestation_protocol.AttestationReplayError("already attested") + ) + scope = { + "type": "http", + "path": "/_attest-connection", + "method": "POST", + "headers": [], + "extensions": {"tls_socket": self.ssl_obj}, + } + + async def receive(): + return {"type": "http.request", "body": b"{}", "more_body": False} + + await self.mw(scope, receive, self.send) + + (response_start_arg,) = self.send.call_args_list[0].args + with self.subTest("Forbidden"): + self.assertEqual( + response_start_arg["status"], http.HTTPStatus.FORBIDDEN + ) + with self.subTest("ConnectionClosed"): + self.assertIn( + (b"connection", b"close"), response_start_arg["headers"] + ) + with self.subTest("AuthorizationRevoked"): + self.assertNotIn(self.ssl_obj, self.mw._attested_sockets) + + # Application traffic on the dropped session is refused. + self.send.reset_mock() + await self.mw( + { + "type": "http", + "path": "/infer", + "method": "GET", + "headers": [], + "extensions": {"tls_socket": self.ssl_obj}, + }, + receive, + self.send, + ) + (unauthorized_start,) = self.send.call_args_list[0].args + with self.subTest("SubsequentTrafficUnauthorized"): + self.assertEqual( + unauthorized_start["status"], http.HTTPStatus.UNAUTHORIZED + ) + self.app.assert_not_called() + + asyncio.run(run()) + def test_handle_attestation_success_with_label(self): async def run(): scope = { @@ -361,7 +417,9 @@ def test_run_uvicorn_app_defaults( ) mock_token_manager_cls.return_value.__enter__.assert_called_once() mock_attested_tls_cls.assert_called_once_with( - mock_token_manager_cls.return_value + mock_token_manager_cls.return_value, + client_policy=None, + require_mutual_attestation=False, ) mock_uvicorn_run.assert_called_once() self.assertIsInstance( @@ -397,7 +455,11 @@ def test_run_uvicorn_app_with_managers( ) mock_token_manager.__enter__.assert_called_once() - mock_attested_tls_cls.assert_called_once_with(mock_token_manager) + mock_attested_tls_cls.assert_called_once_with( + mock_token_manager, + client_policy=None, + require_mutual_attestation=False, + ) mock_uvicorn_run.assert_called_once() self.assertIsInstance( mock_uvicorn_run.call_args[0][0], asgi.PromptEncryptionASGIMiddleware @@ -433,7 +495,9 @@ def test_run_uvicorn_app_with_log_config( ) mock_token_manager_cls.return_value.__enter__.assert_called_once() mock_attested_tls_cls.assert_called_once_with( - mock_token_manager_cls.return_value + mock_token_manager_cls.return_value, + client_policy=None, + require_mutual_attestation=False, ) mock_uvicorn_run.assert_called_once() self.assertIsInstance( diff --git a/tests/server/attestation_test.py b/tests/server/attestation_test.py index a1d44d3..296dfc8 100644 --- a/tests/server/attestation_test.py +++ b/tests/server/attestation_test.py @@ -18,6 +18,8 @@ from unittest import mock from absl.testing import absltest +from prompt_encryption_sdk import attestation as attestation_protocol +from prompt_encryption_sdk.client import exceptions from prompt_encryption_sdk.ekm import exporter from prompt_encryption_sdk.proto import attestation_pb2 from prompt_encryption_sdk.server import attestation @@ -182,5 +184,278 @@ def test_attest_connection_mixed_verifier_types_fails(self): attested_tls_instance.attest_connection(request, ssl_obj=mock_ssl_obj) +def _mutual_request(**overrides): + """Builds a well-formed mutual AttestConnectionRequest.""" + fields = dict( + required_verifier_type=[attestation_pb2.VerifierType.VERIFIER_TYPE_GCA], + protocol_version=( + attestation_protocol.MUTUAL_ATTESTATION_PROTOCOL_VERSION + ), + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + client_attestation=attestation_pb2.AttestationProof( + session_signature=b"client-sig" + ), + ) + fields.update(overrides) + return attestation_pb2.AttestConnectionRequest(**fields) + + +class MutualAttestedTlsTest(absltest.TestCase): + """Covers the single-round-trip mutual exchange.""" + + def setUp(self): + super().setUp() + self.ekm = b"session-ekm" + self.key_manager = mock.create_autospec(keys.KeyManager, instance=True) + self.token_manager = mock.create_autospec(token.TokenManager, instance=True) + self.token_manager.key_manager = self.key_manager + self.token_manager.get_identity_snapshot.return_value = ( + b"server-ecdsa", + b"server-mldsa", + b"server-token", + ) + self.key_manager.sign_payload.return_value = b"server-sig" + self.key_manager.sign_payload_mldsa.return_value = b"server-pqc-sig" + + self.ssl_obj = mock.MagicMock() + self.ssl_obj.export_keying_material.return_value = self.ekm + + self.client_validator = mock.MagicMock() + self.policy = attestation_pb2.AttestationPolicy() + + def _attested_tls(self, **overrides): + kwargs = dict( + client_policy=self.policy, + attestation_validator_cls=lambda policy: self.client_validator, + ) + kwargs.update(overrides) + return attestation.AttestedTLS(self.token_manager, **kwargs) + + def test_verifies_client_proof_then_returns_server_proof(self): + atls = self._attested_tls() + request = _mutual_request() + + response = atls.attest_connection(request, ssl_obj=self.ssl_obj) + + with self.subTest("ClientProofVerifiedAgainstSessionEkm"): + self.client_validator.validate_proof.assert_called_once_with( + request.client_attestation, + self.ekm, + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_CLIENT, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ) + with self.subTest("EkmExportedWithoutContext"): + # Both peers must derive the same value straight from the TLS session. + self.ssl_obj.export_keying_material.assert_called_once_with( + b"EXPORTER-Prompt-Encryption-SDK", 32, context=None + ) + with self.subTest("ServerProofBoundToTheSameEkmAsTheServerRole"): + self.assertEqual( + self.key_manager.sign_payload.call_args[0][0], + attestation_protocol.build_signature_payload( + tls_ekm=self.ekm, + attestation_token=b"server-token", + peer_role=attestation_pb2.ATTESTATION_PEER_ROLE_SERVER, + mode=attestation_pb2.ATTESTATION_MODE_MUTUAL, + ), + ) + with self.subTest("NegotiationEchoed"): + self.assertEqual(response.protocol_version, 1) + self.assertEqual(response.mode, attestation_pb2.ATTESTATION_MODE_MUTUAL) + self.assertTrue(response.mutual_attestation_complete) + self.assertEqual(response.session_signature, b"server-sig") + self.assertEqual(response.pqc_session_signature, b"server-pqc-sig") + + def test_client_proof_is_verified_before_the_server_proves_itself(self): + """A failed client check must not leak a fresh server proof.""" + self.client_validator.validate_proof.side_effect = ( + exceptions.PolicyViolationError("bad client") + ) + atls = self._attested_tls() + + with self.assertRaises(exceptions.PolicyViolationError): + atls.attest_connection(_mutual_request(), ssl_obj=self.ssl_obj) + + self.key_manager.sign_payload.assert_not_called() + self.key_manager.sign_payload_mldsa.assert_not_called() + + def test_rejected_client_can_retry_on_the_same_session(self): + """Only a *successful* exchange consumes the session.""" + self.client_validator.validate_proof.side_effect = [ + exceptions.PolicyViolationError("bad client"), + None, + ] + atls = self._attested_tls() + + with self.assertRaises(exceptions.PolicyViolationError): + atls.attest_connection(_mutual_request(), ssl_obj=self.ssl_obj) + response = atls.attest_connection(_mutual_request(), ssl_obj=self.ssl_obj) + + self.assertTrue(response.mutual_attestation_complete) + + def test_second_attestation_on_an_attested_session_is_refused(self): + """Re-attesting binds to the same EKM, so it proves nothing new.""" + atls = self._attested_tls() + atls.attest_connection(_mutual_request(), ssl_obj=self.ssl_obj) + + with self.assertRaises(attestation_protocol.AttestationReplayError): + atls.attest_connection(_mutual_request(), ssl_obj=self.ssl_obj) + + def test_server_only_request_on_an_attested_session_is_refused(self): + """Replay refusal is not a way to downgrade an attested session.""" + atls = self._attested_tls() + atls.attest_connection(_mutual_request(), ssl_obj=self.ssl_obj) + + with self.assertRaises(attestation_protocol.AttestationReplayError): + atls.attest_connection( + attestation_pb2.AttestConnectionRequest( + required_verifier_type=[ + attestation_pb2.VerifierType.VERIFIER_TYPE_GCA + ] + ), + ssl_obj=self.ssl_obj, + ) + + def test_a_different_session_is_unaffected_by_an_attested_one(self): + atls = self._attested_tls() + atls.attest_connection(_mutual_request(), ssl_obj=self.ssl_obj) + + other_ssl_obj = mock.MagicMock() + other_ssl_obj.export_keying_material.return_value = b"other-ekm" + response = atls.attest_connection( + _mutual_request(), ssl_obj=other_ssl_obj + ) + + self.assertTrue(response.mutual_attestation_complete) + + def test_missing_client_proof_is_rejected(self): + atls = self._attested_tls() + request = _mutual_request() + request.ClearField("client_attestation") + + with self.assertRaisesRegex( + ValueError, "Client attestation proof is missing." + ): + atls.attest_connection(request, ssl_obj=self.ssl_obj) + + def test_nonce_is_rejected_in_mutual_mode(self): + """A caller-chosen EKM context would let the peers derive different keys.""" + atls = self._attested_tls() + + with self.assertRaisesRegex(ValueError, "nonce must be empty"): + atls.attest_connection( + _mutual_request(nonce=b"n" * 32), ssl_obj=self.ssl_obj + ) + + def test_unsupported_protocol_version_is_rejected(self): + atls = self._attested_tls() + + with self.assertRaisesRegex( + ValueError, "Unsupported mutual attestation protocol version" + ): + atls.attest_connection( + _mutual_request(protocol_version=2), ssl_obj=self.ssl_obj + ) + + def test_mutual_request_to_a_server_without_a_client_policy_is_rejected(self): + atls = attestation.AttestedTLS(self.token_manager) + + with self.assertRaisesRegex( + ValueError, "not configured to verify client evidence" + ): + atls.attest_connection(_mutual_request(), ssl_obj=self.ssl_obj) + + def test_server_only_request_is_rejected_when_mutual_is_required(self): + atls = self._attested_tls(require_mutual_attestation=True) + + with self.assertRaisesRegex( + ValueError, "This server requires mutual attestation." + ): + atls.attest_connection( + attestation_pb2.AttestConnectionRequest( + required_verifier_type=[ + attestation_pb2.VerifierType.VERIFIER_TYPE_GCA + ] + ), + ssl_obj=self.ssl_obj, + ) + + def test_server_only_request_is_accepted_when_mutual_is_optional(self): + atls = self._attested_tls() + + response = atls.attest_connection( + attestation_pb2.AttestConnectionRequest( + required_verifier_type=[ + attestation_pb2.VerifierType.VERIFIER_TYPE_GCA + ], + nonce=b"n" * 32, + ), + ssl_obj=self.ssl_obj, + ) + + self.assertFalse(response.mutual_attestation_complete) + self.assertEqual(response.mode, attestation_pb2.ATTESTATION_MODE_SERVER_ONLY) + self.client_validator.validate_proof.assert_not_called() + + def test_server_only_request_may_not_carry_a_client_proof(self): + atls = self._attested_tls() + + with self.assertRaisesRegex( + ValueError, "client_attestation is valid only for mutual attestation." + ): + atls.attest_connection( + attestation_pb2.AttestConnectionRequest( + required_verifier_type=[ + attestation_pb2.VerifierType.VERIFIER_TYPE_GCA + ], + client_attestation=attestation_pb2.AttestationProof(), + ), + ssl_obj=self.ssl_obj, + ) + + def test_unknown_mode_is_rejected(self): + atls = self._attested_tls() + + with self.assertRaisesRegex(ValueError, "Unsupported attestation mode"): + atls.attest_connection( + _mutual_request(mode=99), ssl_obj=self.ssl_obj + ) + + def test_requiring_mutual_attestation_without_a_policy_is_rejected(self): + with self.assertRaisesRegex( + ValueError, "client_policy is required when mutual attestation" + ): + attestation.AttestedTLS( + self.token_manager, require_mutual_attestation=True + ) + + def test_client_validator_is_built_once_and_reused(self): + """OIDC discovery is expensive; it must not run per connection.""" + factory = mock.MagicMock(return_value=self.client_validator) + atls = self._attested_tls(attestation_validator_cls=factory) + + atls.attest_connection(_mutual_request(), ssl_obj=self.ssl_obj) + other_ssl_obj = mock.MagicMock() + other_ssl_obj.export_keying_material.return_value = b"other-ekm" + atls.attest_connection(_mutual_request(), ssl_obj=other_ssl_obj) + + factory.assert_called_once_with(self.policy) + + def test_client_validator_is_not_built_for_server_only_traffic(self): + factory = mock.MagicMock(return_value=self.client_validator) + atls = self._attested_tls(attestation_validator_cls=factory) + + atls.attest_connection( + attestation_pb2.AttestConnectionRequest( + required_verifier_type=[ + attestation_pb2.VerifierType.VERIFIER_TYPE_GCA + ] + ), + ssl_obj=self.ssl_obj, + ) + + factory.assert_not_called() + + if __name__ == "__main__": absltest.main() diff --git a/tests/server/wsgi_test.py b/tests/server/wsgi_test.py index b49720d..d5d4021 100644 --- a/tests/server/wsgi_test.py +++ b/tests/server/wsgi_test.py @@ -21,6 +21,7 @@ from absl.testing import absltest from absl.testing import parameterized +from prompt_encryption_sdk import attestation as attestation_protocol from prompt_encryption_sdk.proto import attestation_pb2 from prompt_encryption_sdk.server import attestation from prompt_encryption_sdk.server import keys @@ -130,6 +131,44 @@ def test_handle_attestation_success(self): self.mw._attested_sockets, ) + def test_replayed_attestation_drops_the_connection(self): + """A mutually attested session may not attest again; it is dropped.""" + mock_socket = mock.create_autospec(ssl.SSLSocket, instance=True) + self.mw._attested_sockets.add(mock_socket) + self.mock_attested_tls.attest_connection.side_effect = ( + attestation_protocol.AttestationReplayError("already attested") + ) + + response = self.client.post( + "/_attest-connection", + json={}, + environ_overrides={"prompt_encryption.socket": mock_socket}, + ) + + with self.subTest("Forbidden"): + self.assertEqual(response.status_code, http.HTTPStatus.FORBIDDEN) + self.assertIn(b"already attested", response.data) + with self.subTest("ConnectionClosed"): + self.assertEqual(response.headers.get("Connection"), "close") + with self.subTest("AuthorizationRevoked"): + self.assertNotIn(mock_socket, self.mw._attested_sockets) + + def test_application_traffic_is_refused_after_a_replay_drop(self): + mock_socket = mock.create_autospec(ssl.SSLSocket, instance=True) + self.mw._attested_sockets.add(mock_socket) + self.mock_attested_tls.attest_connection.side_effect = ( + attestation_protocol.AttestationReplayError("already attested") + ) + environ_overrides = {"prompt_encryption.socket": mock_socket} + self.client.post( + "/_attest-connection", json={}, environ_overrides=environ_overrides + ) + + response = self.client.get("/infer", environ_overrides=environ_overrides) + + self.assertEqual(response.status_code, http.HTTPStatus.UNAUTHORIZED) + self.app.assert_not_called() + @parameterized.named_parameters( dict( testcase_name="missing_socket", @@ -207,7 +246,11 @@ def test_run_gunicorn_app(self): ) with self.subTest("AttestedTLS"): - mock_attested_tls_cls.assert_called_once_with(mock_token_manager) + mock_attested_tls_cls.assert_called_once_with( + mock_token_manager, + client_policy=None, + require_mutual_attestation=False, + ) with self.subTest("StandaloneApplication"): mock_standalone_app_cls.assert_called_once_with(