Skip to content

Add verified TEE inference client with OHTTP and RSA-PSS signature verification#297

Merged
adambalogh merged 3 commits into
mainfrom
claude/jolly-euler-jq01yf
Jun 12, 2026
Merged

Add verified TEE inference client with OHTTP and RSA-PSS signature verification#297
adambalogh merged 3 commits into
mainfrom
claude/jolly-euler-jq01yf

Conversation

@adambalogh

Copy link
Copy Markdown
Collaborator

Summary

This PR adds a complete client-side implementation for verified, private TEE (Trusted Execution Environment) inference through an untrusted relay. It enables end-to-end cryptographic verification of inference responses without trusting the relay, host, or OpenGradient infrastructure.

Key Changes

New Modules

  1. tee_verify.py — Cryptographic verification of TEE responses

    • RSA-PSS signature verification (salt length 32, SHA256) over keccak256(requestHash || outputHash || timestamp)
    • Request canonicalization matching the gateway's hashing (JSON with sort_keys=True)
    • Support for both text and tool-call responses
    • TeeProof dataclass capturing verified provenance (TEE ID, hashes, timestamp, signature)
    • Helper functions: pem_from_der(), tee_id_for_key(), canonical_request_bytes(), response_content_for_hash()
    • build_inner_request() converts OpenAI-style chat bodies to canonical form, stripping attachment bytes from the hash while preserving them in the wire request
  2. tee_ohttp.py — Client-side Oblivious HTTP (RFC 9458) encapsulation

    • HPKE encryption to TEE's X25519 public key (fixed ciphersuite: DHKEM(X25519), HKDF-SHA256, ChaCha20-Poly1305)
    • Single-shot and chunked-streaming response decryption
    • EncapsulatedRequest dataclass with wire bytes and response secrets
    • ChunkedResponseDecrypter for incremental stream decryption with final-frame validation
    • Byte-compatible with tee-gateway recipient implementation
  3. tee_ohttp_client.py — High-level relay client combining all three pieces

    • OhttpRelayClient class for sending verified, private chat completions
    • VerifiedChatResponse dataclass with body, content, proof, and optional stream frames
    • RelayError exception for relay/inner response errors
    • Support for both single-shot and streaming requests
    • Full stream buffering and verification before returning frames to caller
    • Pluggable authentication via auth_headers provider callback
  4. tee_registry.py (modified) — Enhanced with OHTTP config support

    • OhttpConfig dataclass for TEE's HPKE public key and algorithm identifiers
    • TEEEndpoint now carries ohttp_config and signing_public_key_der

Tests

  • tee_verify_test.py — Signature verification against independently-built gateway signatures

    • Tests valid signatures, tampered content/request rejection, wrong key rejection
    • Tool-call output hashing validation
    • Multimodal content canonicalization (attachment bytes stripped from hash)
  • tee_ohttp_test.py — Wire-compatibility round-trips with actual tee-gateway code

    • Single-shot and chunked response encryption/decryption
    • Truncation detection for streaming
  • tee_ohttp_client_test.py — End-to-end integration tests

    • Full client → (fake relay+gateway) → client path for single-shot and streaming
    • Real tee-gateway recipient crypto (when checkout available)

Public API Exports

Updated src/opengradient/client/__init__.py and src/opengradient/__init__.py to export:

  • TEERegistry, TEEEndpoint
  • OhttpRelayClient, VerifiedChatResponse, RelayError
  • TeeProof, VerificationError
  • build_inner_request(), verify_response()

Dependencies

Added to pyproject.toml:

  • pyhpke — HPKE sender context (matches gateway's recipient implementation)
  • cryptography — RSA-PSS signature verification and key serialization
  • `

https://claude.ai/code/session_01PdYbDC47zuBGiZex7ZHMSs

…erification)

Expose the client side of the OpenGradient verified-inference protocol so any
integrator (the chat-app relay, third-party tools, and the new local proxy) can
route an OpenAI-style request to a TEE through an untrusted relay and
cryptographically verify the response, sharing one non-drifting implementation.

- tee_ohttp.py: client-side Oblivious HTTP (RFC 9458) encapsulation + single-shot
  and chunked-streaming response decryption, wire-compatible with the tee-gateway
  recipient and the chat-app browser client.
- tee_verify.py: RSA-PSS response signature verification, request canonicalization
  (build_inner_request), and tee_id/PCR helpers, mirroring the gateway's signing.
- tee_ohttp_client.py: high-level OhttpRelayClient tying registry + OHTTP + verify
  together; verifies before returning (buffers streams) so no unverified token is
  surfaced. Caller supplies relay auth headers.
- tee_registry: surface on-chain pcr_hash on TEEEndpoint for reproducible-build
  PCR pinning.
- Tests round-trip the OHTTP crypto against the real tee-gateway recipient code
  and verify signatures against an independently-constructed gateway signature.

https://claude.ai/code/session_01PdYbDC47zuBGiZex7ZHMSs
@adambalogh adambalogh marked this pull request as ready for review June 12, 2026 15:55
@socket-security

socket-security Bot commented Jun 12, 2026

Copy link
Copy Markdown

This comment was marked as outdated.

claude added 2 commits June 12, 2026 16:07
Minor release adding the OHTTP client + signature-verification utilities
(OhttpRelayClient, verify_response, build_inner_request, TEEEndpoint.pcr_hash).
Lets downstream packages (e.g. opengradient-local) pin opengradient>=1.1.0.

https://claude.ai/code/session_01PdYbDC47zuBGiZex7ZHMSs
…am tool-calls, CI-runnable crypto tests

- tee_ohttp.encapsulate_request: thread key_id (and validate kem/kdf/aead) from
  the registry OhttpConfig so a TEE that rotated its key_id still decapsulates.
- tee_verify.build_inner_request: reject non-list 'tools' and non-dict messages
  with UnsupportedRequestError (was AttributeError / silent dict->keys coercion);
  preserve tool_choice on the wire while keeping it out of the signed hash.
- OhttpRelayClient: wrap chunked-decrypt ValueError as VerificationError; pass the
  config's key/alg ids; aggregate streamed delta.tool_calls so honest tool-call
  streams verify against the gateway's signed json.dumps(tool_calls).
- Tests: add a self-contained OHTTP recipient so encapsulation/decryption +
  end-to-end verification run in CI without a tee-gateway checkout (real gateway
  still cross-checked when present, via a sys.path-restoring fixture).

https://claude.ai/code/session_01PdYbDC47zuBGiZex7ZHMSs
@adambalogh adambalogh merged commit 9d62c4a into main Jun 12, 2026
9 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants