Skip to content

Conversation

@ArmanKolozyan
Copy link

This PR implements the server-side PQXDH (Post-Quantum Extended Diffie-Hellman) protocol in the TEE prover server, completing the quantum-resistant key exchange implementation started in #1315. The TEE now supports hybrid X25519 + ML-KEM-768 (Kyber) key exchange, protecting against quantum computer attacks while maintaining backward compatibility with legacy P-256 clients.

How It Works

Suite Negotiation Flow

  1. Client sends hello with supported_suites: ["Self-PQXDH-1", "legacy-p256"]
  2. TEE selects best supported suite (prefers PQXDH)
  3. TEE generates ephemeral keys and responds with:
    • Attestation JWT containing suite and public keys
    • X25519 public key (32 bytes)
    • Kyber ML-KEM-768 public key (1184 bytes)

PQXDH Handshake

  1. hello - TEE generates X25519 + Kyber keypairs, computes X25519 shared secret, stores pending state
  2. key_exchange - TEE decapsulates Kyber ciphertext, derives session key via HKDF, transitions to complete state
  3. submit_request - TEE uses session key to decrypt client payload

Key Derivation (Signal PQXDH Spec)

  • IKM = F || X25519_shared || Kyber_shared (F = 32 bytes of 0xFF)
  • Salt = 32 zero bytes
  • Info = "Self-PQXDH-1_X25519_SHA-256_ML-KEM-768"
  • Output = HKDF-SHA256(IKM, Salt, Info, 32 bytes)

Changes

New Files

src/lib.rs

  • Library entry point exposing internal modules for testing
  • Enables integration tests to access store, types, and utils

examples/pqxdh_test_server.rs

  • Standalone test server for local PQXDH testing
  • Runs with --features test_mode (no DB, no circuits, mock attestation)
  • Includes debug_get_session_key endpoint for cross-language test verification
  • Used by TypeScript client-side integration tests to spawn real server programmatically

tests/store_tests.rs

  • 8 comprehensive tests for LruStore key material management
  • Tests all KeyMaterial variants: LegacyP256, PqxdhPending, PqxdhComplete
  • Tests state transitions, duplicate UUID rejection, LRU eviction

Modified Files

Cargo.toml

  • Added PQXDH dependencies:
  • Added test_mode feature flag for local testing without TEE infrastructure

src/types.rs

  • Extended HelloResponse with:
    • selected_suite: String - negotiated cipher suite
    • x25519_pubkey: Option<Vec<u8>> - server's X25519 public key (PQXDH only)
    • kyber_pubkey: Option<Vec<u8>> - server's Kyber public key (PQXDH only)

src/store.rs

  • Added KeyMaterial enum for multi-state session management
  • Added methods:
    • get_key_material() - retrieves KeyMaterial for state inspection
    • update_key_material() - transitions PqxdhPending -> PqxdhComplete
    • Updated get_shared_secret() - handles all KeyMaterial variants

src/server.rs

  • Modified hello() RPC method:
    • Added supported_suites parameter
    • Implements suite negotiation (prefers Self-PQXDH-1, fallbacks to legacy-p256)
    • PQXDH branch: generates X25519 + Kyber keypairs, stores PqxdhPending state
    • Legacy branch: existing P-256 flow (unchanged)
    • Validates key sizes (X25519: 32 bytes, P-256: 33 bytes)
  • Added key_exchange() RPC method:
    • Validates Kyber ciphertext length (1088 bytes for ML-KEM-768)
    • Decapsulates ciphertext to recover Kyber shared secret
    • Derives session key using HKDF per Signal spec
    • Updates store: PqxdhPending -> PqxdhComplete
    • Returns "key_exchange_complete" confirmation

src/utils.rs

  • Added #[cfg(feature = "test_mode")] mock attestation function:
    • Returns JWT-like structure with base64url encoding
    • NOT cryptographically valid (for testing only)
    • Allows local development without AWS Nitro Enclaves

Testing

Unit Tests (Rust)

cd tee-prover-server
cargo test
  • 8/8 store tests passing - KeyMaterial state management

Integration Tests (Cross-Language)

cd common
yarn test pqxdh-cross-language
  • 5/5 integration tests passing
  • Spawns Rust server with test_mode feature
  • Tests full PQXDH handshake programmatically
  • Verifies session keys match between TypeScript client and Rust server
  • Tests suite negotiation, error handling, and validation

Copy link

@cursor cursor bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the final PR Bugbot will review for you during this billing cycle

Your free Bugbot reviews will reset on November 17

Details

Your team is on the Bugbot Free tier. On this plan, Bugbot will review limited PRs each billing cycle for each member of your team.

To receive Bugbot reviews on all of your PRs, visit the Cursor dashboard to activate Pro and start your 14-day free trial.

// decapsulating to get shared secret
let kyber_shared = decaps_key.decapsulate(&ct).map_err(|e| {
format!("Decapsulation failed: {:?}", e)
}).unwrap();
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bug: Server Crashes on Kyber Operation Failures

The key_exchange method in both src/server.rs and examples/pqxdh_test_server.rs uses map_err().unwrap() for Kyber cryptographic operations. This pattern causes the server to panic and crash when these operations fail due to invalid client input, instead of returning a proper error response. This creates a denial-of-service vulnerability.

Additional Locations (1)

Fix in Cursor Fix in Web

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.

1 participant