Attestation-bound encrypted tensor transport for confidential ML inference over VSock and TCP.
Disclaimer: This project is under development. All source code and features are not production ready.
Try it in 60 seconds — no project setup needed:
git clone https://github.com/cyntrisec/confidential-ml-transport.git
cd confidential-ml-transport
cargo run --example echo_server --features mockThis runs an encrypted echo server and client over TCP with a full attested handshake (mock attestation). You should see:
echo server listening on 127.0.0.1:9876
accepted connection from 127.0.0.1:...
[client] connected and handshake complete
[client] sending: echo request #0
[server] handshake complete
[server] echoing: echo request #0
[client] received: echo request #0
...
done!
To use in your own project:
cargo add confidential-ml-transport --features mockuse confidential_ml_transport::{
MockProvider, MockVerifier, SecureChannel, SessionConfig,
};
use bytes::Bytes;
// Server: accept an attested encrypted connection
let (stream, _) = listener.accept().await?;
let server_provider = MockProvider::new();
let server_verifier = MockVerifier::new();
let mut server = SecureChannel::accept_with_attestation(
stream,
&server_provider,
&server_verifier,
SessionConfig::development(),
).await?;
// Client: connect with attestation verification
let stream = tokio::net::TcpStream::connect("127.0.0.1:9876").await?;
let client_provider = MockProvider::new();
let client_verifier = MockVerifier::new();
let mut client = SecureChannel::connect_with_attestation(
stream,
&client_provider,
&client_verifier,
SessionConfig::development(),
).await?;
client.send(Bytes::from("hello")).await?;Note:
MockProvider/MockVerifierperform no real attestation and are for testing and development only. For production, useNitroProvider/NitroVerifier(featurenitro),SevSnpProvider/SevSnpVerifier(featuresev-snp),TdxProvider/TdxVerifier(featuretdx), or implement theAttestationProvider/AttestationVerifiertraits for your TEE platform.
confidential-ml-transport is a Rust library that provides secure, binary-framed communication between TEE (Trusted Execution Environment) enclaves and clients. It combines a compact wire protocol with X25519+ChaCha20Poly1305 encryption and a 3-message attested handshake, designed for streaming tensor data in confidential AI inference pipelines.
Key properties:
- Binary framing — 13-byte fixed header, compact tensor sub-headers with 8-byte-aligned data, 32 MiB max payload
- Attestation-bound sessions — session keys are derived from attestation documents, binding the cryptographic channel to a verified TEE identity
- Full channel encryption — all post-handshake frames (data, tensor, heartbeat, shutdown, error) are encrypted and authenticated via AEAD
- Key material protection — symmetric keys zeroized on drop, contributory DH check, domain-separated session ID
- Pluggable transports — TCP and VSock backends via feature flags
- Pluggable attestation — trait-based attestation provider/verifier, with mock, Nitro, SEV-SNP, Azure SEV-SNP, and TDX implementations
- Monotonic sequence enforcement — replay protection on every decrypted message
- Hardened handshake — configurable timeout, mandatory public key binding, sequence validation, confirmation binds both keys
- Measurement verification — verify PCR/measurement registers against expected values during handshake
- Connection retry — exponential backoff with jitter for resilient connection establishment
- Transparent proxy — encrypt-on-the-wire proxy pair for wrapping existing TCP services without code changes
13-byte frame header (big-endian):
0 1 2 3
0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
| magic (0xCF4D) | version (4) | msg_type |
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
| flags | sequence |
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
| | payload_len |
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
| |
+-+-+-+-+-+-+-+-+
| Message Type | Value | Description |
|---|---|---|
| Hello | 0x01 |
Handshake messages |
| Data | 0x02 |
Application data |
| Error | 0x03 |
Error from peer |
| Heartbeat | 0x04 |
Keep-alive |
| Shutdown | 0x05 |
Graceful close |
| Tensor | 0x06 |
Tensor payload with sub-header |
Tensor frames include a sub-header with dtype, shape, name, and 8-byte-aligned raw data. Supported dtypes: F32, F64, F16, BF16, I32, I64, U8, U32.
A 3-message protocol establishes an encrypted session:
Initiator (client) Responder (server/enclave)
| |
|--- Hello { pubkey_c, nonce_c, att_doc_c } ->|
| |
|<-- Hello { pubkey_s, nonce_s, att_doc_s } --|
| |
|--- Hello { confirmation_hash } ----------->|
| |
[session established, encrypted data flows]
- Client sends its ephemeral X25519 public key, nonce, and attestation document binding that key
- Server verifies the client attestation, then responds with its public key, nonce, and attestation document binding its key
- Client verifies the server attestation, both sides derive session keys, and the client sends a confirmation hash proving key agreement
Session keys are derived via HKDF-SHA256 from the X25519 shared secret, salted with a transcript hash that binds both attestation hashes, both ephemeral public keys, both nonces, and the protocol version using explicit labeled framing.
Requires feature
mock:cargo add confidential-ml-transport --features mock
MockProvider/MockVerifierskip real attestation and are for testing only. For production, useNitroProvider/NitroVerifieror implement your ownAttestationProvider/AttestationVerifier.
use bytes::Bytes;
use confidential_ml_transport::{
MockProvider, MockVerifier, SecureChannel, SessionConfig,
session::channel::Message,
};
// Server side
let listener = tokio::net::TcpListener::bind("127.0.0.1:9876").await?;
let (stream, _) = listener.accept().await?;
let server_provider = MockProvider::new();
let server_verifier = MockVerifier::new();
let mut server = SecureChannel::accept_with_attestation(
stream,
&server_provider,
&server_verifier,
SessionConfig::development(),
).await?;
match server.recv().await? {
Message::Data(data) => server.send(data).await?,
_ => {}
}
// Client side
let stream = tokio::net::TcpStream::connect("127.0.0.1:9876").await?;
let client_provider = MockProvider::new();
let client_verifier = MockVerifier::new();
let mut client = SecureChannel::connect_with_attestation(
stream,
&client_provider,
&client_verifier,
SessionConfig::development(),
).await?;
client.send(Bytes::from("hello")).await?;
let response = client.recv().await?; // Message::Data("hello")
client.shutdown().await?;Use the builder pattern to customize session configuration:
use std::time::Duration;
use std::collections::BTreeMap;
use confidential_ml_transport::{
SessionConfig, RetryPolicy, ExpectedMeasurements,
};
let config = SessionConfig::builder()
.handshake_timeout(Duration::from_secs(10))
.retry_policy(RetryPolicy {
max_retries: 5,
initial_delay: Duration::from_millis(100),
max_delay: Duration::from_secs(10),
backoff_multiplier: 2.0,
})
.expected_measurements(ExpectedMeasurements::new({
let mut m = BTreeMap::new();
m.insert(0, vec![0xAA; 48]); // expected PCR0
m
}))
.build()?;Use connect_with_retry to automatically retry failed connections with exponential backoff:
use confidential_ml_transport::{MockProvider, MockVerifier, SecureChannel, SessionConfig, RetryPolicy};
use std::time::Duration;
let config = SessionConfig::builder()
.retry_policy(RetryPolicy::default()) // 3 retries, 1s initial, 2x backoff
.allow_empty_measurements() // mock/dev only
.build()?;
let provider = MockProvider::new();
let verifier = MockVerifier::new();
let mut channel = SecureChannel::connect_with_retry(
|| async { tokio::net::TcpStream::connect("enclave:5000").await },
&provider,
&verifier,
config,
).await?;Verify PCR/measurement registers against expected values during the handshake. If any measurement mismatches, the connection is rejected before any application data flows:
use std::collections::BTreeMap;
use confidential_ml_transport::{
ExpectedMeasurements, SecureChannel, SessionConfig,
};
let mut expected = BTreeMap::new();
expected.insert(0, pcr0_bytes.to_vec()); // PCR0: enclave image hash
expected.insert(1, pcr1_bytes.to_vec()); // PCR1: kernel hash
let config = SessionConfig::builder()
.expected_measurements(ExpectedMeasurements::new(expected))
.build()?;
// Handshake will fail if the enclave's measurements don't match
let mut channel = SecureChannel::connect_with_attestation(
stream,
&provider,
&verifier,
config,
).await?;For testing, use MockVerifierWithMeasurements to simulate an enclave returning specific measurement values:
use confidential_ml_transport::MockVerifierWithMeasurements;
let verifier = MockVerifierWithMeasurements::new(vec![
vec![0xAA; 48], // measurement[0]
vec![0xBB; 48], // measurement[1]
]);Wrap any existing TCP service with encryption, without modifying the service code. The server proxy runs inside the enclave and forwards decrypted traffic to a local backend; the client proxy runs on the host and accepts plaintext TCP connections.
use std::sync::Arc;
use confidential_ml_transport::proxy::server::{run_server_proxy, ServerProxyConfig};
use confidential_ml_transport::proxy::client::{run_client_proxy, ClientProxyConfig};
use confidential_ml_transport::{MockProvider, MockVerifier, SessionConfig};
// Inside the enclave: decrypt and forward to local inference server
let server_provider = Arc::new(MockProvider::new());
let server_verifier = Arc::new(MockVerifier::new());
let server_config = ServerProxyConfig {
listen_addr: "0.0.0.0:5000".parse()?,
backend_addr: "127.0.0.1:8080".parse()?, // local inference server
session_config: SessionConfig::development(),
max_connections: 256,
};
tokio::spawn(run_server_proxy(server_config, server_provider, server_verifier));
// On the host: accept plaintext, encrypt and forward to enclave
let client_provider = Arc::new(MockProvider::new());
let client_verifier = Arc::new(MockVerifier::new());
let client_config = ClientProxyConfig {
listen_addr: "127.0.0.1:9000".parse()?,
enclave_addr: "enclave:5000".parse()?,
session_config: SessionConfig::development(),
max_connections: 256,
};
tokio::spawn(run_client_proxy(client_config, client_provider, client_verifier));
// Now any TCP client connecting to localhost:9000 gets transparent encryptionuse confidential_ml_transport::{TensorRef, DType};
let activations: Vec<f32> = vec![0.0; 128 * 768];
// Safe conversion: collect f32s into a byte vec via to_ne_bytes
let data: Vec<u8> = activations.iter().flat_map(|f| f.to_ne_bytes()).collect();
let tensor = TensorRef {
name: "hidden_state",
dtype: DType::F32,
shape: &[128, 768],
data: &data,
};
channel.send_tensor(tensor).await?;Requires feature
tdx:cargo add confidential-ml-transport --features tdx
TdxProvideruses the Linux configfs-tsm ABI (/sys/kernel/config/tsm/report/) available on kernel 6.7+.TdxVerifierparses TDX v4/v5 quotes. Production verification must use DCAP collateral (require_collateral: true) so the quote is anchored to Intel PCS collateral instead of only checking internal quote consistency. Tested on GCPc3-standard-4confidential VMs with real TDX hardware.
use std::collections::BTreeMap;
use confidential_ml_transport::{
ExpectedMeasurements, SecureChannel, SessionConfig, TcbStatus, TdxCollateral,
TdxProvider, TdxVerifier, TdxVerifyPolicy,
};
// Server (inside TDX TD): generate real TDX attestation
let provider = TdxProvider::new()?;
let verifier = /* verifier for the peer side, if the peer also attests */ ;
let mut server = SecureChannel::accept_with_attestation(
stream,
&provider,
&verifier,
SessionConfig::development(), // use pinned measurements for production
).await?;
// Client: verify TDX with full DCAP collateral and measurement pinning
let collateral = TdxCollateral {
root_ca_der,
pck_chain_der,
crl_der,
qe_identity_json: Some(qe_identity_json),
tcb_info_json: Some(tcb_info_json),
tcb_signing_chain_der: Some(tcb_signing_chain_der),
};
let verifier = TdxVerifier::with_policy(TdxVerifyPolicy {
expected_mrtd: Some(expected_mrtd.to_vec()),
collateral: Some(collateral),
require_collateral: true,
accepted_tcb_statuses: vec![TcbStatus::UpToDate, TcbStatus::SWHardeningNeeded],
..Default::default()
});
let mut expected = BTreeMap::new();
expected.insert(0, expected_mrtd.to_vec()); // measurement[0] is MRTD
let config = SessionConfig::builder()
.expected_measurements(ExpectedMeasurements::new(expected))
.build()?;
let mut client = SecureChannel::connect_with_attestation(
stream,
&client_provider,
&verifier,
config,
).await?;TdxVerifier::new(...) is retained as a backward-compatible convenience constructor. It does not require DCAP collateral by default and should not be the sole production trust decision. TdxVerifier extracts measurements as: MRTD → measurements[0], RTMR0-3 → measurements[1..5].
Requires feature
sev-snp:cargo add confidential-ml-transport --features sev-snp
use std::collections::BTreeMap;
use confidential_ml_transport::{
ExpectedMeasurements, SecureChannel, SessionConfig,
SevSnpProvider, SevSnpVerifier,
};
// Server (inside SEV-SNP VM)
let provider = SevSnpProvider::new()?;
let verifier = /* verifier for the peer side, if the peer also attests */ ;
let mut server = SecureChannel::accept_with_attestation(
stream,
&provider,
&verifier,
SessionConfig::development(), // use pinned measurements for production
).await?;
// Client: verify SEV-SNP attestation and pin the launch measurement
let verifier = SevSnpVerifier::new(Some(expected_measurement.to_vec()));
let mut expected = BTreeMap::new();
expected.insert(0, expected_measurement.to_vec());
let config = SessionConfig::builder()
.expected_measurements(ExpectedMeasurements::new(expected))
.build()?;
let mut client = SecureChannel::connect_with_attestation(
stream,
&client_provider,
&verifier,
config,
).await?;For TEE platforms without a built-in implementation, implement the AttestationProvider and AttestationVerifier traits:
use async_trait::async_trait;
use confidential_ml_transport::{
AttestationProvider, AttestationVerifier,
attestation::types::{AttestationDocument, VerifiedAttestation},
error::AttestError,
};
struct MyProvider { /* platform-specific handle */ }
#[async_trait]
impl AttestationProvider for MyProvider {
async fn attest(
&self,
user_data: Option<&[u8]>,
nonce: Option<&[u8]>,
public_key: Option<&[u8]>,
) -> Result<AttestationDocument, AttestError> {
// Call platform API to generate attestation document
todo!()
}
}Built-in attestation implementations:
| Provider/Verifier | Feature | TEE Platform |
|---|---|---|
NitroProvider / NitroVerifier |
nitro |
AWS Nitro Enclaves |
SevSnpProvider / SevSnpVerifier |
sev-snp |
AMD SEV-SNP (Azure, GCP) |
AzureSevSnpProvider / AzureSevSnpVerifier |
azure-sev-snp |
Azure CVM HCL-backed SEV-SNP |
TdxProvider / TdxVerifier |
tdx |
Intel TDX (GCP, Azure) |
| Feature | Default | Description |
|---|---|---|
mock |
No | Mock attestation provider/verifier for testing |
tcp |
Yes | TCP transport helpers and transparent proxy |
vsock |
No | VSock transport via tokio-vsock |
nitro |
No | AWS Nitro Enclave attestation (NitroProvider/NitroVerifier) |
sev-snp |
No | AMD SEV-SNP attestation (SevSnpProvider/SevSnpVerifier) |
azure-sev-snp |
No | Azure CVM SEV-SNP attestation through the vTPM/HCL report path |
tdx |
No | Intel TDX attestation (TdxProvider/TdxVerifier) |
# Default (tcp only)
cargo build
# With mock attestation for testing
cargo build --features mock
# With VSock support
cargo build --features vsock
# With Nitro attestation (requires libssl-dev)
cargo build --features nitro
# With SEV-SNP attestation (requires libssl-dev)
cargo build --features sev-snp
# With Azure CVM SEV-SNP attestation (requires libssl-dev and tss2 headers)
cargo build --features azure-sev-snp
# With TDX attestation (requires libssl-dev)
cargo build --features tdx
# All features
cargo build --all-featuresTest counts depend on features and environment:
| Command | Tests | Notes |
|---|---|---|
cargo test --features "mock,tcp,tdx" |
~180 | Full suite without platform packages (109 lib + 70 integration + 1 doc-test). Proxy tests need socket permissions. |
cargo test --features "mock,tcp" |
~43 | Without TDX attestation tests |
cargo test --all-features |
~180+ | Requires system TSS2 headers for azure-sev-snp / tss-esapi |
Some integration tests (proxy, session) bind TCP ports and may fail in sandboxed
environments that restrict SO_REUSEADDR or ephemeral port binding. CI is the
authoritative source for full-suite pass/fail.
# Recommended: full suite without system dependencies
cargo test --features "mock,tcp,tdx"
# All tests (requires tss2-sys headers for azure-sev-snp / tss-esapi)
cargo test --all-features
# Property-based tests only
cargo test --test frame_roundtrip
# Retry and measurement tests
cargo test --test session_retry
# Proxy integration tests
cargo test --test proxy_integration
# SEV-SNP attestation tests
cargo test --features sev-snp
# TDX attestation tests
cargo test --features tdx
# Benchmarks
cargo bench --bench frame_codec
# Fuzz targets (requires nightly): frame_decode, tensor_decode, handshake_resp, aead_open, handshake_init
cargo +nightly fuzz run fuzz_frame_decode fuzz/seed_corpus/fuzz_frame_decode -- -max_total_time=60
cargo +nightly fuzz run fuzz_handshake_resp fuzz/seed_corpus/fuzz_handshake_resp -- -max_total_time=60
# Run the echo server example (requires mock feature)
cargo run --example echo_server --features mock| Example | Description | Command |
|---|---|---|
echo_server |
Encrypted echo over TCP with mock attestation | cargo run --example echo_server --features mock |
hostile-host-demo |
A/B proof that SecureChannel stops a hostile relay from reading tensors (threat model) | cargo run --release -p hostile-host-demo |
nitro-inference |
End-to-end confidential MiniLM-L6-v2 inference over SecureChannel | See README |
Benchmarked on 3 cloud platforms — including 2 real confidential VMs with hardware memory encryption. Transport crypto overhead is <0.3% of inference time on all platforms.
| Metric | AWS m6i.xlarge | Azure DC4ads_v5 | GCP c3-standard-4 |
|---|---|---|---|
| CPU | Intel Xeon 8375C | AMD EPYC 7763 (Milan) | Intel Sapphire Rapids |
| Security | Standard VM | SEV-SNP | TDX |
| Handshake p50 | 139 µs | 168 µs | 143 µs |
| 1536B RTT p50 | 33 µs | 124 µs | 108 µs |
| 4KB throughput | 751 MB/s | 618 MB/s | 604 MB/s |
| 384KB throughput | 825 MB/s | 725 MB/s | 918 MB/s |
| Overhead vs 100ms inference | <0.17% | <0.29% | <0.25% |
See benchmark_results/cross_provider_comparison.md for detailed analysis.
| Metric | Value | Environment |
|---|---|---|
| Frame codec roundtrip (4 KB) | 320 ns (12 GiB/s) | Local |
| ChaCha20-Poly1305 seal (4 KB) | 6.2 µs (630 MB/s) | Local |
# Run all benchmarks
bash scripts/bench_transport_performance.sh
# Quick smoke test
bash scripts/bench_transport_performance.sh --quick
# Individual: handshake, overhead, throughput, reconnect, frame_codec
cargo bench --bench reconnectMeasured on m5.xlarge (Intel Xeon 8175M, 2 vCPU enclave, 2 GiB) with real Nitro attestation in production mode (Flags: "NONE", no debug).
| Phase | p50 | p95 | n |
|---|---|---|---|
| Connect + handshake (NSM + COSE + X25519) | 5.699 ms | 5.822 ms | 50 |
| Encrypted transport RTT (64 B echo) | 0.263 ms | 0.286 ms | 200 |
| Inference RTT (MiniLM-L6-v2, 384-dim F32) | 98.332 ms | 99.102 ms | 50 |
Transport overhead is 0.27% of inference time — encryption is not the bottleneck. No measurable performance difference between debug and production enclave modes.
Baseline: bench-baseline-v0.1.1.
See benchmark_results/nitro_enclave/ for full data including raw measurements and debug vs production comparison.
| Component | Algorithm | Purpose |
|---|---|---|
| Key exchange | X25519 (Diffie-Hellman) | Ephemeral shared secret |
| Key derivation | HKDF-SHA256 | Derive send/recv keys from shared secret + transcript |
| Encryption | ChaCha20Poly1305 | Per-message AEAD with AAD = version || msg_type || flags || session_id || sequence |
| Transcript | SHA256 | Bind session to attestation + public keys + nonces |
| Replay protection | Monotonic u64 sequence | Reject any sequence <= last accepted |
The following security measures have been applied based on a comprehensive audit:
- Key zeroization —
SymmetricKeyusesZeroize + ZeroizeOnDropto clear key material from memory when no longer needed.SealingContextandOpeningContextimplementDropto zeroize session IDs and sequence counters. - Contributory key check —
was_contributory()rejects non-contributory DH results (small-subgroup or identity point attacks). - Domain-separated session ID — Session ID is derived from the transcript hash via HKDF with label
"cmt-session-id", preventing reuse as HKDF salt.
- All post-handshake frames encrypted — Heartbeat, shutdown, and error frames are encrypted via AEAD, preventing traffic analysis and injection of unauthenticated control messages. Unencrypted frames in an established session are rejected.
- Unified sequence counters — The AEAD sealer's internal sequence counter is used directly as the frame header sequence number, eliminating desynchronization between the wire format and cryptographic state.
- Bounded read buffer — The read buffer enforces a maximum size to prevent memory exhaustion from oversized or malicious frames.
- Handshake timeout — Configurable via
SessionConfig::handshake_timeout(default: 30 seconds). Prevents resource exhaustion from stalled or slow handshakes. - Mandatory public key binding — Each peer's attestation document must include a public key that matches its handshake key exchange. Missing public keys are rejected.
- Measurement verification — Optional
ExpectedMeasurementschecked during the handshake, before any application data flows. Mismatched PCR/measurement values abort the connection. - Confirmation hash binds both keys — The confirmation message includes both the send and receive keys, ensuring both parties derived identical key pairs.
- Handshake sequence validation — Frame sequence numbers are validated during the handshake (initiator hello=0, responder hello=0, confirmation=1).
- Sanitized error messages — Internal error details are logged via
tracingbut not exposed in protocol-level error messages.
- Exponential backoff with jitter —
RetryPolicyprovides configurable retry with exponential delay (default: 3 retries, 1s initial, 2x multiplier, 30s cap) and random jitter in [0.5x, 1.0x] to avoid thundering herd. - Transport factory pattern —
connect_with_retryaccepts a closure to create fresh transports per attempt, ensuring clean state on each retry.
- Structured attestation logging —
tracing::infoevents emitted after successful attestation verification (document hash, measurement count) and measurement verification (expected count). - Debug-level measurement dumps — Full hex-encoded measurement values logged at
tracing::debugfor forensic analysis. - PCR check logging — Nitro verifier logs PCR check results at debug level.
- Tensor dimension cap —
ndimsis capped at 32 in the decoder, preventing allocation amplification from maliciously crafted tensor headers. - Flags encapsulation — The
Flagsinner field ispub(crate), withfrom_raw()/raw()accessors for external use.
- Mutual attestation depends on real providers on both sides — The protocol carries and verifies attestation documents from both peers. Demo deployments may intentionally pair a real provider on one side with
MockVerifieron the other side; those demos authenticate only the side backed by a real verifier and must not be described as full production mutual identity. - No transport binding — The channel authenticates the data stream but does not bind to a specific transport address (IP, VSock CID). Perform a transport-level identity check separately if required.
- Proxy is TCP-only — The transparent proxy currently supports TCP backends. VSock proxy support can be added by implementing the same pattern with
tokio-vsocklisteners/streams.
MIT OR Apache-2.0