Merge remote-tracking branch 'origin/master'

Updated Reserved entry order. Made DataType ID changes easier in future
(this MAY NOT  happen again once in use).

.github/workflows/ci.yml

@@ -0,0 +1,65 @@
+name: CI
+
+on:
+  push:
+    branches: [master]
+  pull_request:
+
+env:
+  CARGO_TERM_COLOR: always
+  # The example workspace config points the type-map build script at this file.
+  MTP_TYPE_MAPS: example-type-maps.yaml
+
+jobs:
+  fmt:
+    name: rustfmt
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: dtolnay/rust-toolchain@stable
+        with:
+          components: rustfmt
+      - run: cargo fmt --all --check
+
+  clippy:
+    name: clippy
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: dtolnay/rust-toolchain@stable
+        with:
+          components: clippy
+      - uses: Swatinem/rust-cache@v2
+      - run: cargo clippy --workspace --exclude mtp-wasm --all-targets --all-features -- -D warnings
+
+  test:
+    name: test
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: dtolnay/rust-toolchain@stable
+      - uses: Swatinem/rust-cache@v2
+      - run: cargo test --workspace --exclude mtp-wasm --all-features
+
+  wasm:
+    name: wasm build
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: dtolnay/rust-toolchain@stable
+        with:
+          targets: wasm32-unknown-unknown
+      - uses: Swatinem/rust-cache@v2
+      # The WebTransport bindings in web-sys are still gated behind this cfg.
+      - run: cargo build -p mtp-wasm --target wasm32-unknown-unknown
+        env:
+          RUSTFLAGS: --cfg=web_sys_unstable_apis
+
+  deny:
+    name: cargo-deny
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: EmbarkStudios/cargo-deny-action@v2
+        with:
+          command: check

.gitignore

@@ -1,5 +1,18 @@
 **target/
-**Cargo.lock
 *.pem
 *.key
 example-usage/dev-cert/
+
+# The root Cargo.lock and the example-usage workspace lockfile are committed:
+# the workspace ships binaries (example-usage/*) and a committed lockfile keeps
+# their builds reproducible. Workspace members do not own a lockfile (the root
+# one governs them); ignore any stray per-member lockfiles so they are not
+# committed by accident.
+/client/Cargo.lock
+/codec/Cargo.lock
+/common/Cargo.lock
+/crypto/Cargo.lock
+/host/Cargo.lock
+/transport/Cargo.lock
+/type-map/Cargo.lock
+/wasm/Cargo.lock

CONNECTOR.md

@@ -68,14 +68,24 @@ The host's `accept()` method:
 6. Returns `None` if the version is unsupported
 7. Returns an `MTPConnection` with the negotiated version otherwise
 
-### Login/Register Handshake (crypto feature)
+### Login/Register Handshake
 
-When `require_authentication` is set, the host sends a **greeting** first (host ID, public keys, nonce). The client then responds with either:
+When `require_authentication` is set, the parties run a mutually-authenticated
+**challenge-response**. The client speaks first with an *unsigned* hello:
 
-- **Login** (`CommunicationType::Identification`, ID 15): client ID, nonce, signature
-- **Register** (`CommunicationType::Register`, ID 17): public keys, nonce, signature
+- **Login** (`CommunicationType::Identification`, ID 15): version, client ID
+- **Register** (`CommunicationType::Register`, ID 17): version, public keys
 
-The host verifies the client's signature, sends a signed response, and the client verifies the host's signature.
+The host then issues a fresh random `server_challenge` in a signed `Challenge`
+(`CommunicationType::Challenge`, ID 21, carrying `ServerNonce`). The client signs
+that challenge, binding its id (login) or public keys (register), and returns a
+`ChallengeResponse` (ID 22). The host verifies the proof against the challenge it
+issued and sends a signed final response, which the client verifies.
+
+Because the client's proof covers the host-issued `server_challenge` (a one-time
+value held only on the accepting task's stack), a captured proof cannot be
+replayed on another connection. All signed payloads are domain-separated; see
+`mtp::crypto::auth`.
 
 ---
 
@@ -116,8 +126,8 @@ Client (v2.0)         Host (v0.0, v1.0, v2.0)
     |  CommValue{ Ident. }   |
     |    Version -> "2.0"    |
     |    Id -> 8765          |
-    |    Nonce -> ...        |
-    |    Signature -> ...    |
+    |  (unsigned hello; auth |
+    |   challenge follows)   |
     |----------------------->|
     |                        |  registry.negotiate(&[Version(2,0)])
     |                        |  -> Some(Version(2,0))

Cargo.toml

@@ -9,6 +9,24 @@ members = [
     "client",
     "wasm",
 ]
+# `wasm` is a wasm32-only crate: it relies on web-sys unstable APIs
+# (`--cfg=web_sys_unstable_apis`, set in wasm/.cargo/config.toml) and the
+# wasm32 target. Cargo only reads .cargo/config.toml from the invocation
+# directory and its ancestors, so building it for the host target from the
+# workspace root fails. Exclude it from the default set so a bare
+# `cargo build`/`test`/`clippy` at the root matches CI, which always uses
+# `--exclude mtp-wasm`. Build it explicitly with:
+#   cargo build -p mtp-wasm --target wasm32-unknown-unknown
+default-members = [
+    ".",
+    "common",
+    "crypto",
+    "type-map",
+    "codec",
+    "transport",
+    "host",
+    "client",
+]
 resolver = "3"
 
 # =============================================================================
@@ -61,3 +79,8 @@ host = ["dep:mtp-host", "mtp-codec/registry", "mtp-transport/host"]
 
 # MTP client - outgoing QUIC connections to a host.
 client = ["dep:mtp-client"]
+
+[dev-dependencies]
+tokio = { version = "1", features = ["full"] }
+rcgen = "0.14"
+rand = "0.8"

NATIVE-CLIENT.md

@@ -109,15 +109,22 @@ let config = ClientConfig {
 let conn = MTPClient::auth_connect(config, &keys, &host_pk).await?;
 ```
 
-Protocol:
-1. Client generates a random nonce
-2. Builds a signature payload: `version || client_id || client_nonce`
-3. Signs with Ed25519 (and optionally ML-DSA-65)
-4. Sends `Identification` frame containing version, client ID, nonce, signature(s)
-5. Host responds with `IdentificationResponse` containing echoed nonce, host
-   nonce, and host signature
+Protocol (challenge-response, the host issues the freshness):
+1. Client sends an unsigned `Identification` hello (version, client ID)
+2. Host replies with a `Challenge` carrying a fresh random `server_challenge`
+   and the host's signature over it; the client verifies that signature
+3. Client generates a random `client_nonce` and signs
+   `version || client_id || server_challenge || client_nonce` with Ed25519
+   (and optionally ML-DSA-65)
+4. Client sends a `ChallengeResponse` frame (nonce + signature(s))
+5. Host verifies the proof against `server_challenge` and responds with
+   `IdentificationResponse` (echoed nonce + host signature)
 6. Client verifies the host signature and nonce echo
 
+Because the client's signature covers the host-issued `server_challenge`, a
+captured proof cannot be replayed on another connection (each connection gets a
+different challenge).
+
 ### Registration
 
 ```rust
@@ -134,13 +141,16 @@ let id = conn.client_id;
 let keyring_bytes = keyring.to_bytes();
 ```
 
-Protocol:
-1. Client generates a random nonce
-2. Builds a signature payload: `version || client_nonce || public_key_bytes`
-3. Signs with Ed25519 (and optionally ML-DSA-65)
-4. Sends `Register` frame containing version, nonce, public key bundle, signature(s)
-5. Host assigns a new client ID, responds with `RegisterResponse` containing
-   the ID, echoed nonce, host nonce, and host signature
+Protocol (challenge-response):
+1. Client sends an unsigned `Register` hello (version, public key bundle)
+2. Host replies with a `Challenge` carrying a fresh random `server_challenge`
+   (signed by the host); the client verifies that signature
+3. Client generates a random `client_nonce` and signs
+   `version || server_challenge || client_nonce || public_key_bytes` with
+   Ed25519 (and optionally ML-DSA-65)
+4. Client sends a `ChallengeResponse` frame (nonce + signature(s))
+5. Host verifies the proof against `server_challenge`, assigns a new client ID,
+   and responds with `RegisterResponse` (the ID, echoed nonce, host signature)
 6. Client verifies the host signature and nonce echo
 
 ## Key Material
@@ -151,12 +161,12 @@ A `Keyring` bundles all secret and public key material for one identity:
 
 ```rust
 pub struct Keyring {
-    pub kem_secret_key: KemPrivateKey,
     pub kem_public_key: KemPublicKey,
-    pub sig_cl_secret_key: SignaturePrivateKey,   // Ed25519
-    pub sig_cl_public_key: SignaturePublicKey,    // Ed25519
-    pub sig_pq_secret_key: SignaturePqPrivateKey, // ML-DSA-65
+    pub kem_secret_key: KemPrivateKey,
     pub sig_pq_public_key: SignaturePqPublicKey,  // ML-DSA-65
+    pub sig_pq_secret_key: SignaturePqPrivateKey,
+    pub sig_cl_public_key: SignaturePublicKey,    // Ed25519
+    pub sig_cl_secret_key: SignaturePrivateKey,
 }
 ```
 
@@ -238,19 +248,24 @@ force-closes the QUIC connection if the peer has not already done so.
 ## Crypto Containers
 
 With the `crypto` feature, `DataValue` supports encrypted, signed, and
-signed+encrypted containers:
+signed+encrypted containers. Encryption uses ML-KEM to encapsulate to a
+recipient's KEM public key (from their `PublicKeyBundle`); only the holder of
+the matching `Keyring` can decrypt. Signing uses the sender's Ed25519 key.
 
 ```rust
-use mtp::crypto::{ChaCha20Poly1305, Ed25519Signer, SigAlgorithm};
+use mtp::crypto::{EncryptionType, Ed25519Signer, SigAlgorithm};
 
-let cipher = ChaCha20Poly1305::new(derive_encryption_key(...));
+let enc_type = EncryptionType::MlKemChaCha20Poly1305;
 let signer = Ed25519Signer::new(&keyring.sig_cl_secret_key)?;
 
+// `recipient` is the PublicKeyBundle of whoever should be able to decrypt
+// (e.g. the host's bundle, obtained out of band).
+
 // Encrypted container
 let mut enc = DataValue::Container(vec![
     (DataTypeId(1), DataValue::Str("secret".into())),
 ]);
-enc.encrypt_container(&cipher, b"aad");
+enc.encrypt_container(enc_type, &recipient, b"aad");
 
 // Signed container
 let mut sig = DataValue::Container(vec![
@@ -262,11 +277,18 @@ sig.sign_container(SigAlgorithm::ED25519, &signer);
 let mut sec = DataValue::Container(vec![
     (DataTypeId(1), DataValue::Str("both".into())),
 ]);
-sec.sign_and_encrypt_container(SigAlgorithm::ED25519, &signer, &cipher, b"aad");
+sec.sign_and_encrypt_container(SigAlgorithm::ED25519, &signer, enc_type, &recipient, b"aad");
 ```
 
-On the receiving side, use the corresponding `decrypt_into_container`,
-`verify_into_container`, or `decrypt_signed_encrypted_container` methods.
+On the receiving side, the recipient decrypts with its own `Keyring` (each blob
+is self-describing: its leading byte selects the algorithm and the matching KEM
+key from the keyring):
+
+```rust
+enc.decrypt_into_container(&keyring, b"aad");                // -> Container
+sig.verify_into_container(&verifier);                        // verifier: impl SignatureScheme
+sec.decrypt_signed_encrypted_container(&keyring, b"aad");    // -> SignedContainer, then verify_into_container
+```
 
 ## Policy Configuration
 
@@ -321,5 +343,3 @@ version and expects the host to negotiate a compatible version.
 | `AuthenticationFailed`  | Nonce mismatch or invalid host signature   |
 | `ConnectionError`       | QUIC connection failure                    |
 | `UseAfterClosed`        | Attempted send/receive after close         |
-
-

NATIVE-HOST.md

@@ -129,12 +129,19 @@ let negotiated = registry.negotiate(&[Version(1, 0), Version(2, 0)]);
 // -> Some(Version(2, 0)) if both versions are registered
 ```
 
-## Authentication Flow (crypto feature)
+## Authentication Flow
 
-When `require_authentication` is `true`, `accept()` runs an authenticated
-handshake before returning the connection. The flow is:
+When `require_authentication` is `true`, `accept()` runs a mutually-authenticated
+**challenge-response** handshake before returning the connection. The host issues
+a fresh, random `server_challenge` that the client must sign, which is what makes
+the client's proof unreplayable: a captured proof is bound to a one-time challenge
+the host generates per connection and will never reissue. The challenge lives only
+on the accepting task's stack; there is no replay database or shared state.
 
-### Login (existing client)
+All signed payloads begin with a one-byte domain-separation tag (see
+`mtp::crypto::auth`) so a signature for one step can never be reused as another.
+
+### Login
 
 ```
 Client                                    Host
@@ -142,26 +149,35 @@ Client                                    Host
   |  QUIC connect                           |
   |---------------------------------------->|
   |                                         |
-  |  Identification {                       |
-  |    Version, Id, ClientNonce,            |
-  |    Signature, [PqSignature]             |
+  |  Identification { Version, Id }         |  (unsigned hello)
+  |---------------------------------------->|
+  |                                         |  lookup get_existing_user(id)
+  |                                         |  generate random server_challenge
+  |  Challenge {                            |
+  |    ServerNonce(server_challenge),       |
+  |    Signature, [PqSignature]             |  host signs the challenge
+  |  }                                      |
+  |<----------------------------------------|
+  |  ChallengeResponse {                    |
+  |    ClientNonce, Signature, [PqSignature]|  client signs the challenge
   |  }                                      |
   |---------------------------------------->|
-  |                                         |  lookup get_existing_user(client_id)
-  |                                         |  verify Ed25519 (and optional ML-DSA) sig
+  |                                         |  verify proof over server_challenge
   |  IdentificationResponse {              |
-  |    Connected=true, ClientNonce(echoed), |
-  |    Id, Timestamp(new_nonce),            |
+  |    Connected=true, Id,                  |
+  |    ClientNonce(echoed),                 |
   |    Signature, [PqSignature]             |
   |  }                                      |
   |<----------------------------------------|
 ```
 
-The client signature payload is: `version_string || client_id (8 bytes, big-endian) || client_nonce (16 bytes, big-endian)`
+Payloads (`||` is concatenation, integers big-endian; `DS_*` are domain tags):
 
-The host signs: `0x01 || assigned_id (8 bytes, big-endian) || client_nonce (16 bytes) || host_new_nonce (16 bytes)`
+- Host challenge: `DS_CHALLENGE || id (8) || server_challenge (16)`
+- Client proof:   `DS_LOGIN_PROOF || version_string || id (8) || server_challenge (16) || client_nonce (16)`
+- Host final:     `DS_HOST_FINAL || assigned_id (8) || client_nonce (16) || server_challenge (16)`
 
-### Register (new client)
+### Register
 
 ```
 Client                                    Host
@@ -170,23 +186,32 @@ Client                                    Host
   |---------------------------------------->|
   |                                         |
   |  Register {                             |
-  |    Version, ClientNonce,                |
-  |    PublicKeys (serialized PublicKeyBundle),
-  |    Signature, [PqSignature]             |
+  |    Version,                             |  (unsigned hello)
+  |    PublicKeys (serialized PublicKeyBundle)
   |  }                                      |
   |---------------------------------------->|
-  |                                         |  extract PublicKeyBundle from frame
-  |                                         |  verify Ed25519 (and optional ML-DSA) sig
+  |                                         |  generate random server_challenge
+  |  Challenge {                            |
+  |    ServerNonce(server_challenge),       |
+  |    Signature, [PqSignature]             |  (challenge binds id = 0)
+  |  }                                      |
+  |<----------------------------------------|
+  |  ChallengeResponse {                    |
+  |    ClientNonce, Signature, [PqSignature]|
+  |  }                                      |
+  |---------------------------------------->|
+  |                                         |  verify proof over server_challenge
   |                                         |  call complete_register(bundle) -> new_id
   |  RegisterResponse {                    |
-  |    Connected=true, ClientNonce(echoed), |
-  |    Id, Timestamp(new_nonce),            |
+  |    Connected=true, Id(new_id),          |
+  |    ClientNonce(echoed),                 |
   |    Signature, [PqSignature]             |
   |  }                                      |
   |<----------------------------------------|
 ```
 
-The client signature payload is: `version_string || client_nonce (16 bytes) || public_key_bytes`
+The register client proof is:
+`DS_REGISTER_PROOF || version_string || server_challenge (16) || client_nonce (16) || public_key_bytes`
 
 After a successful handshake, `accept()` returns an `MTPConnection` with
 `auth_state = Authenticated`, `client_id` set, and `client_public_key`

README.md

@@ -2,7 +2,7 @@
 
 MTP is a modular transport protocol built on QUIC. It provides version-negotiable type maps, a binary codec, cryptographic primitives (classical and post-quantum), and host/client connection management with mutual authentication.
 
-There are Area specific docs when working with seperate concerns for the [Native-Client](./NATIVE-CLIENT.md), [WASM-Client](./WASM-CLIENT.md) & [Host](./NATIVE-HOST.md)
+See the area-specific docs for [Native Client](./NATIVE-CLIENT.md), [WASM Client](./WASM-CLIENT.md), and [Host](./NATIVE-HOST.md)
 
 ## Getting Started
 

WASM-CLIENT.md

@@ -37,7 +37,7 @@ if (!WasmClient.is_supported()) {
 const config = new ConnectionConfig("https://host.example.com:4433");
 config.client_id = 12345n;                    // optional, for re-authentication
 config.server_certificate_hashes = [           // optional, for certificate pinning
-  "sha256:abc123...",
+  "sha-256:abc123...",
 ];
 ```
 
@@ -55,11 +55,12 @@ providing its hash:
 
 ```typescript
 config.server_certificate_hashes = [
-  "sha256:abcd1234...",  // hex-encoded hash value
+  "sha-256:abcd1234...",  // hex-encoded hash value
 ];
 ```
 
-The hash format is `"<algorithm>:<hex-encoded-hash>"`. When hashes are
+The hash format is `"<algorithm>:<hex-encoded-hash>"`, where the only algorithm
+the browser's WebTransport API currently accepts is `sha-256`. When hashes are
 provided, the browser **only** trusts certificates matching one of the given
 hashes and ignores its root store for this connection.
 
@@ -96,7 +97,7 @@ await client.connect(config);
 
 Sends an `Identification` frame with the protocol version and client ID. The host may accept or reject. No cryptographic handshake occurs.
 
-### Authenticated Login (existing client ID)
+### Authenticated Login
 
 ```typescript
 const confirmedId = await client.auth_connect(
@@ -107,10 +108,14 @@ const confirmedId = await client.auth_connect(
 );
 ```
 
-Exchange: client sends a signed `Identification` frame, the host verifies it and
-responds with a signed `IdentificationResponse`. Returns the confirmed client ID.
+Exchange (challenge-response): the client sends an unsigned `Identification`
+hello, the host replies with a signed `Challenge` carrying a fresh
+`server_challenge`, the client signs that challenge in a `ChallengeResponse`, and
+the host verifies it and replies with a signed `IdentificationResponse`. Signing
+over the host-issued challenge is what prevents a captured proof from being
+replayed on another connection. Returns the confirmed client ID.
 
-### Registration (new client)
+### Registration
 
 ```typescript
 const newId = await client.auth_register(
@@ -120,8 +125,10 @@ const newId = await client.auth_register(
 );
 ```
 
-Exchange: client sends a signed `Register` frame with public keys, the host
-assigns a new ID and responds with a signed `RegisterResponse`. Returns the
+Exchange (challenge-response): the client sends an unsigned `Register` hello with
+its public keys, the host replies with a signed `Challenge`, the client signs it
+(binding the public-key bundle) in a `ChallengeResponse`, and the host verifies
+it, assigns a new ID, and responds with a signed `RegisterResponse`. Returns the
 newly assigned client ID.
 
 ## Sending and Receiving Messages
@@ -172,13 +179,16 @@ binary payload. Useful for health checks and simple messaging.
 function build_demo_message(
   clientId: bigint,
   keyringBytes: Uint8Array,
+  hostBundleBytes: Uint8Array,
 ): Uint8Array;
 ```
 
 Constructs a `Ping` frame that demonstrates encrypted, signed, and
-signed+encrypted containers using a deterministic demo key. The paired host
-handler can decrypt and verify these containers if it knows the same shared
-secret.
+signed+encrypted containers. The containers are ML-KEM-encrypted to the host's
+`PublicKeyBundle` (`hostBundleBytes`, the same bytes passed to `auth_connect` /
+`auth_register`), so the host decrypts them with its own keyring; signatures use
+the client keyring's Ed25519 key. The client keyring only needs its Ed25519
+signing key for this demo.
 
 ### `parse_auth_response`
 

client/src/lib.rs

@@ -1,6 +1,4 @@
-#[cfg(feature = "crypto")]
-use mtp_codec::DataType;
-use mtp_codec::{CommunicationValue, DataTypeId, DataValue, PROTOCOL_VERSION, Version};
+use mtp_codec::{CommunicationValue, DataType, DataValue, PROTOCOL_VERSION, Version};
 use mtp_common::CommunicationError;
 use mtp_transport::{Policy, Receiver, Sender};
 
@@ -24,7 +22,7 @@ pub struct ClientConfig {
     pub client_id: u64,
 }
 
-// Established MTP connection with a single negotiated version.
+/* Established MTP connection with a single negotiated version. */
 pub struct MTPConnection {
     pub version: Version,
     pub sender: Sender,
@@ -59,9 +57,9 @@ impl MTPClient {
         // Build the initial identification message with the protocol version.
         let version_str = format!("{}", PROTOCOL_VERSION);
         let ident = CommunicationValue::new(mtp_codec::CommunicationType::Identification)
-            .add_data(DataTypeId(3), DataValue::Str(version_str))
-            .add_data(
-                DataTypeId(6),
+            .add_typed_default(DataType::Version, DataValue::Str(version_str))
+            .add_typed_default(
+                DataType::Id,
                 DataValue::UnsignedNumber(config.client_id.into()),
             );
 
@@ -80,91 +78,66 @@ impl MTPClient {
 }
 
 /* ===== Authentication ===== */
+
+/*
+ * Verify the host's signature over the challenge it issued (step 2).
+ *
+ * `id` is the client id for a login, or `0` for a registration (the host binds
+ * `0` since no id has been assigned yet). The Ed25519 signature is mandatory;
+ * the ML-DSA signature is checked only when the host included one.
+ */
 #[cfg(feature = "crypto")]
-impl MTPClient {
-    pub async fn auth_connect(
-        config: ClientConfig,
-        keys: &mtp_crypto::Keyring,
-        host_public_key_bundle: &mtp_crypto::PublicKeyBundle,
-    ) -> Result<MTPConnection, CommunicationError> {
-        use mtp_crypto::{
-            Ed25519Signer, MlDsaSigner, SignatureScheme, verify_ed25519, verify_ml_dsa,
-        };
+fn verify_host_challenge(
+    challenge: &CommunicationValue,
+    tm: &mtp_codec::TypeMap,
+    host_pk: &mtp_crypto::PublicKeyBundle,
+    id: u64,
+    server_challenge: u128,
+) -> Result<(), CommunicationError> {
+    use mtp_crypto::{auth, verify_ed25519, verify_ml_dsa};
 
-        let (sender, receiver) =
-            mtp_transport::connect(&config.url, config.server_cert, Policy::default()).await?;
-
-        // 1. Build and send Identification message immediately (no greeting)
-        let client_nonce: u128 = rand::random();
-        let version_str = format!("{}", PROTOCOL_VERSION);
-
-        let mut sig_payload = Vec::new();
-        sig_payload.extend_from_slice(version_str.as_bytes());
-        sig_payload.extend_from_slice(&config.client_id.to_be_bytes());
-        sig_payload.extend_from_slice(&client_nonce.to_be_bytes());
-
-        let signer = Ed25519Signer::new(&keys.sig_cl_secret_key)
-            .map_err(|e| CommunicationError::Other(e.to_string()))?;
-
-        /* ===== Signature ===== */
-        let signature = signer
-            .sign(&sig_payload)
-            .map_err(|e| CommunicationError::Other(e.to_string()))?;
-
-        let mut ident = CommunicationValue::new(mtp_codec::CommunicationType::Identification)
-            .add_typed_default(DataType::Version, DataValue::Str(version_str))
-            .add_typed_default(
-                DataType::Id,
-                DataValue::UnsignedNumber(config.client_id as u128),
-            )
-            .add_typed_default(
-                DataType::ClientNonce,
-                DataValue::UnsignedNumber(client_nonce),
-            )
-            .add_typed_default(DataType::Signature, DataValue::Bytes(signature));
-
-        if !keys.sig_pq_secret_key.as_bytes().is_empty() {
-            let pq_signer = MlDsaSigner::new(&keys.sig_pq_secret_key, &keys.sig_pq_public_key)
-                .map_err(|e| CommunicationError::Other(e.to_string()))?;
-            let pq_signature = pq_signer
-                .sign(&sig_payload)
-                .map_err(|e| CommunicationError::Other(e.to_string()))?;
-            ident = ident.add_typed_default(DataType::PqSignature, DataValue::Bytes(pq_signature));
-        }
-        /* ===== End Signature ===== */
-
-        sender.send(&ident).await?;
-
-        // 2. Receive host response (single message)
-        let response = receiver.receive().await?;
-
-        let expected_type = mtp_codec::CommunicationTypeId(16); // IdentificationResponse
-        if response.get_type() != expected_type {
-            return Err(unexpected_response_type_error(
-                "auth_connect",
-                expected_type,
-                &response,
-            ));
-        }
-
-        let connected = response.get_data(DataTypeId(11));
-        match connected {
-            DataValue::BoolTrue => {}
-            DataValue::BoolFalse => {
-                return Err(CommunicationError::AuthenticationFailed(
-                    "Server rejected authentication".into(),
-                ));
-            }
+    let sig = match challenge.get_data(DataType::Signature.to_id(tm)) {
+        DataValue::Bytes(b) => b.clone(),
         _ => {
             return Err(CommunicationError::AuthenticationFailed(
-                    "Invalid response".into(),
+                "Missing host challenge signature".into(),
             ));
         }
+    };
+    let pq_sig = match challenge.get_data(DataType::PqSignature.to_id(tm)) {
+        DataValue::Bytes(b) => b.clone(),
+        _ => vec![],
+    };
+
+    let payload = auth::challenge_payload(id, server_challenge);
+    verify_ed25519(&host_pk.sig_cl_public_key, &payload, &sig).map_err(|_| {
+        CommunicationError::AuthenticationFailed("Host challenge signature invalid".into())
+    })?;
+    if !pq_sig.is_empty() && verify_ml_dsa(&host_pk.sig_pq_public_key, &payload, &pq_sig).is_err() {
+        return Err(CommunicationError::AuthenticationFailed(
+            "Host challenge PQ signature invalid".into(),
+        ));
+    }
+    Ok(())
 }
 
-        let echo_nonce = response.get_data(DataTypeId(7));
-        match echo_nonce {
-            DataValue::UnsignedNumber(n) if *n == client_nonce as u128 => {}
+/*
+ * Verify the host's final confirmation (step 4): the echoed `client_nonce` and
+ * the host signature over the handshake transcript.
+ */
+#[cfg(feature = "crypto")]
+fn verify_host_final(
+    response: &CommunicationValue,
+    tm: &mtp_codec::TypeMap,
+    host_pk: &mtp_crypto::PublicKeyBundle,
+    id: u64,
+    client_nonce: u128,
+    server_challenge: u128,
+) -> Result<(), CommunicationError> {
+    use mtp_crypto::{auth, verify_ed25519, verify_ml_dsa};
+
+    match response.get_data(DataType::ClientNonce.to_id(tm)) {
+        DataValue::UnsignedNumber(n) if *n == client_nonce => {}
         _ => {
             return Err(CommunicationError::AuthenticationFailed(
                 "Nonce mismatch".into(),
@@ -172,16 +145,7 @@ impl MTPClient {
         }
     }
 
-        let host_new_nonce = match response.get_data(DataTypeId(5)) {
-            DataValue::UnsignedNumber(n) => *n as u128,
-            _ => {
-                return Err(CommunicationError::AuthenticationFailed(
-                    "Missing new nonce".into(),
-                ));
-            }
-        };
-
-        let host_sig = match response.get_data(DataTypeId(10)) {
+    let sig = match response.get_data(DataType::Signature.to_id(tm)) {
         DataValue::Bytes(b) => b.clone(),
         _ => {
             return Err(CommunicationError::AuthenticationFailed(
@@ -189,38 +153,138 @@ impl MTPClient {
             ));
         }
     };
-
-        let host_pq_sig = match response.get_data(DataTypeId(12)) {
+    let pq_sig = match response.get_data(DataType::PqSignature.to_id(tm)) {
         DataValue::Bytes(b) => b.clone(),
         _ => vec![],
     };
 
-        let mut host_sig_payload = Vec::new();
-        host_sig_payload.push(0x01);
-        host_sig_payload.extend_from_slice(&config.client_id.to_be_bytes());
-        host_sig_payload.extend_from_slice(&client_nonce.to_be_bytes());
-        host_sig_payload.extend_from_slice(&host_new_nonce.to_be_bytes());
-
-        /* ===== Signature ===== */
-        verify_ed25519(
-            &host_public_key_bundle.sig_cl_public_key,
-            &host_sig_payload,
-            &host_sig,
-        )
+    let payload = auth::host_final_payload(id, client_nonce, server_challenge);
+    verify_ed25519(&host_pk.sig_cl_public_key, &payload, &sig)
         .map_err(|_| CommunicationError::AuthenticationFailed("Host signature invalid".into()))?;
-        if !host_pq_sig.is_empty()
-            && verify_ml_dsa(
-                &host_public_key_bundle.sig_pq_public_key,
-                &host_sig_payload,
-                &host_pq_sig,
-            )
-            .is_err()
-        {
+    if !pq_sig.is_empty() && verify_ml_dsa(&host_pk.sig_pq_public_key, &payload, &pq_sig).is_err() {
         return Err(CommunicationError::AuthenticationFailed(
             "Host PQ signature invalid".into(),
         ));
     }
-        /* ===== End Signature ===== */
+    Ok(())
+}
+
+/* Interpret the host's `Connected` flag. */
+#[cfg(feature = "crypto")]
+fn check_connected(
+    response: &CommunicationValue,
+    tm: &mtp_codec::TypeMap,
+    reject_msg: &str,
+) -> Result<(), CommunicationError> {
+    match response.get_data(DataType::Connected.to_id(tm)) {
+        DataValue::BoolTrue => Ok(()),
+        DataValue::BoolFalse => Err(CommunicationError::AuthenticationFailed(reject_msg.into())),
+        _ => Err(CommunicationError::AuthenticationFailed(
+            "Invalid response".into(),
+        )),
+    }
+}
+
+#[cfg(feature = "crypto")]
+impl MTPClient {
+    pub async fn auth_connect(
+        config: ClientConfig,
+        keys: &mtp_crypto::Keyring,
+        host_public_key_bundle: &mtp_crypto::PublicKeyBundle,
+    ) -> Result<MTPConnection, CommunicationError> {
+        use mtp_crypto::{Ed25519Signer, MlDsaSigner, SignatureScheme, auth};
+
+        let (sender, receiver) =
+            mtp_transport::connect(&config.url, config.server_cert, Policy::default()).await?;
+
+        let tm = mtp_codec::TypeMap::latest();
+        let version_str = format!("{}", PROTOCOL_VERSION);
+
+        // 1. Send the unsigned Identification hello (version + claimed id).
+        let ident = CommunicationValue::new(mtp_codec::CommunicationType::Identification)
+            .add_typed_default(DataType::Version, DataValue::Str(version_str.clone()))
+            .add_typed_default(
+                DataType::Id,
+                DataValue::UnsignedNumber(config.client_id as u128),
+            );
+        sender.send(&ident).await?;
+
+        // 2. Receive and verify the host's challenge.
+        let challenge = receiver.receive().await?;
+        let expected = mtp_codec::CommunicationType::Challenge.to_id(&tm);
+        if challenge.get_type() != expected {
+            return Err(unexpected_response_type_error(
+                "auth_connect challenge",
+                expected,
+                &challenge,
+            ));
+        }
+        let server_challenge = match challenge.get_data(DataType::ServerNonce.to_id(&tm)) {
+            DataValue::UnsignedNumber(n) => *n,
+            _ => {
+                return Err(CommunicationError::AuthenticationFailed(
+                    "Missing server challenge".into(),
+                ));
+            }
+        };
+        verify_host_challenge(
+            &challenge,
+            &tm,
+            host_public_key_bundle,
+            config.client_id,
+            server_challenge,
+        )?;
+
+        // 3. Sign the host's challenge and send the proof.
+        let client_nonce: u128 = rand::random();
+        let proof_payload = auth::login_proof_payload(
+            &version_str,
+            config.client_id,
+            server_challenge,
+            client_nonce,
+        );
+
+        let signer = Ed25519Signer::new(&keys.sig_cl_secret_key)
+            .map_err(|e| CommunicationError::Other(e.to_string()))?;
+        let signature = signer
+            .sign(&proof_payload)
+            .map_err(|e| CommunicationError::Other(e.to_string()))?;
+
+        let mut proof = CommunicationValue::new(mtp_codec::CommunicationType::ChallengeResponse)
+            .add_typed_default(
+                DataType::ClientNonce,
+                DataValue::UnsignedNumber(client_nonce),
+            )
+            .add_typed_default(DataType::Signature, DataValue::Bytes(signature));
+        if !keys.sig_pq_secret_key.as_bytes().is_empty() {
+            let pq_signer = MlDsaSigner::new(&keys.sig_pq_secret_key, &keys.sig_pq_public_key)
+                .map_err(|e| CommunicationError::Other(e.to_string()))?;
+            let pq_signature = pq_signer
+                .sign(&proof_payload)
+                .map_err(|e| CommunicationError::Other(e.to_string()))?;
+            proof = proof.add_typed_default(DataType::PqSignature, DataValue::Bytes(pq_signature));
+        }
+        sender.send(&proof).await?;
+
+        // 4. Receive and verify the host's final confirmation.
+        let response = receiver.receive().await?;
+        let expected_type = mtp_codec::CommunicationType::IdentificationResponse.to_id(&tm);
+        if response.get_type() != expected_type {
+            return Err(unexpected_response_type_error(
+                "auth_connect",
+                expected_type,
+                &response,
+            ));
+        }
+        check_connected(&response, &tm, "Server rejected authentication")?;
+        verify_host_final(
+            &response,
+            &tm,
+            host_public_key_bundle,
+            config.client_id,
+            client_nonce,
+            server_challenge,
+        )?;
 
         Ok(MTPConnection {
             version: PROTOCOL_VERSION,
@@ -236,58 +300,73 @@ impl MTPClient {
         keys: &mtp_crypto::Keyring,
         host_public_key_bundle: &mtp_crypto::PublicKeyBundle,
     ) -> Result<MTPConnection, CommunicationError> {
-        use mtp_crypto::{
-            Ed25519Signer, MlDsaSigner, SignatureScheme, verify_ed25519, verify_ml_dsa,
-        };
+        use mtp_crypto::{Ed25519Signer, MlDsaSigner, SignatureScheme, auth};
 
         let (sender, receiver) =
             mtp_transport::connect(&config.url, config.server_cert, Policy::default()).await?;
 
-        // 1. Build and send Register message immediately
-        let client_nonce: u128 = rand::random();
+        let tm = mtp_codec::TypeMap::latest();
         let version_str = format!("{}", PROTOCOL_VERSION);
         let pk_bundle = keys.public_key_bundle();
         let pk_bytes = pk_bundle.as_bytes();
 
-        let mut sig_payload = Vec::new();
-        sig_payload.extend_from_slice(version_str.as_bytes());
-        sig_payload.extend_from_slice(&client_nonce.to_be_bytes());
-        sig_payload.extend_from_slice(&pk_bytes);
+        // 1. Send the unsigned Register hello (version + public-key bundle).
+        let register = CommunicationValue::new(mtp_codec::CommunicationType::Register)
+            .add_typed_default(DataType::Version, DataValue::Str(version_str.clone()))
+            .add_typed_default(DataType::PublicKeys, DataValue::Bytes(pk_bytes.clone()));
+        sender.send(&register).await?;
+
+        // 2. Receive and verify the host's challenge (register binds id = 0).
+        let challenge = receiver.receive().await?;
+        let expected = mtp_codec::CommunicationType::Challenge.to_id(&tm);
+        if challenge.get_type() != expected {
+            return Err(unexpected_response_type_error(
+                "auth_register challenge",
+                expected,
+                &challenge,
+            ));
+        }
+        let server_challenge = match challenge.get_data(DataType::ServerNonce.to_id(&tm)) {
+            DataValue::UnsignedNumber(n) => *n,
+            _ => {
+                return Err(CommunicationError::AuthenticationFailed(
+                    "Missing server challenge".into(),
+                ));
+            }
+        };
+        verify_host_challenge(&challenge, &tm, host_public_key_bundle, 0, server_challenge)?;
+
+        // 3. Sign the host's challenge over the bundle and send the proof.
+        let client_nonce: u128 = rand::random();
+        let proof_payload =
+            auth::register_proof_payload(&version_str, &pk_bytes, server_challenge, client_nonce);
 
         let signer = Ed25519Signer::new(&keys.sig_cl_secret_key)
             .map_err(|e| CommunicationError::Other(e.to_string()))?;
-
-        /* ===== Signature ===== */
         let signature = signer
-            .sign(&sig_payload)
+            .sign(&proof_payload)
             .map_err(|e| CommunicationError::Other(e.to_string()))?;
 
-        let mut register = CommunicationValue::new(mtp_codec::CommunicationType::Register)
-            .add_typed_default(DataType::Version, DataValue::Str(version_str))
+        let mut proof = CommunicationValue::new(mtp_codec::CommunicationType::ChallengeResponse)
             .add_typed_default(
                 DataType::ClientNonce,
                 DataValue::UnsignedNumber(client_nonce),
             )
-            .add_typed_default(DataType::PublicKeys, DataValue::Bytes(pk_bytes))
             .add_typed_default(DataType::Signature, DataValue::Bytes(signature));
-
         if !keys.sig_pq_secret_key.as_bytes().is_empty() {
             let pq_signer = MlDsaSigner::new(&keys.sig_pq_secret_key, &keys.sig_pq_public_key)
                 .map_err(|e| CommunicationError::Other(e.to_string()))?;
             let pq_signature = pq_signer
-                .sign(&sig_payload)
+                .sign(&proof_payload)
                 .map_err(|e| CommunicationError::Other(e.to_string()))?;
-            register =
-                register.add_typed_default(DataType::PqSignature, DataValue::Bytes(pq_signature));
+            proof = proof.add_typed_default(DataType::PqSignature, DataValue::Bytes(pq_signature));
         }
-        /* ===== End Signature ===== */
+        sender.send(&proof).await?;
 
-        sender.send(&register).await?;
-
-        // 2. Receive host response (single message)
+        // 4. Receive the host's final confirmation; extract the assigned id and
+        //    verify the host signature binds to it.
         let response = receiver.receive().await?;
-
-        let expected_type = mtp_codec::CommunicationTypeId(18); // RegisterResponse
+        let expected_type = mtp_codec::CommunicationType::RegisterResponse.to_id(&tm);
         if response.get_type() != expected_type {
             return Err(unexpected_response_type_error(
                 "auth_register",
@@ -295,97 +374,30 @@ impl MTPClient {
                 &response,
             ));
         }
-
-        let connected = response.get_data(DataTypeId(11));
-        match connected {
-            DataValue::BoolTrue => {}
-            DataValue::BoolFalse => {
-                return Err(CommunicationError::AuthenticationFailed(
-                    "Server rejected registration".into(),
-                ));
-            }
-            _ => {
-                return Err(CommunicationError::AuthenticationFailed(
-                    "Invalid response".into(),
-                ));
-            }
-        }
-
-        let assigned_id = match response.get_data(DataTypeId(6)) {
-            DataValue::UnsignedNumber(n) => *n as u128,
+        check_connected(&response, &tm, "Server rejected registration")?;
+        let assigned_id = match response.get_data(DataType::Id.to_id(&tm)) {
+            DataValue::UnsignedNumber(n) => *n as u64,
             _ => {
                 return Err(CommunicationError::AuthenticationFailed(
                     "Missing assigned ID".into(),
                 ));
             }
         };
-
-        let echo_nonce = response.get_data(DataTypeId(7));
-        match echo_nonce {
-            DataValue::UnsignedNumber(n) if *n == client_nonce as u128 => {}
-            _ => {
-                return Err(CommunicationError::AuthenticationFailed(
-                    "Nonce mismatch".into(),
-                ));
-            }
-        }
-
-        let host_new_nonce = match response.get_data(DataTypeId(5)) {
-            DataValue::UnsignedNumber(n) => *n as u128,
-            _ => {
-                return Err(CommunicationError::AuthenticationFailed(
-                    "Missing new nonce".into(),
-                ));
-            }
-        };
-
-        let host_sig = match response.get_data(DataTypeId(10)) {
-            DataValue::Bytes(b) => b.clone(),
-            _ => {
-                return Err(CommunicationError::AuthenticationFailed(
-                    "Missing signature".into(),
-                ));
-            }
-        };
-
-        let host_pq_sig = match response.get_data(DataTypeId(12)) {
-            DataValue::Bytes(b) => b.clone(),
-            _ => vec![],
-        };
-
-        let mut host_sig_payload = Vec::new();
-        host_sig_payload.push(0x01);
-        host_sig_payload.extend_from_slice(&(assigned_id as u64).to_be_bytes());
-        host_sig_payload.extend_from_slice(&client_nonce.to_be_bytes());
-        host_sig_payload.extend_from_slice(&host_new_nonce.to_be_bytes());
-
-        /* ===== Signature ===== */
-        verify_ed25519(
-            &host_public_key_bundle.sig_cl_public_key,
-            &host_sig_payload,
-            &host_sig,
-        )
-        .map_err(|_| CommunicationError::AuthenticationFailed("Host signature invalid".into()))?;
-        if !host_pq_sig.is_empty()
-            && verify_ml_dsa(
-                &host_public_key_bundle.sig_pq_public_key,
-                &host_sig_payload,
-                &host_pq_sig,
-            )
-            .is_err()
-        {
-            return Err(CommunicationError::AuthenticationFailed(
-                "Host PQ signature invalid".into(),
-            ));
-        }
-        /* ===== End Signature ===== */
+        verify_host_final(
+            &response,
+            &tm,
+            host_public_key_bundle,
+            assigned_id,
+            client_nonce,
+            server_challenge,
+        )?;
 
         Ok(MTPConnection {
             version: PROTOCOL_VERSION,
             sender,
             receiver,
             auth_state: AuthState::Authenticated,
-            client_id: assigned_id as u64,
+            client_id: assigned_id,
         })
     }
 }

codec/Cargo.toml

@@ -14,4 +14,4 @@ rand = { version = "0.8", features = ["std", "std_rng"] }
 [features]
 default = []
 registry = ["mtp-type-map/registry"]
-crypto = ["dep:mtp-crypto"]
+crypto = ["dep:mtp-crypto", "mtp-crypto/mlkem-tls"]

codec/src/communication_value.rs

@@ -146,13 +146,30 @@ impl CommunicationValue {
      * bit3 => is data encrypted If so data bytes will be an encrypted container
      * bit4 => is communication value signed
      */
-    pub fn to_bytes(&self) -> Vec<u8> {
+    /*
+     * Build the canonical metadata header and data payload shared by both
+     * `to_bytes` and `build_signed_payload`. Keeping a single source here
+     * guarantees the serialized frame and the signed-over bytes stay in sync.
+     *
+     * Returns `(metadata, data_bytes)` where
+     *   metadata = comm_type || flags || id? || sender? || receiver?
+     *
+     * `force_signed` forces the `FLAG_SIGNED` bit on regardless of whether a
+     * signature is currently attached. The signed-payload path passes `true` so
+     * that the bytes signed by `sign_frame` (before the signature is stored) and
+     * the bytes verified by `verify_frame` (after it is stored) are identical.
+     */
+    fn build_metadata_and_data(
+        &self,
+        force_signed: bool,
+    ) -> Result<(Vec<u8>, Vec<u8>), CodecError> {
         let has_sender = self.sender != 0;
         let has_receiver = self.receiver != 0;
         let has_id = self.id != 0;
 
         #[cfg(feature = "crypto")]
-        let is_encrypted = self.data.len() == 1 && self.data.values().any(|v| {
+        let is_encrypted = self.data.len() == 1
+            && self.data.values().any(|v| {
                 matches!(
                     v,
                     DataValue::EncryptedContainer(_) | DataValue::SignedEncryptedContainer(_)
@@ -179,7 +196,7 @@ impl CommunicationValue {
         if is_encrypted {
             flags |= FLAG_ENCRYPTED;
         }
-        if has_frame_sig {
+        if has_frame_sig || force_signed {
             flags |= FLAG_SIGNED;
         }
 
@@ -212,34 +229,39 @@ impl CommunicationValue {
                 })
                 .unwrap_or_default()
         } else {
-            let container_value = DataValue::container_from_map(&self.data);
-            container_value.to_bytes()
+            DataValue::container_from_map(&self.data).to_bytes()?
         };
 
         #[cfg(not(feature = "crypto"))]
-        let data_bytes = {
-            let container_value = DataValue::container_from_map(&self.data);
-            container_value.to_bytes()
-        };
+        let data_bytes = DataValue::container_from_map(&self.data).to_bytes()?;
+
+        Ok((metadata, data_bytes))
+    }
+
+    pub fn to_bytes(&self) -> Result<Vec<u8>, CodecError> {
+        let (metadata, data_bytes) = self.build_metadata_and_data(false)?;
 
         let mut payload = Vec::new();
         payload.extend_from_slice(&metadata);
 
         #[cfg(feature = "crypto")]
-        if let Some((_alg, _sig)) = &self.frame_signature {
+        if let Some((alg, sig)) = &self.frame_signature {
             // algorithm and signature are computed by sign_frame() and stored.
             // The frame bytes are built by using the pre-computed signature.
-            payload.push(*_alg);
-            payload.extend_from_slice(_sig);
+            payload.push(*alg);
+            payload.extend_from_slice(sig);
         }
 
         payload.extend_from_slice(&data_bytes);
 
+        let len = u32::try_from(payload.len()).map_err(|_| CodecError::TooManyEntries)?;
         let mut frame = Vec::with_capacity(4 + payload.len());
-        let _ = frame.write_u32::<BigEndian>(payload.len() as u32);
+        frame
+            .write_u32::<BigEndian>(len)
+            .map_err(|_| CodecError::InvalidEncoding)?;
         frame.extend_from_slice(&payload);
 
-        frame
+        Ok(frame)
     }
 
     pub fn from_bytes(bytes: &[u8]) -> Result<Self, CodecError> {
@@ -302,8 +324,7 @@ impl CommunicationValue {
         #[cfg(feature = "crypto")]
         let frame_signature = if is_signed {
             let alg = cursor.read_u8().map_err(|_| CodecError::InvalidEncoding)?;
-            let sig_len =
-                SigAlgorithm::length(alg).ok_or(CodecError::InvalidEncoding)?;
+            let sig_len = SigAlgorithm::length(alg).ok_or(CodecError::InvalidEncoding)?;
             let mut sig = vec![0u8; sig_len];
             cursor
                 .read_exact(&mut sig)
@@ -324,7 +345,7 @@ impl CommunicationValue {
         let data = if is_encrypted {
             let mut map = BTreeMap::new();
             map.insert(
-                DataTypeId(0),
+                DataType::Version.to_id(&TypeMap::latest()),
                 DataValue::EncryptedContainer(data_bytes.to_vec()),
             );
             map
@@ -367,12 +388,8 @@ impl CommunicationValue {
      * signature before the data payload.
      */
     #[cfg(feature = "crypto")]
-    pub fn sign_frame(
-        &mut self,
-        algorithm: u8,
-        signer: &impl SignatureScheme,
-    ) -> Option<()> {
-        let signed_payload = self.build_signed_payload();
+    pub fn sign_frame(&mut self, algorithm: u8, signer: &impl SignatureScheme) -> Option<()> {
+        let signed_payload = self.build_signed_payload().ok()?;
         let sig = signer.sign(&signed_payload).ok()?;
         self.frame_signature = Some((algorithm, sig));
         Some(())
@@ -389,7 +406,7 @@ impl CommunicationValue {
             .as_ref()
             .ok_or(CodecError::InvalidEncoding)?;
 
-        let signed_payload = self.build_signed_payload();
+        let signed_payload = self.build_signed_payload()?;
         verifier
             .verify(&signed_payload, sig)
             .map_err(|_| CodecError::InvalidEncoding)
@@ -400,75 +417,11 @@ impl CommunicationValue {
      *   comm_type || flags || id? || sender? || receiver? || data_bytes
      */
     #[cfg(feature = "crypto")]
-    fn build_signed_payload(&self) -> Vec<u8> {
-        let has_sender = self.sender != 0;
-        let has_receiver = self.receiver != 0;
-        let has_id = self.id != 0;
-
-        let is_encrypted = self.data.len() == 1 && self.data.values().any(|v| {
-            matches!(
-                v,
-                DataValue::EncryptedContainer(_) | DataValue::SignedEncryptedContainer(_)
-            )
-        });
-
-        let mut flags: u8 = 0;
-        if has_sender {
-            flags |= FLAG_HAS_SENDER;
-        }
-        if has_receiver {
-            flags |= FLAG_HAS_RECEIVER;
-        }
-        if has_id {
-            flags |= FLAG_HAS_ID;
-        }
-        if is_encrypted {
-            flags |= FLAG_ENCRYPTED;
-        }
-        if self.frame_signature.is_some() {
-            flags |= FLAG_SIGNED;
-        }
-
-        let mut metadata = Vec::new();
-        let _ = metadata.write_u16::<BigEndian>(self.comm_type.0);
-        metadata.push(flags);
-
-        if has_id {
-            let _ = metadata.write_u32::<BigEndian>(self.id);
-        }
-
-        if has_sender {
-            let sender_be = self.sender.to_be_bytes();
-            metadata.extend_from_slice(&sender_be[2..]);
-        }
-
-        if has_receiver {
-            let receiver_be = self.receiver.to_be_bytes();
-            metadata.extend_from_slice(&receiver_be[2..]);
-        }
-
-        #[cfg(feature = "crypto")]
-        let data_bytes = if is_encrypted {
-            self.data
-                .values()
-                .find_map(|v| match v {
-                    DataValue::EncryptedContainer(ct) => Some(ct.clone()),
-                    DataValue::SignedEncryptedContainer(ct) => Some(ct.clone()),
-                    _ => None,
-                })
-                .unwrap_or_default()
-        } else {
-            let container_value = DataValue::container_from_map(&self.data);
-            container_value.to_bytes()
-        };
-
-        #[cfg(not(feature = "crypto"))]
-        let data_bytes = {
-            let container_value = DataValue::container_from_map(&self.data);
-            container_value.to_bytes()
-        };
-
-        [metadata, data_bytes].concat()
+    fn build_signed_payload(&self) -> Result<Vec<u8>, CodecError> {
+        // Force FLAG_SIGNED on so the signed bytes match whether or not the
+        // signature has been attached yet (sign_frame runs before storing it).
+        let (metadata, data_bytes) = self.build_metadata_and_data(true)?;
+        Ok([metadata, data_bytes].concat())
     }
 
     #[cfg(feature = "crypto")]
@@ -610,9 +563,9 @@ mod tests {
     use crate::data_value::DataValue;
 
     fn roundtrip(cv: CommunicationValue) -> CommunicationValue {
-        let bytes = cv.to_bytes();
+        let bytes = cv.to_bytes().expect("encode failed");
         let decoded = CommunicationValue::from_bytes(&bytes).expect("failed to deserialize");
-        let bytes2 = decoded.to_bytes();
+        let bytes2 = decoded.to_bytes().expect("encode failed");
         assert_eq!(bytes, bytes2);
         decoded
     }
@@ -620,7 +573,7 @@ mod tests {
     #[test]
     fn test_flags_and_order_without_optional() {
         let cv = CommunicationValue::new(CommunicationType::ErrorParsing).with_id(0);
-        let bytes = cv.to_bytes();
+        let bytes = cv.to_bytes().expect("encode failed");
 
         // [u32 len][u16 type][flags]...
         assert!(bytes.len() >= 7);
@@ -629,7 +582,7 @@ mod tests {
         assert_eq!(total_len as usize + 4, bytes.len());
 
         let typ = c.read_u16::<BigEndian>().expect("read type");
-        assert_eq!(typ, 1);
+        assert_eq!(typ, 12);
 
         let flags = c.read_u8().expect("read flags");
         assert_eq!(flags & 0b0000_0111, 0);
@@ -642,14 +595,14 @@ mod tests {
             .with_sender(0x0000_1122_3344_5566)
             .with_receiver(0x0000_6677_8899_AABB);
 
-        let bytes = cv.to_bytes();
+        let bytes = cv.to_bytes().expect("encode failed");
         let mut c = Cursor::new(bytes.as_slice());
 
         let total_len = c.read_u32::<BigEndian>().expect("len");
         assert_eq!(total_len as usize + 4, bytes.len());
 
         let typ = c.read_u16::<BigEndian>().expect("read type");
-        assert_eq!(typ, 2);
+        assert_eq!(typ, 13);
 
         let flags = c.read_u8().expect("read flags");
         assert_eq!(flags & 0b0000_0111, 0b0000_0111);
@@ -668,15 +621,16 @@ mod tests {
 
     #[test]
     fn test_roundtrip_complex() {
+        let tm = TypeMap::latest();
         let cv = CommunicationValue::new(CommunicationType::Disconnect)
             .with_id(1234)
             .with_sender(111)
             .with_receiver(222)
-            .add_data(DataTypeId(1), DataValue::Str("alice".to_string()))
-            .add_data(DataTypeId(2), DataValue::SignedNumber(42))
-            .add_data(DataTypeId(3), DataValue::BoolTrue)
-            .add_data(
-                DataTypeId(4),
+            .add_typed_default(DataType::Id, DataValue::Str("alice".to_string()))
+            .add_typed_default(DataType::ClientNonce, DataValue::SignedNumber(42))
+            .add_typed_default(DataType::ServerNonce, DataValue::BoolTrue)
+            .add_typed_default(
+                DataType::PublicKeys,
                 DataValue::Array(vec![DataValue::SignedNumber(1), DataValue::SignedNumber(2)]),
             );
 
@@ -685,13 +639,13 @@ mod tests {
         assert_eq!(decoded.get_id(), 1234);
         assert_eq!(decoded.get_sender(), 111);
         assert_eq!(decoded.get_receiver(), 222);
-        assert_eq!(decoded.get_type(), CommunicationTypeId(3));
+        assert_eq!(decoded.get_type(), CommunicationType::Disconnect.to_id(&tm));
         assert_eq!(
-            decoded.get_data(DataTypeId(1)),
+            decoded.get_data(DataType::Id.to_id(&tm)),
             &DataValue::Str("alice".to_string())
         );
         assert_eq!(
-            decoded.get_data(DataTypeId(2)),
+            decoded.get_data(DataType::ClientNonce.to_id(&tm)),
             &DataValue::SignedNumber(42)
         );
     }
@@ -703,4 +657,45 @@ mod tests {
         bad[0..4].copy_from_slice(&(1000u32.to_be_bytes()));
         assert!(CommunicationValue::from_bytes(&bad).is_err());
     }
+
+    #[cfg(feature = "crypto")]
+    #[test]
+    fn test_sign_verify_frame_roundtrip() {
+        use mtp_crypto::{Ed25519Signer, SigAlgorithm};
+
+        let (signer, sk, _pk) = Ed25519Signer::generate();
+
+        let mut cv = CommunicationValue::new(CommunicationType::Ping)
+            .with_id(7)
+            .with_sender(1)
+            .with_receiver(2)
+            .add_typed_default(DataType::PqSignature, DataValue::UnsignedNumber(42));
+
+        assert!(cv.sign_frame(SigAlgorithm::ED25519, &signer).is_some());
+
+        // Same in-memory value verifies (FLAG_SIGNED forced on both sides).
+        let verifier = Ed25519Signer::new(&sk).unwrap();
+        assert!(cv.verify_frame(&verifier).is_ok());
+
+        // Survives a wire round-trip.
+        let bytes = cv.to_bytes().expect("encode failed");
+        let decoded = CommunicationValue::from_bytes(&bytes).expect("decode failed");
+        assert!(decoded.verify_frame(&verifier).is_ok());
+    }
+
+    #[cfg(feature = "crypto")]
+    #[test]
+    fn test_verify_frame_wrong_key_fails() {
+        use mtp_crypto::{Ed25519Signer, SigAlgorithm};
+
+        let (signer, _, _) = Ed25519Signer::generate();
+        let (_, other_sk, _) = Ed25519Signer::generate();
+
+        let mut cv = CommunicationValue::new(CommunicationType::Ping)
+            .add_typed_default(DataType::PqSignature, DataValue::UnsignedNumber(42));
+        assert!(cv.sign_frame(SigAlgorithm::ED25519, &signer).is_some());
+
+        let wrong = Ed25519Signer::new(&other_sk).unwrap();
+        assert!(cv.verify_frame(&wrong).is_err());
+    }
 }

codec/src/data_value.rs

@@ -6,10 +6,14 @@ use std::fmt;
 use std::hash::{Hash, Hasher};
 use std::io::Cursor;
 
+use mtp_common::CodecError;
 use mtp_type_map::DataTypeId;
 
+#[cfg(test)]
+use mtp_type_map::{DataType, TypeMap};
+
 #[cfg(feature = "crypto")]
-use mtp_crypto::{AeadDecrypt, AeadEncrypt, SigAlgorithm, SignatureScheme};
+use mtp_crypto::{EncryptionType, Keyring, PublicKeyBundle, SigAlgorithm, SignatureScheme};
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum DataKind {
@@ -124,9 +128,9 @@ impl DataValue {
      * 0x07 => Bytes
      * 0x08 => Array
      * 0x09 => Container
-     * 0x0A => EncryptedContainer (4 bytes u32 len + encrypted bytes)
-     * 0x0B => SignedContainer (4 bytes u32 len + 3373 bytes signature)
-     * 0x0C => SignedEncryptedContainer (4 bytes u32 len + 3373 bytes signature + encrypted bytes)
+     * 0x0A => EncryptedContainer (1 byte EncryptionType + KEM ciphertext + AEAD payload)
+     * 0x0B => SignedContainer (1 byte SigAlgorithm + signature + serialized container)
+     * 0x0C => SignedEncryptedContainer (encrypted blob that decrypts to a SignedContainer)
      * 0xFF => Null
      */
     const KIND_BOOL_TRUE: u8 = 0x01;
@@ -150,10 +154,19 @@ impl DataValue {
 
     const KIND_NULL: u8 = 0xFF;
 
+    /*
+     * Smallest possible encoded entry, used to cap pre-reservation when
+     * decoding containers/arrays so a small frame cannot force a huge
+     * allocation from an attacker-controlled count. A bool/null entry in a
+     * container is 3 bytes (1 kind + 2 key); a bare value in an array is 1
+     * byte, so 1 is the safe lower bound shared by both.
+     */
+    const MIN_ENTRY_BYTES: usize = 1;
+
     pub fn container_from_map(map: &BTreeMap<DataTypeId, DataValue>) -> DataValue {
         let mut container = Vec::new();
         for (key, value) in map {
-            container.push((key.clone(), value.clone()));
+            container.push((*key, value.clone()));
         }
         DataValue::Container(container)
     }
@@ -272,13 +285,14 @@ impl DataValue {
 
     /*
      * Decrypt an `EncryptedContainer` in-place, replacing it with the
-     * deserialized `Container`. Returns `None` if decryption or
-     * deserialization fails.
+     * deserialized `Container`. The algorithm (and which keypair to use) is read
+     * from the blob's leading `EncryptionType` byte; the matching key is taken
+     * from `keyring`. Returns `None` if decryption or deserialization fails.
      */
     #[cfg(feature = "crypto")]
-    pub fn decrypt_into_container(&mut self, cipher: &impl AeadDecrypt, aad: &[u8]) -> Option<()> {
+    pub fn decrypt_into_container(&mut self, keyring: &Keyring, aad: &[u8]) -> Option<()> {
         let data = self.as_encrypted_container()?;
-        let plaintext = cipher.decrypt(&data, aad).ok()?;
+        let plaintext = mtp_crypto::decrypt_with(&data, keyring, aad).ok()?;
         let dv = DataValue::from_bytes(&plaintext)?;
         match dv {
             DataValue::Container(entries) => {
@@ -291,13 +305,21 @@ impl DataValue {
 
     /*
      * Encrypt a `Container` into an `EncryptedContainer` in-place.
+     * `enc_type` selects the algorithm and `recipient` provides the public key
+     * encapsulated to. The resulting blob is self-describing: its leading byte
+     * is `enc_type`, so `decrypt_into_container` needs only a `Keyring`.
      * Returns `None` if the value is not a `Container` or encryption fails.
      */
     #[cfg(feature = "crypto")]
-    pub fn encrypt_container(&mut self, cipher: &impl AeadEncrypt, aad: &[u8]) -> Option<()> {
+    pub fn encrypt_container(
+        &mut self,
+        enc_type: EncryptionType,
+        recipient: &PublicKeyBundle,
+        aad: &[u8],
+    ) -> Option<()> {
         let entries = self.as_container()?;
-        let plaintext = DataValue::Container(entries).to_bytes();
-        let ct = cipher.encrypt(&plaintext, aad).ok()?;
+        let plaintext = DataValue::Container(entries).to_bytes().ok()?;
+        let ct = mtp_crypto::encrypt_for(enc_type, recipient, &plaintext, aad).ok()?;
         *self = DataValue::EncryptedContainer(ct);
         Some(())
     }
@@ -311,7 +333,7 @@ impl DataValue {
     #[cfg(feature = "crypto")]
     pub fn sign_container(&mut self, algorithm: u8, signer: &impl SignatureScheme) -> Option<()> {
         let entries = self.as_container()?;
-        let container_bytes = Self::encode_container(&entries);
+        let container_bytes = Self::encode_container(&entries).ok()?;
 
         let sig = signer.sign(&container_bytes).ok()?;
 
@@ -354,35 +376,40 @@ impl DataValue {
 
     /*
      * Encrypt a `Container` into a `SignedEncryptedContainer` in-place.
-     * The result is an opaque ciphertext that decrypts to a `SignedContainer`.
+     * The container is first signed (with `algorithm`/`signer`), then the signed
+     * blob is encrypted with `enc_type` to `recipient`. The result is an opaque
+     * ciphertext that decrypts to a `SignedContainer`.
      */
     #[cfg(feature = "crypto")]
     pub fn sign_and_encrypt_container(
         &mut self,
         algorithm: u8,
         signer: &impl SignatureScheme,
-        cipher: &impl AeadEncrypt,
+        enc_type: EncryptionType,
+        recipient: &PublicKeyBundle,
         aad: &[u8],
     ) -> Option<()> {
         self.sign_container(algorithm, signer)?;
         let blob = self.as_signed_container()?;
-        let ct = cipher.encrypt(&blob, aad).ok()?;
+        let ct = mtp_crypto::encrypt_for(enc_type, recipient, &blob, aad).ok()?;
         *self = DataValue::SignedEncryptedContainer(ct);
         Some(())
     }
 
     /*
      * Decrypt a `SignedEncryptedContainer` in-place, replacing it with a
-     * `SignedContainer`. Does NOT verify; call `verify_into_container` next.
+     * `SignedContainer`. The algorithm and keypair are resolved from the blob's
+     * leading byte and `keyring`. Does NOT verify; call `verify_into_container`
+     * next.
      */
     #[cfg(feature = "crypto")]
     pub fn decrypt_signed_encrypted_container(
         &mut self,
-        cipher: &impl AeadDecrypt,
+        keyring: &Keyring,
         aad: &[u8],
     ) -> Option<()> {
         let data = self.as_signed_encrypted_container()?;
-        let plaintext = cipher.decrypt(&data, aad).ok()?;
+        let plaintext = mtp_crypto::decrypt_with(&data, keyring, aad).ok()?;
         *self = DataValue::SignedContainer(plaintext);
         Some(())
     }
@@ -392,7 +419,7 @@ impl DataValue {
             DataValue::Container(c) => {
                 let mut out = BTreeMap::new();
                 for (k, v) in c {
-                    out.insert(k.clone(), v.clone());
+                    out.insert(*k, v.clone());
                 }
                 Some(out)
             }
@@ -400,16 +427,14 @@ impl DataValue {
         }
     }
 
-    pub fn to_bytes(&self) -> Vec<u8> {
+    pub fn to_bytes(&self) -> Result<Vec<u8>, CodecError> {
         match self {
             DataValue::Container(entries) => Self::encode_container(entries),
             DataValue::Array(arr) => Self::encode_array(arr),
             _ => {
                 let mut out = Vec::new();
-                if Self::write_value_payload(&mut out, self).is_none() {
-                    return Vec::new();
-                }
-                out
+                Self::write_value_payload(&mut out, self)?;
+                Ok(out)
             }
         }
     }
@@ -423,8 +448,8 @@ impl DataValue {
         Some(value)
     }
 
-    pub fn to_base64(&self) -> String {
-        general_purpose::STANDARD.encode(self.to_bytes())
+    pub fn to_base64(&self) -> Result<String, CodecError> {
+        Ok(general_purpose::STANDARD.encode(self.to_bytes()?))
     }
 
     pub fn from_base64(base64_str: &str) -> Option<Self> {
@@ -432,144 +457,138 @@ impl DataValue {
         Self::from_bytes(&bytes)
     }
 
-    fn encode_container(entries: &[(DataTypeId, DataValue)]) -> Vec<u8> {
+    fn encode_container(entries: &[(DataTypeId, DataValue)]) -> Result<Vec<u8>, CodecError> {
         let mut out = Vec::new();
-        if out
-            .write_u16::<BigEndian>(u16::try_from(entries.len()).ok().unwrap_or(0))
-            .is_err()
-        {
-            return Vec::new();
-        }
+        let count = u16::try_from(entries.len()).map_err(|_| CodecError::TooManyEntries)?;
+        out.write_u16::<BigEndian>(count)
+            .map_err(|_| CodecError::InvalidEncoding)?;
 
         for (key, value) in entries {
-            if !Self::write_container_entry(&mut out, key.clone(), value) {
-                return Vec::new();
+            Self::write_container_entry(&mut out, *key, value)?;
         }
-        }
-        out
+        Ok(out)
     }
 
-    fn write_container_entry(buf: &mut Vec<u8>, key: DataTypeId, value: &DataValue) -> bool {
+    fn write_container_entry(
+        buf: &mut Vec<u8>,
+        key: DataTypeId,
+        value: &DataValue,
+    ) -> Result<(), CodecError> {
         let kind = Self::kind_marker(value);
         buf.push(kind);
 
         if kind == Self::KIND_BOOL_TRUE || kind == Self::KIND_BOOL_FALSE || kind == Self::KIND_NULL
         {
-            let _ = buf.write_u16::<BigEndian>(key.0);
-            return true;
+            buf.write_u16::<BigEndian>(key.0)
+                .map_err(|_| CodecError::InvalidEncoding)?;
+            return Ok(());
         }
 
         let mut payload = Vec::new();
-        if Self::write_value_payload(&mut payload, value).is_none() {
-            return false;
-        }
+        Self::write_value_payload(&mut payload, value)?;
 
-        if buf.write_u32::<BigEndian>(payload.len() as u32).is_err() {
-            return false;
-        }
-        let _ = buf.write_u16::<BigEndian>(key.0);
+        let len = u32::try_from(payload.len()).map_err(|_| CodecError::TooManyEntries)?;
+        buf.write_u32::<BigEndian>(len)
+            .map_err(|_| CodecError::InvalidEncoding)?;
+        buf.write_u16::<BigEndian>(key.0)
+            .map_err(|_| CodecError::InvalidEncoding)?;
         buf.extend_from_slice(&payload);
-        true
+        Ok(())
     }
 
-    fn encode_array(arr: &[DataValue]) -> Vec<u8> {
+    fn encode_array(arr: &[DataValue]) -> Result<Vec<u8>, CodecError> {
         let mut out = Vec::new();
-        if out
-            .write_u16::<BigEndian>(u16::try_from(arr.len()).ok().unwrap_or(0))
-            .is_err()
-        {
-            return Vec::new();
-        }
+        let count = u16::try_from(arr.len()).map_err(|_| CodecError::TooManyEntries)?;
+        out.write_u16::<BigEndian>(count)
+            .map_err(|_| CodecError::InvalidEncoding)?;
 
         for value in arr {
-            if !Self::write_array_entry(&mut out, value) {
-                return Vec::new();
-            }
+            Self::write_array_entry(&mut out, value)?;
         }
 
-        out
+        Ok(out)
     }
 
-    fn write_array_entry(buf: &mut Vec<u8>, value: &DataValue) -> bool {
+    fn write_array_entry(buf: &mut Vec<u8>, value: &DataValue) -> Result<(), CodecError> {
         let kind = Self::kind_marker(value);
         buf.push(kind);
 
         if kind == Self::KIND_BOOL_TRUE || kind == Self::KIND_BOOL_FALSE || kind == Self::KIND_NULL
         {
-            return true;
+            return Ok(());
         }
 
         let mut payload = Vec::new();
-        if Self::write_value_payload(&mut payload, value).is_none() {
-            return false;
-        }
+        Self::write_value_payload(&mut payload, value)?;
 
-        if buf.write_u32::<BigEndian>(payload.len() as u32).is_err() {
-            return false;
-        }
+        let len = u32::try_from(payload.len()).map_err(|_| CodecError::TooManyEntries)?;
+        buf.write_u32::<BigEndian>(len)
+            .map_err(|_| CodecError::InvalidEncoding)?;
         buf.extend_from_slice(&payload);
-        true
+        Ok(())
     }
 
-    fn write_value_payload(buf: &mut Vec<u8>, value: &DataValue) -> Option<()> {
+    fn write_value_payload(buf: &mut Vec<u8>, value: &DataValue) -> Result<(), CodecError> {
         match value {
-            DataValue::BoolTrue => Some(()),
-            DataValue::BoolFalse => Some(()),
+            DataValue::BoolTrue => Ok(()),
+            DataValue::BoolFalse => Ok(()),
+            #[allow(clippy::if_same_then_else)]
             DataValue::Bool(v) => {
-                if *v {
-                    Some(())
-                } else {
-                    Some(())
-                }
+                // Kept intentionally: the kind marker already encodes the boolean,
+                // so both arms carry no payload. Retained for clear compatibility.
+                if *v { Ok(()) } else { Ok(()) }
             }
             DataValue::SignedNumber(n) => {
-                buf.write_i128::<BigEndian>(*n).ok()?;
-                Some(())
+                buf.write_i128::<BigEndian>(*n)
+                    .map_err(|_| CodecError::InvalidEncoding)?;
+                Ok(())
             }
             DataValue::UnsignedNumber(n) => {
-                buf.write_u128::<BigEndian>(*n).ok()?;
-                Some(())
+                buf.write_u128::<BigEndian>(*n)
+                    .map_err(|_| CodecError::InvalidEncoding)?;
+                Ok(())
             }
             DataValue::Float(a, b) => {
-                buf.write_u8(*a).ok()?;
-                buf.write_u32::<BigEndian>(*b).ok()?;
-                Some(())
+                buf.write_u8(*a).map_err(|_| CodecError::InvalidEncoding)?;
+                buf.write_u32::<BigEndian>(*b)
+                    .map_err(|_| CodecError::InvalidEncoding)?;
+                Ok(())
             }
             DataValue::Str(s) => {
                 buf.extend_from_slice(s.as_bytes());
-                Some(())
+                Ok(())
             }
             DataValue::Array(arr) => {
-                let bytes = Self::encode_array(arr);
+                let bytes = Self::encode_array(arr)?;
                 buf.extend_from_slice(&bytes);
-                Some(())
+                Ok(())
             }
             DataValue::Bytes(b) => {
                 buf.extend_from_slice(b);
-                Some(())
+                Ok(())
             }
             DataValue::Container(entries) => {
-                let bytes = Self::encode_container(entries);
+                let bytes = Self::encode_container(entries)?;
                 buf.extend_from_slice(&bytes);
-                Some(())
+                Ok(())
             }
             #[cfg(feature = "crypto")]
             DataValue::EncryptedContainer(data) => {
                 buf.extend_from_slice(data);
-                Some(())
+                Ok(())
             }
             #[cfg(feature = "crypto")]
             DataValue::SignedContainer(data) => {
                 buf.extend_from_slice(data);
-                Some(())
+                Ok(())
             }
             #[cfg(feature = "crypto")]
             DataValue::SignedEncryptedContainer(data) => {
                 buf.extend_from_slice(data);
-                Some(())
+                Ok(())
             }
 
-            DataValue::Null => Some(()),
+            DataValue::Null => Ok(()),
         }
     }
 
@@ -597,7 +616,11 @@ impl DataValue {
 
     fn try_read_container(cursor: &mut Cursor<&[u8]>) -> Option<Self> {
         let count = cursor.read_u16::<BigEndian>().ok()? as usize;
-        let mut entries = Vec::with_capacity(count);
+        let remaining = cursor
+            .get_ref()
+            .len()
+            .saturating_sub(cursor.position() as usize);
+        let mut entries = Vec::with_capacity(count.min(remaining / Self::MIN_ENTRY_BYTES));
 
         for _ in 0..count {
             let kind = cursor.read_u8().ok()?;
@@ -643,7 +666,11 @@ impl DataValue {
 
     fn read_array(cursor: &mut Cursor<&[u8]>) -> Option<Self> {
         let count = cursor.read_u16::<BigEndian>().ok()? as usize;
-        let mut out = Vec::with_capacity(count);
+        let remaining = cursor
+            .get_ref()
+            .len()
+            .saturating_sub(cursor.position() as usize);
+        let mut out = Vec::with_capacity(count.min(remaining / Self::MIN_ENTRY_BYTES));
 
         for _ in 0..count {
             let kind = cursor.read_u8().ok()?;
@@ -901,62 +928,28 @@ impl PartialEq for DataValue {
 impl Hash for DataValue {
     fn hash<H: Hasher>(&self, state: &mut H) {
         use DataValue::*;
+        // Use the wire kind marker as the per-variant discriminant. It is unique
+        // per kind and maps BoolTrue/Bool(true) (and BoolFalse/Bool(false)) to the
+        // same marker, keeping the hash consistent with the Eq bool equivalence.
+        Self::kind_marker(self).hash(state);
         match self {
-            BoolTrue | Bool(true) => {
-                0u8.hash(state);
-                true.hash(state);
-            }
-            BoolFalse | Bool(false) => {
-                0u8.hash(state);
-                false.hash(state);
-            }
-            SignedNumber(n) => {
-                1u8.hash(state);
-                n.hash(state);
-            }
-            UnsignedNumber(n) => {
-                2u8.hash(state);
-                n.hash(state);
-            }
+            BoolTrue | BoolFalse | Bool(_) | Null => {}
+            SignedNumber(n) => n.hash(state),
+            UnsignedNumber(n) => n.hash(state),
             Float(n, m) => {
-                3u8.hash(state);
                 n.hash(state);
                 m.hash(state);
             }
-            Str(s) => {
-                2u8.hash(state);
-                s.hash(state);
-            }
-            Array(a) => {
-                3u8.hash(state);
-                a.hash(state);
-            }
-            Bytes(a) => {
-                4u8.hash(state);
-                a.hash(state);
-            }
-            Container(c) => {
-                5u8.hash(state);
-                c.hash(state);
-            }
+            Str(s) => s.hash(state),
+            Array(a) => a.hash(state),
+            Bytes(a) => a.hash(state),
+            Container(c) => c.hash(state),
             #[cfg(feature = "crypto")]
-            EncryptedContainer(c) => {
-                6u8.hash(state);
-                c.hash(state);
-            }
+            EncryptedContainer(c) => c.hash(state),
             #[cfg(feature = "crypto")]
-            SignedContainer(c) => {
-                7u8.hash(state);
-                c.hash(state);
-            }
+            SignedContainer(c) => c.hash(state),
             #[cfg(feature = "crypto")]
-            SignedEncryptedContainer(c) => {
-                8u8.hash(state);
-                c.hash(state);
-            }
-            Null => {
-                9u8.hash(state);
-            }
+            SignedEncryptedContainer(c) => c.hash(state),
         }
     }
 }
@@ -966,27 +959,26 @@ impl Hash for DataValue {
 mod tests {
     use super::*;
 
-    /// Only Container and Array can be top-level serialized forms.
-    /// Scalars must be tested inside a container.
     fn container_roundtrip(values: Vec<(DataTypeId, DataValue)>) {
         let dv = DataValue::Container(values.clone());
-        let bytes = dv.to_bytes();
+        let bytes = dv.to_bytes().expect("encode failed");
         let decoded = DataValue::from_bytes(&bytes).expect("roundtrip failed");
         assert_eq!(dv, decoded, "container roundtrip mismatch");
     }
 
     fn array_roundtrip(values: Vec<DataValue>) {
         let dv = DataValue::Array(values.clone());
-        let bytes = dv.to_bytes();
+        let bytes = dv.to_bytes().expect("encode failed");
         let decoded = DataValue::from_bytes(&bytes).expect("roundtrip failed");
         assert_eq!(dv, decoded, "array roundtrip mismatch");
     }
 
     #[test]
     fn test_bool_in_container() {
+        let tm = TypeMap::latest();
         container_roundtrip(vec![
-            (DataTypeId(1), DataValue::BoolTrue),
-            (DataTypeId(2), DataValue::BoolFalse),
+            (DataType::Id.to_id(&tm), DataValue::BoolTrue),
+            (DataType::ClientNonce.to_id(&tm), DataValue::BoolFalse),
         ]);
     }
 
@@ -1007,54 +999,60 @@ mod tests {
 
     #[test]
     fn test_signed_number_in_container() {
+        let tm = TypeMap::latest();
         container_roundtrip(vec![
-            (DataTypeId(1), DataValue::SignedNumber(0)),
-            (DataTypeId(2), DataValue::SignedNumber(42)),
-            (DataTypeId(3), DataValue::SignedNumber(-42)),
-            (DataTypeId(4), DataValue::SignedNumber(i128::MAX)),
-            (DataTypeId(5), DataValue::SignedNumber(i128::MIN)),
+            (DataType::Version.to_id(&tm), DataValue::SignedNumber(0)),
+            (DataType::Id.to_id(&tm), DataValue::SignedNumber(42)),
+            (DataType::ClientNonce.to_id(&tm), DataValue::SignedNumber(-42)),
+            (DataType::ServerNonce.to_id(&tm), DataValue::SignedNumber(i128::MAX)),
+            (DataType::PublicKeys.to_id(&tm), DataValue::SignedNumber(i128::MIN)),
         ]);
     }
 
     #[test]
     fn test_unsigned_number_in_container() {
+        let tm = TypeMap::latest();
         container_roundtrip(vec![
-            (DataTypeId(1), DataValue::UnsignedNumber(0)),
-            (DataTypeId(2), DataValue::UnsignedNumber(42)),
-            (DataTypeId(3), DataValue::UnsignedNumber(u128::MAX)),
+            (DataType::Version.to_id(&tm), DataValue::UnsignedNumber(0)),
+            (DataType::Id.to_id(&tm), DataValue::UnsignedNumber(42)),
+            (DataType::ClientNonce.to_id(&tm), DataValue::UnsignedNumber(u128::MAX)),
         ]);
     }
 
     #[test]
     fn test_float_in_container() {
+        let tm = TypeMap::latest();
         container_roundtrip(vec![
-            (DataTypeId(1), DataValue::Float(0, 0)),
-            (DataTypeId(2), DataValue::Float(2, 12345)),
-            (DataTypeId(3), DataValue::Float(255, 4294967295)),
+            (DataType::Version.to_id(&tm), DataValue::Float(0, 0)),
+            (DataType::Id.to_id(&tm), DataValue::Float(2, 12345)),
+            (DataType::ClientNonce.to_id(&tm), DataValue::Float(255, 4294967295)),
         ]);
     }
 
     #[test]
     fn test_str_in_container() {
+        let tm = TypeMap::latest();
         container_roundtrip(vec![
-            (DataTypeId(1), DataValue::Str(String::new())),
-            (DataTypeId(2), DataValue::Str("hello".to_string())),
-            (DataTypeId(3), DataValue::Str("a".repeat(1000))),
+            (DataType::Version.to_id(&tm), DataValue::Str(String::new())),
+            (DataType::Id.to_id(&tm), DataValue::Str("hello".to_string())),
+            (DataType::ClientNonce.to_id(&tm), DataValue::Str("a".repeat(1000))),
         ]);
     }
 
     #[test]
     fn test_bytes_in_container() {
+        let tm = TypeMap::latest();
         container_roundtrip(vec![
-            (DataTypeId(1), DataValue::Bytes(vec![])),
-            (DataTypeId(2), DataValue::Bytes(vec![0x00, 0xFF, 0xAB])),
-            (DataTypeId(3), DataValue::Bytes(vec![0x42; 100])),
+            (DataType::Version.to_id(&tm), DataValue::Bytes(vec![])),
+            (DataType::Id.to_id(&tm), DataValue::Bytes(vec![0x00, 0xFF, 0xAB])),
+            (DataType::ClientNonce.to_id(&tm), DataValue::Bytes(vec![0x42; 100])),
         ]);
     }
 
     #[test]
     fn test_null_in_container() {
-        container_roundtrip(vec![(DataTypeId(1), DataValue::Null)]);
+        let tm = TypeMap::latest();
+        container_roundtrip(vec![(DataType::Version.to_id(&tm), DataValue::Null)]);
     }
 
     #[test]
@@ -1082,24 +1080,26 @@ mod tests {
 
     #[test]
     fn test_container_mixed_roundtrip() {
+        let tm = TypeMap::latest();
         container_roundtrip(vec![
-            (DataTypeId(1), DataValue::BoolTrue),
-            (DataTypeId(2), DataValue::SignedNumber(-100)),
-            (DataTypeId(3), DataValue::Str("test".to_string())),
-            (DataTypeId(4), DataValue::UnsignedNumber(u128::MAX)),
-            (DataTypeId(5), DataValue::Null),
+            (DataType::Version.to_id(&tm), DataValue::BoolTrue),
+            (DataType::Id.to_id(&tm), DataValue::SignedNumber(-100)),
+            (DataType::ClientNonce.to_id(&tm), DataValue::Str("test".to_string())),
+            (DataType::ServerNonce.to_id(&tm), DataValue::UnsignedNumber(u128::MAX)),
+            (DataType::PublicKeys.to_id(&tm), DataValue::Null),
         ]);
     }
 
     #[test]
     fn test_container_nested_roundtrip() {
+        let tm = TypeMap::latest();
         container_roundtrip(vec![
             (
-                DataTypeId(1),
-                DataValue::Container(vec![(DataTypeId(10), DataValue::BoolTrue)]),
+                DataType::Version.to_id(&tm),
+                DataValue::Container(vec![(DataType::Error.to_id(&tm), DataValue::BoolTrue)]),
             ),
             (
-                DataTypeId(2),
+                DataType::Id.to_id(&tm),
                 DataValue::Array(vec![DataValue::SignedNumber(1), DataValue::SignedNumber(2)]),
             ),
         ]);
@@ -1107,11 +1107,12 @@ mod tests {
 
     #[test]
     fn test_container_base64_roundtrip() {
+        let tm = TypeMap::latest();
         let dv = DataValue::Container(vec![(
-            DataTypeId(7),
+            DataType::Description.to_id(&tm),
             DataValue::Bytes(vec![0xDE, 0xAD, 0xBE, 0xEF]),
         )]);
-        let b64 = dv.to_base64();
+        let b64 = dv.to_base64().expect("encode failed");
         let decoded = DataValue::from_base64(&b64).expect("base64 roundtrip failed");
         assert_eq!(dv, decoded);
     }
@@ -1138,28 +1139,29 @@ mod tests {
 
     #[test]
     fn test_as_accessors() {
+        let tm = TypeMap::latest();
         let dv = DataValue::Container(vec![
-            (DataTypeId(1), DataValue::Str("alice".to_string())),
-            (DataTypeId(2), DataValue::SignedNumber(42)),
-            (DataTypeId(3), DataValue::Bytes(vec![0x01, 0x02])),
-            (DataTypeId(4), DataValue::Array(vec![DataValue::BoolTrue])),
+            (DataType::Version.to_id(&tm), DataValue::Str("alice".to_string())),
+            (DataType::Id.to_id(&tm), DataValue::SignedNumber(42)),
+            (DataType::ClientNonce.to_id(&tm), DataValue::Bytes(vec![0x01, 0x02])),
+            (DataType::ServerNonce.to_id(&tm), DataValue::Array(vec![DataValue::BoolTrue])),
         ]);
 
         let map = dv.as_map().expect("should be a container");
         assert_eq!(
-            map.get(&DataTypeId(1)).and_then(|v| v.as_str()),
+            map.get(&DataType::Version.to_id(&tm)).and_then(|v| v.as_str()),
             Some("alice")
         );
         assert_eq!(
-            map.get(&DataTypeId(2)).and_then(|v| v.as_signed_number()),
+            map.get(&DataType::Id.to_id(&tm)).and_then(|v| v.as_signed_number()),
             Some(42)
         );
         assert_eq!(
-            map.get(&DataTypeId(3)).and_then(|v| v.as_bytes()),
+            map.get(&DataType::ClientNonce.to_id(&tm)).and_then(|v| v.as_bytes()),
             Some(vec![0x01, 0x02])
         );
         assert_eq!(
-            map.get(&DataTypeId(4)).and_then(|v| v.as_array()),
+            map.get(&DataType::ServerNonce.to_id(&tm)).and_then(|v| v.as_array()),
             Some(vec![DataValue::BoolTrue])
         );
     }
@@ -1180,9 +1182,10 @@ mod tests {
 
     #[test]
     fn test_container_from_map() {
+        let tm = TypeMap::latest();
         let mut map = BTreeMap::new();
-        map.insert(DataTypeId(1), DataValue::BoolTrue);
-        map.insert(DataTypeId(2), DataValue::SignedNumber(99));
+        map.insert(DataType::Version.to_id(&tm), DataValue::BoolTrue);
+        map.insert(DataType::Id.to_id(&tm), DataValue::SignedNumber(99));
         let dv = DataValue::container_from_map(&map);
         let container = dv.as_container().expect("should be container");
         assert_eq!(container.len(), 2);
@@ -1202,13 +1205,30 @@ mod tests {
 
     #[test]
     fn test_truncated_container_rejected() {
-        let dv = DataValue::Container(vec![(DataTypeId(1), DataValue::Str("hello".to_string()))]);
-        let bytes = dv.to_bytes();
+        let tm = TypeMap::latest();
+        let dv = DataValue::Container(vec![(DataType::Version.to_id(&tm), DataValue::Str("hello".to_string()))]);
+        let bytes = dv.to_bytes().expect("encode failed");
         // Truncate to fewer than 2 bytes so neither container nor array can be read
         assert!(DataValue::from_bytes(&bytes[..1]).is_none());
         assert!(DataValue::from_bytes(&bytes[..0]).is_none());
     }
 
+    #[test]
+    fn test_oversized_count_does_not_overallocate() {
+        // A frame declaring 65535 entries but carrying almost no payload must be
+        // rejected without pre-reserving a Vec for 65535 entries. The capacity is
+        // capped against remaining bytes, so these decode attempts allocate at
+        // most a handful of slots before failing.
+        // Container path: count = 0xFFFF, no entries follow.
+        assert!(DataValue::from_bytes(&[0xFF, 0xFF]).is_none());
+        // Container path with one stray byte after the count.
+        assert!(DataValue::from_bytes(&[0xFF, 0xFF, 0x01]).is_none());
+        // Array path: force the container parse to fail first, then the array
+        // parse also sees the oversized count. A leading kind byte that is not a
+        // valid container entry makes try_read_container bail to the array path.
+        assert!(DataValue::from_bytes(&[0xFF, 0xFF, 0x08, 0xFF, 0xFF]).is_none());
+    }
+
     #[test]
     fn test_display_basic() {
         assert_eq!(format!("{}", DataValue::BoolTrue), "true");
@@ -1242,9 +1262,10 @@ mod tests {
 
     #[test]
     fn test_container_display() {
+        let tm = TypeMap::latest();
         let dv = DataValue::Container(vec![
-            (DataTypeId(3), DataValue::Str("v2.0".to_string())),
-            (DataTypeId(6), DataValue::UnsignedNumber(42)),
+            (DataType::ServerNonce.to_id(&tm), DataValue::Str("v2.0".to_string())),
+            (DataType::PqSignature.to_id(&tm), DataValue::UnsignedNumber(42)),
         ]);
         let s = format!("{}", dv);
         assert!(s.contains("3:"));
@@ -1263,19 +1284,23 @@ mod tests {
     #[cfg(feature = "crypto")]
     #[test]
     fn test_encrypt_decrypt_container_roundtrip() {
-        use mtp_crypto::ChaCha20Poly1305;
-        let key = [0xAB; 32];
-        let cipher = ChaCha20Poly1305::new(key);
+        use mtp_crypto::{EncryptionType, Keyring};
+        let tm = TypeMap::latest();
+        let keyring = Keyring::generate();
+        let bundle = keyring.public_key_bundle();
 
         let mut dv = DataValue::Container(vec![
-            (DataTypeId(1), DataValue::Str("secret".to_string())),
-            (DataTypeId(2), DataValue::UnsignedNumber(42)),
+            (DataType::Version.to_id(&tm), DataValue::Str("secret".to_string())),
+            (DataType::Id.to_id(&tm), DataValue::UnsignedNumber(42)),
         ]);
 
-        assert!(dv.encrypt_container(&cipher, b"aad").is_some());
+        assert!(
+            dv.encrypt_container(EncryptionType::MlKemChaCha20Poly1305, &bundle, b"aad")
+                .is_some()
+        );
         assert!(matches!(dv, DataValue::EncryptedContainer(_)));
 
-        assert!(dv.decrypt_into_container(&cipher, b"aad").is_some());
+        assert!(dv.decrypt_into_container(&keyring, b"aad").is_some());
         assert!(matches!(dv, DataValue::Container(_)));
 
         let entries = dv.as_container().unwrap();
@@ -1285,60 +1310,91 @@ mod tests {
     #[cfg(feature = "crypto")]
     #[test]
     fn test_encrypt_container_wrong_key_fails() {
-        use mtp_crypto::ChaCha20Poly1305;
-        let cipher_a = ChaCha20Poly1305::new([0xAB; 32]);
-        let cipher_b = ChaCha20Poly1305::new([0xCD; 32]);
+        use mtp_crypto::{EncryptionType, Keyring};
+        let tm = TypeMap::latest();
+        let keyring_a = Keyring::generate();
+        let keyring_b = Keyring::generate();
 
         let mut dv =
-            DataValue::Container(vec![(DataTypeId(1), DataValue::Str("secret".to_string()))]);
+            DataValue::Container(vec![(DataType::Version.to_id(&tm), DataValue::Str("secret".to_string()))]);
 
-        assert!(dv.encrypt_container(&cipher_a, b"aad").is_some());
-        assert!(dv.decrypt_into_container(&cipher_b, b"aad").is_none());
+        assert!(
+            dv.encrypt_container(
+                EncryptionType::MlKemChaCha20Poly1305,
+                &keyring_a.public_key_bundle(),
+                b"aad"
+            )
+            .is_some()
+        );
+        assert!(dv.decrypt_into_container(&keyring_b, b"aad").is_none());
     }
 
     #[cfg(feature = "crypto")]
     #[test]
     fn test_encrypt_container_wrong_aad_fails() {
-        use mtp_crypto::ChaCha20Poly1305;
-        let cipher = ChaCha20Poly1305::new([0xAB; 32]);
+        use mtp_crypto::{EncryptionType, Keyring};
+        let tm = TypeMap::latest();
+        let keyring = Keyring::generate();
 
         let mut dv =
-            DataValue::Container(vec![(DataTypeId(1), DataValue::Str("secret".to_string()))]);
+            DataValue::Container(vec![(DataType::Version.to_id(&tm), DataValue::Str("secret".to_string()))]);
 
-        assert!(dv.encrypt_container(&cipher, b"correct-aad").is_some());
-        assert!(dv.decrypt_into_container(&cipher, b"wrong-aad").is_none());
+        assert!(
+            dv.encrypt_container(
+                EncryptionType::MlKemChaCha20Poly1305,
+                &keyring.public_key_bundle(),
+                b"correct-aad"
+            )
+            .is_some()
+        );
+        assert!(dv.decrypt_into_container(&keyring, b"wrong-aad").is_none());
     }
 
     #[cfg(feature = "crypto")]
     #[test]
     fn test_encrypt_non_container_fails() {
-        let cipher = mtp_crypto::ChaCha20Poly1305::new([0xAB; 32]);
+        use mtp_crypto::{EncryptionType, Keyring};
+        let keyring = Keyring::generate();
 
         let mut dv = DataValue::Str("not a container".to_string());
-        assert!(dv.encrypt_container(&cipher, b"aad").is_none());
+        assert!(
+            dv.encrypt_container(
+                EncryptionType::MlKemChaCha20Poly1305,
+                &keyring.public_key_bundle(),
+                b"aad"
+            )
+            .is_none()
+        );
     }
 
     #[cfg(feature = "crypto")]
     #[test]
     fn test_sign_verify_container_roundtrip() {
-        use mtp_crypto::{ChaCha20Poly1305, Ed25519Signer, SigAlgorithm};
+        use mtp_crypto::{Ed25519Signer, EncryptionType, Keyring, SigAlgorithm};
+        let tm = TypeMap::latest();
 
+        let keyring = Keyring::generate();
         let (signer, sk, _pk) = Ed25519Signer::generate();
-        let cipher = ChaCha20Poly1305::new([0xAB; 32]);
 
         let mut dv = DataValue::Container(vec![(
-            DataTypeId(1),
+            DataType::Version.to_id(&tm),
             DataValue::Str("signed data".to_string()),
         )]);
 
         assert!(
-            dv.sign_and_encrypt_container(SigAlgorithm::ED25519, &signer, &cipher, b"aad")
+            dv.sign_and_encrypt_container(
+                SigAlgorithm::ED25519,
+                &signer,
+                EncryptionType::MlKemChaCha20Poly1305,
+                &keyring.public_key_bundle(),
+                b"aad"
+            )
             .is_some()
         );
         assert!(matches!(dv, DataValue::SignedEncryptedContainer(_)));
 
         assert!(
-            dv.decrypt_signed_encrypted_container(&cipher, b"aad")
+            dv.decrypt_signed_encrypted_container(&keyring, b"aad")
                 .is_some()
         );
         assert!(matches!(dv, DataValue::SignedContainer(_)));
@@ -1355,13 +1411,14 @@ mod tests {
     #[test]
     fn test_sign_container_wrong_key_fails() {
         use mtp_crypto::{Ed25519Signer, SigAlgorithm};
+        let tm = TypeMap::latest();
 
         let (signer, _, _) = Ed25519Signer::generate();
         let (_, sk2, _) = Ed25519Signer::generate();
         let wrong_verifier = Ed25519Signer::new(&sk2).unwrap();
 
         let mut dv = DataValue::Container(vec![(
-            DataTypeId(1),
+            DataType::Version.to_id(&tm),
             DataValue::Str("signed data".to_string()),
         )]);
 

codec/src/lib.rs

@@ -6,8 +6,8 @@ pub use data_value::{DataKind, DataValue};
 pub use mtp_common::CodecError;
 
 pub use mtp_type_map::{
-    communication_type_name, data_type_name, CommunicationType, CommunicationTypeId, DataType,
-    DataTypeId, TypeMap, Version, PROTOCOL_VERSION,
+    CommunicationType, CommunicationTypeId, DataType, DataTypeId, PROTOCOL_VERSION, TypeMap,
+    Version, communication_type_name, data_type_name,
 };
 
 pub(crate) fn rand_u32() -> u32 {

common/src/lib.rs

@@ -12,6 +12,8 @@ pub enum CodecError {
     ReservedCommunicationType(u16),
     #[error("Invalid encoding")]
     InvalidEncoding,
+    #[error("Too many entries to encode")]
+    TooManyEntries,
     #[error("Crypto failed: {0}")]
     CryptoFailed(String),
 }
@@ -40,13 +42,11 @@ mod tests {
     }
 }
 
-// ===========================================================================
-// CommunicationError
-//
-// On native targets the full variant set (including quinn / wtransport
-// wrappers) is available.  On WASM only the transport-independent subset is
-// compiled.
-// ===========================================================================
+/* CommunicationError
+ *
+ * On native targets the full variant set (including quinn / wtransport
+ * wrappers) is available.  On WASM only the transport-independent subset is
+ * compiled. */
 
 #[cfg(not(target_arch = "wasm32"))]
 #[derive(Debug, Error, Clone)]
@@ -69,6 +69,9 @@ pub enum CommunicationError {
     #[error("ParseCommunicationValue error")]
     ParseCommunicationValue,
 
+    #[error("Encode error")]
+    Encode,
+
     #[error("Parse Certificate error")]
     CertificateParseFailed,
 
@@ -139,6 +142,9 @@ pub enum CommunicationError {
     #[error("ParseCommunicationValue error")]
     ParseCommunicationValue,
 
+    #[error("Encode error")]
+    Encode,
+
     #[error("Parse Certificate error")]
     CertificateParseFailed,
 
@@ -182,6 +188,7 @@ impl PartialEq for CommunicationError {
             (Self::ConnectionLost, Self::ConnectionLost) => true,
             (Self::Quinn(_), Self::Quinn(_)) => true,
             (Self::ParseCommunicationValue, Self::ParseCommunicationValue) => true,
+            (Self::Encode, Self::Encode) => true,
             (Self::CertificateParseFailed, Self::CertificateParseFailed) => true,
             (Self::CertificateLoadFailed, Self::CertificateLoadFailed) => true,
             (Self::ParseError(a), Self::ParseError(b)) => a == b,
@@ -213,6 +220,7 @@ impl PartialEq for CommunicationError {
             (Self::ClosedByPeer, Self::ClosedByPeer) => true,
             (Self::ConnectionLost, Self::ConnectionLost) => true,
             (Self::ParseCommunicationValue, Self::ParseCommunicationValue) => true,
+            (Self::Encode, Self::Encode) => true,
             (Self::CertificateParseFailed, Self::CertificateParseFailed) => true,
             (Self::CertificateLoadFailed, Self::CertificateLoadFailed) => true,
             (Self::ParseError(a), Self::ParseError(b)) => a == b,
@@ -241,9 +249,18 @@ mod communication_error_tests {
 
     #[test]
     fn test_communication_error_display() {
-        assert_eq!(format!("{}", CommunicationError::UseAfterClosed), "Use after Closed");
-        assert_eq!(format!("{}", CommunicationError::StreamClosed), "Stream Closed");
-        assert_eq!(format!("{}", CommunicationError::StreamError), "Stream Error");
+        assert_eq!(
+            format!("{}", CommunicationError::UseAfterClosed),
+            "Use after Closed"
+        );
+        assert_eq!(
+            format!("{}", CommunicationError::StreamClosed),
+            "Stream Closed"
+        );
+        assert_eq!(
+            format!("{}", CommunicationError::StreamError),
+            "Stream Error"
+        );
     }
 
     #[test]

crypto/README.md

@@ -24,7 +24,7 @@ let pt = cipher.decrypt(&ct, b"aad")?;
 
 ## Signatures
 
-### Ed25519 (classical)
+### Ed25519
 
 ```rust
 use mtp_crypto::{Ed25519Signer, SignatureScheme};
@@ -34,7 +34,7 @@ let sig = signer.sign(b"message")?;
 signer.verify(b"message", &sig)?;
 ```
 
-### ML-DSA-65 (post-quantum, requires `pqc`)
+### ML-DSA-65
 
 ```rust
 use mtp_crypto::{MlDsaSigner, SignatureScheme};
@@ -47,7 +47,7 @@ signer.verify(b"message", &sig)?;
 let signer = MlDsaSigner::new(&sk, &pk)?;
 ```
 
-### Dual signatures (requires `pqc`)
+### Dual signatures
 
 ```rust
 use mtp_crypto::{sign_dual, DualSignature, Ed25519Signer, MlDsaSigner};
@@ -58,7 +58,7 @@ let dual = sign_dual(ed_signer.signing_key(), ml_signer.signing_key(), b"msg");
 dual.verify(ed_signer.verifying_key(), ml_signer.verifying_key(), b"msg")?;
 ```
 
-## Hybrid KEM (requires `pqc`)
+## Hybrid KEM
 
 X25519 + ML-KEM-768. 64-byte shared secret. Feed into HKDF before use.
 

crypto/src/aead.rs

@@ -33,9 +33,9 @@ impl ChaCha20Poly1305 {
 #[cfg(feature = "chacha20poly1305")]
 impl AeadEncrypt for ChaCha20Poly1305 {
     fn encrypt(&self, plaintext: &[u8], aad: &[u8]) -> Result<Vec<u8>, CryptoError> {
-        use chacha20poly1305::aead::{Aead, KeyInit, Payload};
         use chacha20poly1305::XChaCha20Poly1305;
         use chacha20poly1305::XNonce;
+        use chacha20poly1305::aead::{Aead, KeyInit, Payload};
 
         let key = chacha20poly1305::Key::from_slice(&self.key);
         let cipher = XChaCha20Poly1305::new(key);
@@ -63,9 +63,9 @@ impl AeadEncrypt for ChaCha20Poly1305 {
 #[cfg(feature = "chacha20poly1305")]
 impl AeadDecrypt for ChaCha20Poly1305 {
     fn decrypt(&self, ciphertext: &[u8], aad: &[u8]) -> Result<Vec<u8>, CryptoError> {
-        use chacha20poly1305::aead::{Aead, KeyInit, Payload};
         use chacha20poly1305::XChaCha20Poly1305;
         use chacha20poly1305::XNonce;
+        use chacha20poly1305::aead::{Aead, KeyInit, Payload};
 
         if ciphertext.len() < 24 {
             return Err(CryptoError::InvalidNonceLength);
@@ -76,10 +76,7 @@ impl AeadDecrypt for ChaCha20Poly1305 {
         let cipher = XChaCha20Poly1305::new(key);
         let nonce_ref = XNonce::from_slice(nonce);
 
-        let payload = Payload {
-            msg: ct,
-            aad,
-        };
+        let payload = Payload { msg: ct, aad };
 
         cipher
             .decrypt(nonce_ref, payload)
@@ -109,9 +106,9 @@ impl Aes256Gcm {
 #[cfg(feature = "aes-gcm")]
 impl AeadEncrypt for Aes256Gcm {
     fn encrypt(&self, plaintext: &[u8], aad: &[u8]) -> Result<Vec<u8>, CryptoError> {
-        use aes_gcm::aead::{Aead, KeyInit, Payload};
         use aes_gcm::Aes256Gcm as AesGcmInner;
         use aes_gcm::Nonce;
+        use aes_gcm::aead::{Aead, KeyInit, Payload};
 
         let key = aes_gcm::Key::<AesGcmInner>::from_slice(&self.key);
         let cipher = AesGcmInner::new(key);
@@ -139,9 +136,9 @@ impl AeadEncrypt for Aes256Gcm {
 #[cfg(feature = "aes-gcm")]
 impl AeadDecrypt for Aes256Gcm {
     fn decrypt(&self, ciphertext: &[u8], aad: &[u8]) -> Result<Vec<u8>, CryptoError> {
-        use aes_gcm::aead::{Aead, KeyInit, Payload};
         use aes_gcm::Aes256Gcm as AesGcmInner;
         use aes_gcm::Nonce;
+        use aes_gcm::aead::{Aead, KeyInit, Payload};
 
         if ciphertext.len() < 12 {
             return Err(CryptoError::InvalidNonceLength);
@@ -152,10 +149,7 @@ impl AeadDecrypt for Aes256Gcm {
         let cipher = AesGcmInner::new(key);
         let nonce_ref = Nonce::from_slice(nonce);
 
-        let payload = Payload {
-            msg: ct,
-            aad,
-        };
+        let payload = Payload { msg: ct, aad };
 
         cipher
             .decrypt(nonce_ref, payload)

crypto/src/auth.rs

@@ -0,0 +1,128 @@
+/*
+ * Canonical signed payloads for the authentication handshake.
+ *
+ * Each payload starts with a unique [`domain`] tag so a signature for one step
+ * cannot be replayed as another.
+ *
+ * Handshake:
+ * Step 1. Client -> Host : Identification { version, id }                 (unsigned hello)
+ * Step 2. Host -> Client : Challenge      { server_challenge, host_sig }  host_sig over challenge_payload
+ * Step 3. Client -> Host : ChallengeResponse { client_nonce, sig }        sig over login_proof_payload
+ * Step 4. Host -> Client : IdentificationResponse { connected, id, host_sig } host_sig over host_final_payload
+ */
+
+/// Domain-separation tags — a distinct leading byte per signed context.
+pub mod domain {
+    /// Host's signature over the challenge it issues (step 2).
+    pub const CHALLENGE: u8 = 0x10;
+    /// Client's authenticating proof for a login (step 3).
+    pub const LOGIN_PROOF: u8 = 0x11;
+    /// Client's authenticating proof for a registration (step 3).
+    pub const REGISTER_PROOF: u8 = 0x12;
+    /// Host's final confirmation signature (step 4).
+    pub const HOST_FINAL: u8 = 0x13;
+}
+
+/*
+ * Host's challenge (step 2): binds `id` and `server_challenge` to prove host
+ * key possession before the client reveals its proof.
+ */
+pub fn challenge_payload(id: u64, server_challenge: u128) -> Vec<u8> {
+    let mut p = Vec::with_capacity(1 + 8 + 16);
+    p.push(domain::CHALLENGE);
+    p.extend_from_slice(&id.to_be_bytes());
+    p.extend_from_slice(&server_challenge.to_be_bytes());
+    p
+}
+
+/*
+ * Client's login proof (step 3): binds version, id, server_challenge, client_nonce.
+ */
+pub fn login_proof_payload(
+    version: &str,
+    id: u64,
+    server_challenge: u128,
+    client_nonce: u128,
+) -> Vec<u8> {
+    let mut p = Vec::with_capacity(1 + version.len() + 8 + 16 + 16);
+    p.push(domain::LOGIN_PROOF);
+    p.extend_from_slice(version.as_bytes());
+    p.extend_from_slice(&id.to_be_bytes());
+    p.extend_from_slice(&server_challenge.to_be_bytes());
+    p.extend_from_slice(&client_nonce.to_be_bytes());
+    p
+}
+
+/*
+ * Client's registration proof (step 3): binds version, public_keys, server_challenge, client_nonce.
+ */
+pub fn register_proof_payload(
+    version: &str,
+    public_keys: &[u8],
+    server_challenge: u128,
+    client_nonce: u128,
+) -> Vec<u8> {
+    let mut p = Vec::with_capacity(1 + version.len() + 16 + 16 + public_keys.len());
+    p.push(domain::REGISTER_PROOF);
+    p.extend_from_slice(version.as_bytes());
+    p.extend_from_slice(&server_challenge.to_be_bytes());
+    p.extend_from_slice(&client_nonce.to_be_bytes());
+    p.extend_from_slice(public_keys);
+    p
+}
+
+/*
+ * Host's final confirmation (step 4): binds id, client_nonce, server_challenge
+ * to prove host liveness over a value the client chose.
+ */
+pub fn host_final_payload(id: u64, client_nonce: u128, server_challenge: u128) -> Vec<u8> {
+    let mut p = Vec::with_capacity(1 + 8 + 16 + 16);
+    p.push(domain::HOST_FINAL);
+    p.extend_from_slice(&id.to_be_bytes());
+    p.extend_from_slice(&client_nonce.to_be_bytes());
+    p.extend_from_slice(&server_challenge.to_be_bytes());
+    p
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn domain_tags_are_distinct() {
+        let tags = [
+            domain::CHALLENGE,
+            domain::LOGIN_PROOF,
+            domain::REGISTER_PROOF,
+            domain::HOST_FINAL,
+        ];
+        for (i, a) in tags.iter().enumerate() {
+            for b in &tags[i + 1..] {
+                assert_ne!(a, b, "domain tags must be unique");
+            }
+        }
+    }
+
+    #[test]
+    fn login_and_register_proofs_never_collide() {
+        let login = login_proof_payload("1.0", 7, 1, 2);
+        let register = register_proof_payload("1.0", &7u64.to_be_bytes(), 1, 2);
+        assert_ne!(login, register);
+        assert_ne!(login[0], register[0]);
+    }
+
+    #[test]
+    fn challenge_binds_id_and_value() {
+        assert_ne!(challenge_payload(1, 9), challenge_payload(2, 9));
+        assert_ne!(challenge_payload(1, 9), challenge_payload(1, 8));
+        assert_eq!(challenge_payload(1, 9)[0], domain::CHALLENGE);
+    }
+
+    #[test]
+    fn proofs_bind_the_server_challenge() {
+        assert_ne!(
+            login_proof_payload("1.0", 3, 100, 200),
+            login_proof_payload("1.0", 3, 101, 200),
+        );
+    }
+}

crypto/src/enc.rs

@@ -0,0 +1,246 @@
+#[cfg(all(feature = "mlkem-tls", feature = "hkdf"))]
+use crate::error::CryptoError;
+
+#[cfg(all(feature = "mlkem-tls", feature = "hkdf"))]
+use crate::kdf::derive_encryption_key;
+#[cfg(all(feature = "mlkem-tls", feature = "hkdf"))]
+use crate::kem::HybridKem;
+#[cfg(all(feature = "mlkem-tls", feature = "hkdf"))]
+use crate::keypair::{Keyring, PublicKeyBundle};
+
+/*
+ * Algorithm selector for encrypted containers.
+ *
+ * Mirrors `SigAlgorithm` for signatures: a single marking byte identifies the
+ * key-encapsulation mechanism and the AEAD used to seal a container. The byte
+ * is stored as the first byte of every encrypted blob so the decryptor can pick
+ * the matching algorithm (and the matching keypair from a `Keyring`) without
+ * any out-of-band agreement.
+ *
+ * All variants currently use ML-KEM (X25519MlKem768) for key encapsulation and
+ * differ only in the AEAD. AES-256-GCM variants require the `aes-gcm` feature
+ * (enabled via the crate's `full` feature); sealing/opening with a variant whose
+ * AEAD feature is not compiled in returns an error.
+ */
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum EncryptionType {
+    /// ML-KEM (X25519MlKem768) key encapsulation with XChaCha20-Poly1305 AEAD.
+    MlKemChaCha20Poly1305,
+    /// ML-KEM (X25519MlKem768) key encapsulation with AES-256-GCM AEAD.
+    MlKemAes256Gcm,
+}
+
+impl EncryptionType {
+    pub const ML_KEM_CHACHA20POLY1305: u8 = 0x01;
+    pub const ML_KEM_AES256_GCM: u8 = 0x02;
+
+    /// The marking byte written at the front of an encrypted blob.
+    pub const fn to_byte(self) -> u8 {
+        match self {
+            Self::MlKemChaCha20Poly1305 => Self::ML_KEM_CHACHA20POLY1305,
+            Self::MlKemAes256Gcm => Self::ML_KEM_AES256_GCM,
+        }
+    }
+
+    /// Recover an `EncryptionType` from its marking byte, or `None` if unknown.
+    pub const fn from_byte(b: u8) -> Option<Self> {
+        match b {
+            Self::ML_KEM_CHACHA20POLY1305 => Some(Self::MlKemChaCha20Poly1305),
+            Self::ML_KEM_AES256_GCM => Some(Self::MlKemAes256Gcm),
+            _ => None,
+        }
+    }
+}
+
+/*
+ * Seal `plaintext` with a 32-byte AEAD key chosen by `enc_type`.
+ *
+ * Returns `EncryptionFailed` when the selected AEAD's feature is not compiled in.
+ */
+#[cfg(all(feature = "mlkem-tls", feature = "hkdf"))]
+#[allow(unused_variables)]
+fn aead_seal(
+    enc_type: EncryptionType,
+    key: [u8; 32],
+    plaintext: &[u8],
+    aad: &[u8],
+) -> Result<Vec<u8>, CryptoError> {
+    #[allow(unused_imports)]
+    use crate::aead::AeadEncrypt;
+    match enc_type {
+        #[cfg(feature = "chacha20poly1305")]
+        EncryptionType::MlKemChaCha20Poly1305 => {
+            crate::aead::ChaCha20Poly1305::new(key).encrypt(plaintext, aad)
+        }
+        #[cfg(feature = "aes-gcm")]
+        EncryptionType::MlKemAes256Gcm => crate::aead::Aes256Gcm::new(key).encrypt(plaintext, aad),
+        #[allow(unreachable_patterns)]
+        _ => Err(CryptoError::EncryptionFailed),
+    }
+}
+
+/*
+ * Open `ciphertext` with a 32-byte AEAD key chosen by `enc_type`.
+ *
+ * Returns `DecryptionFailed` when the selected AEAD's feature is not compiled in.
+ */
+#[cfg(all(feature = "mlkem-tls", feature = "hkdf"))]
+#[allow(unused_variables)]
+fn aead_open(
+    enc_type: EncryptionType,
+    key: [u8; 32],
+    ciphertext: &[u8],
+    aad: &[u8],
+) -> Result<Vec<u8>, CryptoError> {
+    #[allow(unused_imports)]
+    use crate::aead::AeadDecrypt;
+    match enc_type {
+        #[cfg(feature = "chacha20poly1305")]
+        EncryptionType::MlKemChaCha20Poly1305 => {
+            crate::aead::ChaCha20Poly1305::new(key).decrypt(ciphertext, aad)
+        }
+        #[cfg(feature = "aes-gcm")]
+        EncryptionType::MlKemAes256Gcm => crate::aead::Aes256Gcm::new(key).decrypt(ciphertext, aad),
+        #[allow(unreachable_patterns)]
+        _ => Err(CryptoError::DecryptionFailed),
+    }
+}
+
+#[cfg(all(feature = "mlkem-tls", feature = "hkdf"))]
+const ENC_KDF_SALT: &[u8] = b"mtp-container-enc";
+#[cfg(all(feature = "mlkem-tls", feature = "hkdf"))]
+const ENC_KDF_CONTEXT: &[u8] = b"single-recipient";
+
+/*
+ * Encrypt `plaintext` for a single recipient, selecting the algorithm with
+ * `enc_type` and the recipient's KEM public key from `recipient`.
+ *
+ * The returned, self-describing blob is laid out as:
+ *   [1 byte EncryptionType] [2 bytes u16 kem_ct_len] [kem_ciphertext] [aead_payload]
+ * where `aead_payload` is the AEAD output (nonce + ciphertext + tag). The AEAD
+ * key is derived from the KEM shared secret via HKDF, so no separate content key
+ * is transmitted.
+ *
+ * Requires the `mlkem-tls` and `hkdf` features, plus the AEAD feature backing
+ * `enc_type`.
+ */
+#[cfg(all(feature = "mlkem-tls", feature = "hkdf"))]
+pub fn encrypt_for(
+    enc_type: EncryptionType,
+    recipient: &PublicKeyBundle,
+    plaintext: &[u8],
+    aad: &[u8],
+) -> Result<Vec<u8>, CryptoError> {
+    let enc = HybridKem::encapsulate(&recipient.kem_public_key)?;
+    let key = derive_encryption_key(&enc.shared_secret, ENC_KDF_SALT, ENC_KDF_CONTEXT)?;
+    let aead_payload = aead_seal(enc_type, key, plaintext, aad)?;
+
+    let kem_ct = enc.ciphertext;
+    let mut out = Vec::with_capacity(1 + 2 + kem_ct.len() + aead_payload.len());
+    out.push(enc_type.to_byte());
+    out.extend_from_slice(&(kem_ct.len() as u16).to_be_bytes());
+    out.extend_from_slice(&kem_ct);
+    out.extend_from_slice(&aead_payload);
+    Ok(out)
+}
+
+/*
+ * Decrypt a blob produced by [`encrypt_for`] using `keyring`.
+ *
+ * The leading byte selects the `EncryptionType` (and thus which keypair to use
+ * from the keyring); for the current ML-KEM variants that is `kem_secret_key`.
+ * Returns `DecryptionFailed` on any malformed input or authentication failure.
+ *
+ * Requires the `mlkem-tls` and `hkdf` features, plus the AEAD feature backing
+ * the blob's algorithm.
+ */
+#[cfg(all(feature = "mlkem-tls", feature = "hkdf"))]
+pub fn decrypt_with(blob: &[u8], keyring: &Keyring, aad: &[u8]) -> Result<Vec<u8>, CryptoError> {
+    if blob.len() < 3 {
+        return Err(CryptoError::DecryptionFailed);
+    }
+    let enc_type = EncryptionType::from_byte(blob[0]).ok_or(CryptoError::DecryptionFailed)?;
+    let kem_ct_len = u16::from_be_bytes([blob[1], blob[2]]) as usize;
+    let kem_end = 3usize
+        .checked_add(kem_ct_len)
+        .ok_or(CryptoError::DecryptionFailed)?;
+    let kem_ct = blob.get(3..kem_end).ok_or(CryptoError::DecryptionFailed)?;
+    let aead_payload = blob.get(kem_end..).ok_or(CryptoError::DecryptionFailed)?;
+
+    let shared_secret = HybridKem::decapsulate(&keyring.kem_secret_key, kem_ct)?;
+    let key = derive_encryption_key(&shared_secret, ENC_KDF_SALT, ENC_KDF_CONTEXT)?;
+    aead_open(enc_type, key, aead_payload, aad)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn encryption_type_byte_roundtrip() {
+        for t in [
+            EncryptionType::MlKemChaCha20Poly1305,
+            EncryptionType::MlKemAes256Gcm,
+        ] {
+            assert_eq!(EncryptionType::from_byte(t.to_byte()), Some(t));
+        }
+        assert_eq!(EncryptionType::from_byte(0x00), None);
+        assert_eq!(EncryptionType::from_byte(0xFF), None);
+    }
+
+    #[cfg(all(feature = "mlkem-tls", feature = "hkdf", feature = "chacha20poly1305"))]
+    #[test]
+    fn encrypt_for_roundtrip() {
+        let kr = Keyring::generate();
+        let blob = encrypt_for(
+            EncryptionType::MlKemChaCha20Poly1305,
+            &kr.public_key_bundle(),
+            b"secret payload",
+            b"aad",
+        )
+        .unwrap();
+        assert_eq!(blob[0], EncryptionType::ML_KEM_CHACHA20POLY1305);
+
+        let pt = decrypt_with(&blob, &kr, b"aad").unwrap();
+        assert_eq!(pt, b"secret payload");
+    }
+
+    #[cfg(all(feature = "mlkem-tls", feature = "hkdf", feature = "chacha20poly1305"))]
+    #[test]
+    fn decrypt_with_wrong_keyring_fails() {
+        let kr = Keyring::generate();
+        let other = Keyring::generate();
+        let blob = encrypt_for(
+            EncryptionType::MlKemChaCha20Poly1305,
+            &kr.public_key_bundle(),
+            b"secret",
+            b"aad",
+        )
+        .unwrap();
+        assert!(decrypt_with(&blob, &other, b"aad").is_err());
+    }
+
+    #[cfg(all(feature = "mlkem-tls", feature = "hkdf", feature = "chacha20poly1305"))]
+    #[test]
+    fn decrypt_with_wrong_aad_fails() {
+        let kr = Keyring::generate();
+        let blob = encrypt_for(
+            EncryptionType::MlKemChaCha20Poly1305,
+            &kr.public_key_bundle(),
+            b"secret",
+            b"right",
+        )
+        .unwrap();
+        assert!(decrypt_with(&blob, &kr, b"wrong").is_err());
+    }
+
+    #[cfg(all(feature = "mlkem-tls", feature = "hkdf", feature = "chacha20poly1305"))]
+    #[test]
+    fn decrypt_with_malformed_fails() {
+        let kr = Keyring::generate();
+        assert!(decrypt_with(b"", &kr, b"").is_err());
+        assert!(decrypt_with(&[0x01, 0x00], &kr, b"").is_err());
+        // Unknown algorithm byte.
+        assert!(decrypt_with(&[0x7F, 0x00, 0x00], &kr, b"").is_err());
+    }
+}

crypto/src/hash.rs

@@ -13,6 +13,12 @@ pub fn sha256_double(data: &[u8]) -> [u8; 32] {
 
 pub struct Sha256Hasher(sha2::Sha256);
 
+impl Default for Sha256Hasher {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 impl Sha256Hasher {
     pub fn new() -> Self {
         Self(sha2::Sha256::new())

crypto/src/helper.rs

@@ -196,5 +196,3 @@ pub fn decrypt_multi(
     }
     Err(CryptoError::DecryptionFailed)
 }
-
-

crypto/src/kdf.rs

@@ -16,10 +16,14 @@ pub fn hkdf_expand(
 }
 
 pub fn hkdf_extract(ikm: &[u8], salt: &[u8]) -> [u8; 32] {
-    let (_, hk) = Hkdf::<Sha256>::extract(Some(salt), ikm);
-    let mut okm = [0u8; 32];
-    hk.expand(&[], &mut okm).expect("hkdf expand failed");
-    okm
+    /*
+     * Return the pseudo-random key (PRK) produced by HKDF-Extract directly.
+     * Extract cannot fail, so this avoids the panicking expand step entirely.
+     */
+    let (prk, _) = Hkdf::<Sha256>::extract(Some(salt), ikm);
+    let mut out = [0u8; 32];
+    out.copy_from_slice(&prk);
+    out
 }
 
 pub fn derive_encryption_key(

crypto/src/kem.rs

@@ -12,8 +12,7 @@ pub struct HybridKem;
 #[cfg(feature = "mlkem-tls")]
 impl HybridKem {
     pub fn generate_keypair() -> (KemPrivateKey, KemPublicKey) {
-        let (ek, dk) =
-            mlkem_tls::X25519MlKem768::keygen(&mut rand_core::OsRng);
+        let (ek, dk) = mlkem_tls::X25519MlKem768::keygen(&mut rand_core::OsRng);
         (
             KemPrivateKey::new(dk.as_bytes().to_vec()),
             KemPublicKey::new(ek.as_bytes().to_vec()),
@@ -23,8 +22,7 @@ impl HybridKem {
     pub fn encapsulate(recipient_pk: &KemPublicKey) -> Result<Encapsulated, CryptoError> {
         let ek = mlkem_tls::EncapsKey768::try_from(recipient_pk.as_bytes())
             .map_err(|_| CryptoError::KemEncapsulationFailed)?;
-        let (ct, ss) =
-            mlkem_tls::X25519MlKem768::encapsulate(&ek, &mut rand_core::OsRng);
+        let (ct, ss) = mlkem_tls::X25519MlKem768::encapsulate(&ek, &mut rand_core::OsRng);
         Ok(Encapsulated {
             ciphertext: ct.as_bytes().to_vec(),
             shared_secret: ss.as_bytes().to_vec(),

crypto/src/keypair.rs

@@ -348,7 +348,7 @@ impl Keyring {
                     .get(*offset..*offset + 2)
                     .ok_or(CryptoError::InvalidKeyLength)?
                     .try_into()
-                    .unwrap(),
+                    .expect("slice is 2 bytes, verified above"),
             ) as usize;
             *offset += 2;
             let key = bytes

crypto/src/lib.rs

@@ -1,4 +1,5 @@
 pub mod aead;
+pub mod auth;
 pub mod error;
 pub mod keypair;
 
@@ -17,6 +18,8 @@ pub use sign::SigAlgorithm;
 #[cfg(feature = "mlkem-tls")]
 pub mod kem;
 
+pub mod enc;
+
 pub mod helper;
 
 pub use aead::{AeadCipher, AeadDecrypt, AeadEncrypt};
@@ -34,16 +37,16 @@ pub use aead::ChaCha20Poly1305;
 pub use aead::Aes256Gcm;
 
 #[cfg(feature = "ed25519-dalek")]
-pub use sign::{verify_ed25519, Ed25519Signer, SignatureScheme};
+pub use sign::{Ed25519Signer, SignatureScheme, verify_ed25519};
 
 #[cfg(feature = "ml-dsa")]
-pub use sign::{verify_ml_dsa, MlDsaSigner};
+pub use sign::{MlDsaSigner, verify_ml_dsa};
 
 #[cfg(all(feature = "ed25519-dalek", feature = "ml-dsa"))]
-pub use sign::{sign_dual, DualSignature};
+pub use sign::{DualSignature, sign_dual};
 
 #[cfg(feature = "sha2")]
-pub use hash::{sha256, sha256_double, Sha256Hasher};
+pub use hash::{Sha256Hasher, sha256, sha256_double};
 
 #[cfg(feature = "hkdf")]
 pub use kdf::{derive_encryption_key, hkdf_expand, hkdf_extract};
@@ -51,8 +54,13 @@ pub use kdf::{derive_encryption_key, hkdf_expand, hkdf_extract};
 #[cfg(feature = "mlkem-tls")]
 pub use kem::{Encapsulated, HybridKem};
 
+pub use enc::EncryptionType;
+
+#[cfg(all(feature = "mlkem-tls", feature = "hkdf"))]
+pub use enc::{decrypt_with, encrypt_for};
+
 #[cfg(all(feature = "mlkem-tls", feature = "chacha20poly1305", feature = "hkdf"))]
-pub use helper::{decrypt_multi, encrypt_multi, MultiEncryptedMessage, RecipientEntry};
+pub use helper::{MultiEncryptedMessage, RecipientEntry, decrypt_multi, encrypt_multi};
 
 /* ================================ TESTS ================================ */
 #[cfg(test)]
@@ -140,13 +148,8 @@ mod tests {
 
         let (ed_signer, _, _) = Ed25519Signer::generate();
         let (ml_signer, _, _) = MlDsaSigner::generate();
-        let dual = sign_dual(ed_signer.signing_key(), ml_signer.signing_key(), b"msg");
-        dual
-            .verify(
-                ed_signer.verifying_key(),
-                ml_signer.verifying_key(),
-                b"msg",
-            )
+        let dual = sign_dual(ed_signer.signing_key(), ml_signer.signing_key(), b"msg").unwrap();
+        dual.verify(ed_signer.verifying_key(), ml_signer.verifying_key(), b"msg")
             .unwrap();
     }
 
@@ -157,10 +160,15 @@ mod tests {
 
         let (ed_signer, _, _) = Ed25519Signer::generate();
         let (ml_signer, _, _) = MlDsaSigner::generate();
-        let dual = sign_dual(ed_signer.signing_key(), ml_signer.signing_key(), b"msg");
-        assert!(dual
-            .verify(ed_signer.verifying_key(), ml_signer.verifying_key(), b"wrong")
-            .is_err());
+        let dual = sign_dual(ed_signer.signing_key(), ml_signer.signing_key(), b"msg").unwrap();
+        assert!(
+            dual.verify(
+                ed_signer.verifying_key(),
+                ml_signer.verifying_key(),
+                b"wrong"
+            )
+            .is_err()
+        );
     }
 
     #[cfg(feature = "hkdf")]
@@ -248,9 +256,18 @@ mod tests {
         let kr = Keyring::generate();
         let bytes = kr.to_bytes();
         let loaded = Keyring::from_bytes(&bytes).unwrap();
-        assert_eq!(kr.kem_public_key.as_bytes(), loaded.kem_public_key.as_bytes());
-        assert_eq!(kr.sig_pq_public_key.as_bytes(), loaded.sig_pq_public_key.as_bytes());
-        assert_eq!(kr.sig_cl_public_key.as_bytes(), loaded.sig_cl_public_key.as_bytes());
+        assert_eq!(
+            kr.kem_public_key.as_bytes(),
+            loaded.kem_public_key.as_bytes()
+        );
+        assert_eq!(
+            kr.sig_pq_public_key.as_bytes(),
+            loaded.sig_pq_public_key.as_bytes()
+        );
+        assert_eq!(
+            kr.sig_cl_public_key.as_bytes(),
+            loaded.sig_cl_public_key.as_bytes()
+        );
     }
 
     #[cfg(all(feature = "mlkem-tls", feature = "ml-dsa", feature = "ed25519-dalek"))]
@@ -260,7 +277,10 @@ mod tests {
         let bundle = kr.public_key_bundle();
         let bytes = bundle.as_bytes();
         let loaded = PublicKeyBundle::from_bytes(&bytes).unwrap();
-        assert_eq!(bundle.kem_public_key.as_bytes(), loaded.kem_public_key.as_bytes());
+        assert_eq!(
+            bundle.kem_public_key.as_bytes(),
+            loaded.kem_public_key.as_bytes()
+        );
         assert_eq!(
             bundle.sig_pq_public_key.as_bytes(),
             loaded.sig_pq_public_key.as_bytes()
@@ -283,8 +303,8 @@ mod tests {
     #[cfg(all(feature = "mlkem-tls", feature = "chacha20poly1305", feature = "hkdf"))]
     #[test]
     fn encrypt_multi_roundtrip() {
+        use crate::helper::{decrypt_multi, encrypt_multi};
         use crate::keypair::Keyring;
-        use crate::helper::{encrypt_multi, decrypt_multi};
 
         let kr = Keyring::generate();
         let entities = vec![kr.public_key_bundle()];

crypto/src/sign.rs

@@ -172,7 +172,9 @@ impl MlDsaSigner {
 impl SignatureScheme for MlDsaSigner {
     fn sign(&self, msg: &[u8]) -> Result<Vec<u8>, CryptoError> {
         use ml_dsa::Signer;
-        let signature = self.secret.try_sign(msg)
+        let signature = self
+            .secret
+            .try_sign(msg)
             .map_err(|_| CryptoError::SigningFailed)?;
         Ok(signature.encode().to_vec())
     }
@@ -181,7 +183,8 @@ impl SignatureScheme for MlDsaSigner {
         use ml_dsa::Verifier;
         let sig = ml_dsa::Signature::<ml_dsa::MlDsa65>::try_from(signature)
             .map_err(|_| CryptoError::InvalidSignature)?;
-        self.public.verify(msg, &sig)
+        self.public
+            .verify(msg, &sig)
             .map_err(|_| CryptoError::VerificationFailed)
     }
 }
@@ -218,19 +221,20 @@ pub fn sign_dual(
     ed25519_sk: &ed25519_dalek::SigningKey,
     mldsa_sk: &ml_dsa::SigningKey<ml_dsa::MlDsa65>,
     message: &[u8],
-) -> DualSignature {
+) -> Result<DualSignature, CryptoError> {
     let ed25519 = {
         use ed25519_dalek::Signer;
         ed25519_sk.sign(message).to_bytes().to_vec()
     };
     let mldsa = {
         use ml_dsa::Signer;
-        mldsa_sk.try_sign(message)
-            .expect("ML-DSA signing failed")
+        mldsa_sk
+            .try_sign(message)
+            .map_err(|_| CryptoError::SigningFailed)?
             .encode()
             .to_vec()
     };
-    DualSignature { ed25519, mldsa }
+    Ok(DualSignature { ed25519, mldsa })
 }
 
 impl DualSignature {

deny.toml

@@ -0,0 +1,32 @@
+# cargo-deny configuration. See https://embarkstudios.github.io/cargo-deny/
+# Run locally with: cargo deny check
+
+[advisories]
+# Fail on any security advisory affecting the dependency tree.
+yanked = "deny"
+ignore = []
+
+[bans]
+# Flag multiple versions of the same crate so duplicate trees are visible.
+multiple-versions = "warn"
+wildcards = "deny"
+
+[licenses]
+# Allowlist of licenses acceptable for this project's dependencies.
+allow = [
+    "MIT",
+    "Apache-2.0",
+    "Apache-2.0 WITH LLVM-exception",
+    "BSD-2-Clause",
+    "BSD-3-Clause",
+    "ISC",
+    "Unicode-3.0",
+    "Zlib",
+    "MPL-2.0",
+]
+confidence-threshold = 0.8
+
+[sources]
+unknown-registry = "deny"
+unknown-git = "deny"
+allow-registry = ["https://github.com/rust-lang/crates.io-index"]

example-usage/client/src/main.rs

@@ -47,9 +47,10 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         client_id: 0,
     };
 
+    let server_bundle = host_public_key.clone();
     let (conn, keyring) =
         auth::connect_or_register(config, host_public_key, "client_keys.json").await?;
-    messages::send_and_receive(&conn, &keyring).await?;
+    messages::send_and_receive(&conn, &keyring, &server_bundle).await?;
 
     println!("\nDone");
     Ok(())

example-usage/client/src/messages.rs

@@ -1,44 +1,43 @@
 use mtp::client::MTPConnection;
-use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataTypeId, DataValue};
-use mtp::crypto::{ChaCha20Poly1305, Ed25519Signer, Keyring, SigAlgorithm};
+use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
+use mtp::crypto::{Ed25519Signer, EncryptionType, Keyring, PublicKeyBundle, SigAlgorithm};
+use mtp::type_map::TypeMap;
 
-fn derive_demo_key() -> [u8; 32] {
-    mtp::crypto::derive_encryption_key(
-        b"MTP-demo-shared-secret",
-        b"MTP-demo-salt",
-        b"encrypted-container-demo",
-    )
-    .expect("key derivation must succeed")
-}
+pub fn build_demo_message(
+    client_id: u64,
+    keyring: &Keyring,
+    server_bundle: &PublicKeyBundle,
+) -> CommunicationValue {
+    // Encrypt to the server's KEM public key; the server decrypts with its keyring.
+    let enc_type = EncryptionType::MlKemChaCha20Poly1305;
+    let signer = Ed25519Signer::new(&keyring.sig_cl_secret_key)
+        .expect("Ed25519 signer from keyring");
 
-pub fn build_demo_message(client_id: u64, keyring: &Keyring) -> CommunicationValue {
-    let cipher = ChaCha20Poly1305::new(derive_demo_key());
-    let signer =
-        Ed25519Signer::new(&keyring.sig_cl_secret_key).expect("Ed25519 signer from keyring");
+    let tm = TypeMap::latest();
 
     let inner_enc = DataValue::Container(vec![
-        (DataTypeId(1), DataValue::Str("secret inner data".into())),
-        (DataTypeId(2), DataValue::UnsignedNumber(42)),
+        (DataType::Version.to_id(&tm), DataValue::Str("secret inner data".into())),
+        (DataType::Id.to_id(&tm), DataValue::UnsignedNumber(42)),
     ]);
     let mut dv_enc = inner_enc;
-    dv_enc.encrypt_container(&cipher, b"demo-aad");
+    dv_enc.encrypt_container(enc_type, server_bundle, b"demo-aad");
 
     let inner_sig = DataValue::Container(vec![
-        (DataTypeId(1), DataValue::Str("signed by client".into())),
-        (DataTypeId(2), DataValue::UnsignedNumber(99)),
+        (DataType::Version.to_id(&tm), DataValue::Str("signed by client".into())),
+        (DataType::Id.to_id(&tm), DataValue::UnsignedNumber(99)),
     ]);
     let mut dv_sig = inner_sig;
     dv_sig.sign_container(SigAlgorithm::ED25519, &signer);
 
     let inner_sec = DataValue::Container(vec![
         (
-            DataTypeId(1),
+            DataType::Version.to_id(&tm),
             DataValue::Str("signed+encrypted payload".into()),
         ),
-        (DataTypeId(2), DataValue::UnsignedNumber(7)),
+        (DataType::Id.to_id(&tm), DataValue::UnsignedNumber(7)),
     ]);
     let mut dv_sec = inner_sec;
-    dv_sec.sign_and_encrypt_container(SigAlgorithm::ED25519, &signer, &cipher, b"demo-aad");
+    dv_sec.sign_and_encrypt_container(SigAlgorithm::ED25519, &signer, enc_type, server_bundle, b"demo-aad");
 
     let timestamp = std::time::SystemTime::now()
         .duration_since(std::time::UNIX_EPOCH)
@@ -79,8 +78,9 @@ pub fn build_demo_message(client_id: u64, keyring: &Keyring) -> CommunicationVal
 pub async fn send_and_receive(
     conn: &MTPConnection,
     keyring: &Keyring,
+    server_bundle: &PublicKeyBundle,
 ) -> Result<(), Box<dyn std::error::Error>> {
-    let msg = build_demo_message(conn.client_id, keyring);
+    let msg = build_demo_message(conn.client_id, keyring, server_bundle);
     println!("Sending: {msg}");
     conn.sender.send(&msg).await?;
 

example-usage/server/src/handlers.rs

@@ -1,6 +1,6 @@
 use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataTypeId, DataValue, TypeMap};
 use mtp::crypto::{
-    ChaCha20Poly1305, CryptoError, SignaturePublicKey, SignatureScheme, verify_ed25519,
+    CryptoError, Keyring, SignaturePublicKey, SignatureScheme, verify_ed25519,
 };
 
 struct Ed25519Verifier(SignaturePublicKey);
@@ -14,19 +14,11 @@ impl SignatureScheme for Ed25519Verifier {
     }
 }
 
-fn derive_demo_key() -> [u8; 32] {
-    mtp::crypto::derive_encryption_key(
-        b"MTP-demo-shared-secret",
-        b"MTP-demo-salt",
-        b"encrypted-container-demo",
-    )
-    .expect("key derivation must succeed")
-}
-
 pub fn process_and_respond(
     msg: &CommunicationValue,
     tm: &TypeMap,
     client_pk: Option<&mtp::crypto::PublicKeyBundle>,
+    host_keyring: &Keyring,
 ) -> CommunicationValue {
     let desc_id = DataTypeId(tm.data_id_enum(DataType::Description).unwrap());
     let ts_id = DataTypeId(tm.data_id_enum(DataType::Timestamp).unwrap());
@@ -58,8 +50,6 @@ pub fn process_and_respond(
     println!("  Binary: {:?}", binary.as_bytes());
     println!("  Items: {:?}", items.as_array());
 
-    let cipher = ChaCha20Poly1305::new(derive_demo_key());
-
     let mut enc_status = String::from("EncryptedPayload: not present");
     let mut sig_status = String::from("SignedPayload: not present");
     let mut secure_status = String::from("SecurePayload: not present");
@@ -67,7 +57,7 @@ pub fn process_and_respond(
     let enc = msg.get_data(enc_id);
     if matches!(enc, DataValue::EncryptedContainer(_)) {
         let mut dv = enc.clone();
-        if dv.decrypt_into_container(&cipher, b"demo-aad").is_some() {
+        if dv.decrypt_into_container(host_keyring, b"demo-aad").is_some() {
             if let Some(entries) = dv.as_container() {
                 println!("  Decrypted EncryptedPayload: {:?}", entries);
                 enc_status = format!("EncryptedPayload decrypted OK ({} entries)", entries.len());
@@ -103,9 +93,7 @@ pub fn process_and_respond(
         if let Some(pk_bundle) = client_pk {
             let verifier = Ed25519Verifier(pk_bundle.sig_cl_public_key.clone());
             let mut dv = secure.clone();
-            if dv
-                .decrypt_signed_encrypted_container(&cipher, b"demo-aad")
-                .is_some()
+            if dv.decrypt_signed_encrypted_container(host_keyring, b"demo-aad").is_some()
                 && dv.verify_into_container(&verifier).is_some()
             {
                 if let Some(entries) = dv.as_container() {

example-usage/server/src/main.rs

@@ -36,6 +36,11 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let (host_id, host_keyring) = keys::load_or_generate_host_keys("host_keys.json")?;
     keys::export_host_public_keys(&host_keyring)?;
 
+    // The keyring is moved into the host config; keep a copy for decrypting the
+    // demo payloads clients encrypt to our KEM public key.
+    let decrypt_keyring = mtp::crypto::Keyring::from_bytes(&host_keyring.to_bytes())
+        .expect("re-load host keyring for decryption");
+
     let (clients, next_id) = clients::load_client_db("clients.json")?;
 
     let clients_for_get = clients.clone();
@@ -98,8 +103,12 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         match conn.receiver.receive().await {
             Ok(msg) => {
                 println!("Received: {msg}");
-                let response =
-                    handlers::process_and_respond(&msg, tm, conn.client_public_key.as_ref());
+                let response = handlers::process_and_respond(
+                    &msg,
+                    tm,
+                    conn.client_public_key.as_ref(),
+                    &decrypt_keyring,
+                );
                 println!("Sending: {response}");
                 conn.sender.send(&response).await?;
             }

example-usage/web-client/src/main.ts

@@ -213,7 +213,7 @@ async function connect() {
     }
 
     log("\nSending demo message...");
-    const frame = build_demo_message(activeClientId, keyringBytes);
+    const frame = build_demo_message(activeClientId, keyringBytes, hostPk);
     log(`Sending: ${format_frame(frame)}`, "state");
     await client.send(frame);
     log(`Sent ${frame.length} bytes`);

example-usage/web-client/src/messages.ts

@@ -12,8 +12,12 @@ export function buildAuthResponse(
   return parse_auth_response(response);
 }
 
-export function buildDemoMessage(clientId: bigint, keyringBytes: Uint8Array): Uint8Array {
-  return build_demo_message(clientId, keyringBytes);
+export function buildDemoMessage(
+  clientId: bigint,
+  keyringBytes: Uint8Array,
+  hostBundle: Uint8Array,
+): Uint8Array {
+  return build_demo_message(clientId, keyringBytes, hostBundle);
 }
 
 export function buildPingFrame(

flake.lock

@@ -20,11 +20,11 @@
     },
     "nixpkgs": {
       "locked": {
-        "lastModified": 1781577229,
-        "narHash": "sha256-lrp67w8AulE9Ks53n27I45ADSzbOCn4H+CNW1Ck8B+8=",
+        "lastModified": 1782467914,
+        "narHash": "sha256-pGvFkM8N0xEkIIXDe5YYfbEAvHrk4IxBrjB/x8OomhE=",
         "owner": "NixOS",
         "repo": "nixpkgs",
-        "rev": "567a49d1913ce81ac6e9582e3553dd90a955875f",
+        "rev": "e73de5be04e0eff4190a1432b946d469c794e7b4",
         "type": "github"
       },
       "original": {
@@ -62,11 +62,11 @@
         "nixpkgs": "nixpkgs_2"
       },
       "locked": {
-        "lastModified": 1782357464,
-        "narHash": "sha256-mXgoT1qDHCdSfF9IvhMtEEFNy9dxrmUfSViwP7RpzOQ=",
+        "lastModified": 1782443907,
+        "narHash": "sha256-P+pADLtK7qC1mz0/5Xq9uF77oahUR4zYLTaitiHsUHg=",
         "owner": "oxalica",
         "repo": "rust-overlay",
-        "rev": "77a8263847fb02dc49dbe377278ef6b952f1c6bb",
+        "rev": "4b06ff4acf3491ff69721df852507fcc51d0a13d",
         "type": "github"
       },
       "original": {

flake.nix

@@ -94,7 +94,7 @@
               mprocs
             ];
             shellHook = ''
-              nix develop --command bash -c "mprocs 'cd wasm && wasm-pack build --target web --out-dir pkg && cd ../example-usage/web-client && bun dev' 'cargo b && cd example-usage && cargo r --bin server'"
+              nix develop --command bash -c "mprocs 'cd wasm && wasm-pack build --target web --out-dir pkg && cd ../example-usage/web-client && bun dev' 'cargo b && cd example-usage && cargo r --bin server' 'cd example-usage && cargo r --bin client'"
               exit
             '';
           };

host/Cargo.toml

@@ -10,6 +10,7 @@ mtp-transport = { path = "../transport", features = ["host"] }
 mtp-crypto = { path = "../crypto", optional = true }
 tokio = { version = "1", features = ["full"] }
 rand = "0.8"
+log = "0.4"
 
 [features]
 crypto = ["dep:mtp-crypto", "mtp-codec/crypto"]

host/src/lib.rs

@@ -1,14 +1,12 @@
-#[cfg(feature = "crypto")]
-use mtp_codec::DataType;
 use mtp_codec::{
-    CommunicationValue, DataTypeId, DataValue, Version,
+    CommunicationValue, DataType, DataValue, TypeMap, Version,
     registry::{Registry, VersionedCodec},
 };
 use mtp_common::CommunicationError;
 use mtp_transport::{Policy, Receiver, Sender};
 use std::net::IpAddr;
 
-// Host configuration.
+/* Host configuration. */
 pub struct HostConfig {
     pub ip: IpAddr,
     pub port: u16,
@@ -35,7 +33,7 @@ pub enum AuthState {
     Failed,
 }
 
-// A connection that has completed version negotiation.
+/* A connection that has completed version negotiation. */
 pub struct MTPConnection {
     pub version: Version,
     pub codec: VersionedCodec,
@@ -49,7 +47,7 @@ pub struct MTPConnection {
     pub client_public_key: Option<mtp_crypto::PublicKeyBundle>,
 }
 
-// High-level MTP host with built-in version negotiation.
+/* High-level MTP host with built-in version negotiation. */
 pub struct MTPHost {
     transport: mtp_transport::Host,
     registry: Registry,
@@ -86,7 +84,13 @@ impl MTPHost {
      * incompatible.
      */
     pub async fn accept(&mut self) -> Option<MTPConnection> {
-        let (sender, receiver) = self.transport.next().await?;
+        let (sender, receiver) = match self.transport.next().await {
+            Some(pair) => pair,
+            None => {
+                log::warn!("accept: transport.next() returned None (listener closed)");
+                return None;
+            }
+        };
 
         #[cfg(feature = "crypto")]
         if self.config.require_authentication {
@@ -94,7 +98,13 @@ impl MTPHost {
         }
 
         // Read the first message (always encoded with reserved types).
-        let first_msg = receiver.receive().await.ok()?;
+        let first_msg = match receiver.receive().await {
+            Ok(m) => m,
+            Err(e) => {
+                log::warn!("accept: receive failed: {e:?}");
+                return None;
+            }
+        };
 
         /*
          * Extract the client's version from the first message.
@@ -103,9 +113,24 @@ impl MTPHost {
          *
          * Then negotiate the version for single-version clients
          */
-        let client_version = extract_version(&first_msg)?;
+        let client_version = match extract_version(&first_msg) {
+            Some(v) => v,
+            None => {
+                log::warn!(
+                    "accept: extract_version failed on msg type {:?}",
+                    first_msg.get_type()
+                );
+                return None;
+            }
+        };
 
-        let negotiated = self.registry.negotiate(&[client_version])?;
+        let negotiated = match self.registry.negotiate(&[client_version.clone()]) {
+            Some(v) => v,
+            None => {
+                log::warn!("accept: negotiate failed for client version {client_version:?}");
+                return None;
+            }
+        };
 
         let codec = VersionedCodec::new(self.registry.clone());
 
@@ -134,119 +159,107 @@ impl MTPHost {
 
 #[cfg(feature = "crypto")]
 impl MTPHost {
+    /*
+     * Mutually-authenticated handshake with a server-issued challenge.
+     *
+     *   1. C -> H : Identification { version, id }   (or Register { version, public_keys })
+     *   2. H -> C : Challenge { server_challenge, host_sig }
+     *   3. C -> H : ChallengeResponse { client_nonce, sig }
+     *   4. H -> C : IdentificationResponse / RegisterResponse { connected, id, sig }
+     *
+     * The client's authenticating signature (step 3) covers `server_challenge`,
+     * a fresh value generated here in step 2 and kept on this task's stack for
+     * the lifetime of the connection. It is therefore one-time per connection
+     * with no shared replay state, and a captured proof cannot be replayed on
+     * any other connection.
+     */
     async fn accept_authenticated(
         &mut self,
         sender: Sender,
         receiver: Receiver,
     ) -> Option<MTPConnection> {
         use mtp_crypto::{
-            Ed25519Signer, PublicKeyBundle, SignatureScheme, verify_ed25519, verify_ml_dsa,
+            Ed25519Signer, MlDsaSigner, PublicKeyBundle, SignatureScheme, auth, verify_ed25519,
+            verify_ml_dsa,
         };
 
-        // 1. Receive client message first (no host greeting)
-        let msg = receiver.receive().await.ok()?;
-        let version_str = match msg.get_data(DataTypeId(3)) {
+        // Flow-specific state resolved from the client's opening hello.
+        enum Flow {
+            Login {
+                id: u64,
+                bundle: PublicKeyBundle,
+            },
+            Register {
+                bundle: PublicKeyBundle,
+                pk_bytes: Vec<u8>,
+            },
+        }
+
+        let tm = TypeMap::latest();
+        let pq_enabled = !self
+            .config
+            .host_keyring
+            .sig_pq_secret_key
+            .as_bytes()
+            .is_empty();
+
+        // Sign `payload` with the host keys (Ed25519 always, ML-DSA when configured).
+        let host_sign = |payload: &[u8]| -> Option<(Vec<u8>, Vec<u8>)> {
+            let signer = Ed25519Signer::new(&self.config.host_keyring.sig_cl_secret_key).ok()?;
+            let sig = signer.sign(payload).ok()?;
+            let pq_sig = if pq_enabled {
+                let pq = MlDsaSigner::new(
+                    &self.config.host_keyring.sig_pq_secret_key,
+                    &self.config.host_keyring.sig_pq_public_key,
+                )
+                .ok()?;
+                pq.sign(payload).ok()?
+            } else {
+                Vec::new()
+            };
+            Some((sig, pq_sig))
+        };
+
+        // ===== Step 1: receive the client's unsigned hello =====
+        let hello = receiver.receive().await.ok()?;
+        let version_str = match hello.get_data(DataType::Version.to_id(&tm)) {
             DataValue::Str(s) => s.clone(),
             _ => {
                 sender.close();
                 return None;
             }
         };
-
         let client_version = Version::parse(&version_str)?;
 
-        let client_nonce = match msg.get_data(DataTypeId(7)) {
-            DataValue::UnsignedNumber(n) => *n,
-            _ => {
-                sender.close();
-                return None;
-            }
-        };
-
-        let sig_bytes = match msg.get_data(DataTypeId(10)) {
-            DataValue::Bytes(b) => b.clone(),
-            _ => {
-                sender.close();
-                return None;
-            }
-        };
-
-        let pq_sig_bytes: Vec<u8> = match msg.get_data(DataTypeId(12)) {
-            DataValue::Bytes(b) => b.clone(),
-            _ => vec![],
-        };
-
-        let (assigned_id, client_bundle, response_type) = if msg.get_type()
-            == mtp_codec::CommunicationTypeId(15)
-        {
-            // LOGIN
-            let cid = match msg.get_data(DataTypeId(6)) {
+        let (flow, response_type) =
+            if hello.get_type() == mtp_codec::CommunicationType::Identification.to_id(&tm) {
+                // LOGIN: look up the claimed user before issuing a challenge.
+                let cid = match hello.get_data(DataType::Id.to_id(&tm)) {
                     DataValue::UnsignedNumber(n) => *n as u64,
                     _ => {
                         sender.close();
                         return None;
                     }
                 };
-
                 let bundle = match (self.config.get_existing_user)(cid) {
                     Some(b) => b,
                     None => {
                         let rejection = CommunicationValue::new(
                             mtp_codec::CommunicationType::IdentificationResponse,
                         )
-                    .add_typed_default(DataType::Connected, DataValue::BoolFalse)
-                    .add_typed_default(
-                        DataType::ClientNonce,
-                        DataValue::UnsignedNumber(client_nonce),
-                    );
+                        .add_typed_default(DataType::Connected, DataValue::BoolFalse);
                         let _ = sender.send(&rejection).await;
                         sender.close();
                         return None;
                     }
                 };
-
-            let mut sig_payload = Vec::new();
-            sig_payload.extend_from_slice(version_str.as_bytes());
-            sig_payload.extend_from_slice(&cid.to_be_bytes());
-            sig_payload.extend_from_slice(&client_nonce.to_be_bytes());
-
-            /* ===== Signature ===== */
-            if verify_ed25519(&bundle.sig_cl_public_key, &sig_payload, &sig_bytes).is_err() {
-                let rejection =
-                    CommunicationValue::new(mtp_codec::CommunicationType::IdentificationResponse)
-                        .add_typed_default(DataType::Connected, DataValue::BoolFalse)
-                        .add_typed_default(
-                            DataType::ClientNonce,
-                            DataValue::UnsignedNumber(client_nonce),
-                        );
-                let _ = sender.send(&rejection).await;
-                sender.close();
-                return None;
-            }
-            if !pq_sig_bytes.is_empty()
-                && verify_ml_dsa(&bundle.sig_pq_public_key, &sig_payload, &pq_sig_bytes).is_err()
-            {
-                let rejection =
-                    CommunicationValue::new(mtp_codec::CommunicationType::IdentificationResponse)
-                        .add_typed_default(DataType::Connected, DataValue::BoolFalse)
-                        .add_typed_default(
-                            DataType::ClientNonce,
-                            DataValue::UnsignedNumber(client_nonce),
-                        );
-                let _ = sender.send(&rejection).await;
-                sender.close();
-                return None;
-            }
-            /* ===== End Signature ===== */
-
                 (
-                cid,
-                bundle,
+                    Flow::Login { id: cid, bundle },
                     mtp_codec::CommunicationType::IdentificationResponse,
                 )
-        } else if msg.get_type() == mtp_codec::CommunicationTypeId(17) {
-            // REGISTER
-            let bundle = match msg.get_data(DataTypeId(9)) {
+            } else if hello.get_type() == mtp_codec::CommunicationType::Register.to_id(&tm) {
+                // REGISTER: the client presents the bundle it wants to register.
+                let bundle = match hello.get_data(DataType::PublicKeys.to_id(&tm)) {
                     DataValue::Bytes(b) => PublicKeyBundle::from_bytes(b).ok()?,
                     _ => {
                         sender.close();
@@ -254,45 +267,8 @@ impl MTPHost {
                     }
                 };
                 let pk_bytes = bundle.as_bytes();
-
-            let mut sig_payload = Vec::new();
-            sig_payload.extend_from_slice(version_str.as_bytes());
-            sig_payload.extend_from_slice(&client_nonce.to_be_bytes());
-            sig_payload.extend_from_slice(&pk_bytes);
-
-            /* ===== Signature ===== */
-            if verify_ed25519(&bundle.sig_cl_public_key, &sig_payload, &sig_bytes).is_err() {
-                let rejection =
-                    CommunicationValue::new(mtp_codec::CommunicationType::RegisterResponse)
-                        .add_typed_default(DataType::Connected, DataValue::BoolFalse)
-                        .add_typed_default(
-                            DataType::ClientNonce,
-                            DataValue::UnsignedNumber(client_nonce),
-                        );
-                let _ = sender.send(&rejection).await;
-                sender.close();
-                return None;
-            }
-            if !pq_sig_bytes.is_empty()
-                && verify_ml_dsa(&bundle.sig_pq_public_key, &sig_payload, &pq_sig_bytes).is_err()
-            {
-                let rejection =
-                    CommunicationValue::new(mtp_codec::CommunicationType::RegisterResponse)
-                        .add_typed_default(DataType::Connected, DataValue::BoolFalse)
-                        .add_typed_default(
-                            DataType::ClientNonce,
-                            DataValue::UnsignedNumber(client_nonce),
-                        );
-                let _ = sender.send(&rejection).await;
-                sender.close();
-                return None;
-            }
-            /* ===== End Signature ===== */
-
-            let new_id = (self.config.complete_register)(bundle.clone());
                 (
-                new_id,
-                bundle,
+                    Flow::Register { bundle, pk_bytes },
                     mtp_codec::CommunicationType::RegisterResponse,
                 )
             } else {
@@ -300,53 +276,118 @@ impl MTPHost {
                 return None;
             };
 
-        // 2. Send success response (single host message)
-        let new_nonce: u128 = rand::random();
+        // The id bound into the challenge (0 for register: none assigned yet).
+        let challenge_id = match &flow {
+            Flow::Login { id, .. } => *id,
+            Flow::Register { .. } => 0,
+        };
 
-        let mut host_sig_payload = Vec::new();
-        host_sig_payload.push(0x01);
-        host_sig_payload.extend_from_slice(&assigned_id.to_be_bytes());
-        host_sig_payload.extend_from_slice(&client_nonce.to_be_bytes());
-        host_sig_payload.extend_from_slice(&new_nonce.to_be_bytes());
+        // ===== Step 2: issue a fresh, host-signed challenge =====
+        let server_challenge: u128 = rand::random();
+        let (chal_sig, chal_pq_sig) =
+            host_sign(&auth::challenge_payload(challenge_id, server_challenge))?;
 
-        let host_signer = Ed25519Signer::new(&self.config.host_keyring.sig_cl_secret_key).ok()?;
+        let mut challenge_msg = CommunicationValue::new(mtp_codec::CommunicationType::Challenge)
+            .add_typed_default(
+                DataType::ServerNonce,
+                DataValue::UnsignedNumber(server_challenge),
+            )
+            .add_typed_default(DataType::Signature, DataValue::Bytes(chal_sig));
+        if pq_enabled {
+            challenge_msg = challenge_msg
+                .add_typed_default(DataType::PqSignature, DataValue::Bytes(chal_pq_sig));
+        }
+        sender.send(&challenge_msg).await.ok()?;
 
-        /* ===== Signature ===== */
-        let host_sig = host_signer.sign(&host_sig_payload).ok()?;
+        // ===== Step 3: receive and verify the client's proof =====
+        let proof = receiver.receive().await.ok()?;
+        if proof.get_type() != mtp_codec::CommunicationType::ChallengeResponse.to_id(&tm) {
+            sender.close();
+            return None;
+        }
+        let client_nonce = match proof.get_data(DataType::ClientNonce.to_id(&tm)) {
+            DataValue::UnsignedNumber(n) => *n,
+            _ => {
+                sender.close();
+                return None;
+            }
+        };
+        let sig_bytes = match proof.get_data(DataType::Signature.to_id(&tm)) {
+            DataValue::Bytes(b) => b.clone(),
+            _ => {
+                sender.close();
+                return None;
+            }
+        };
+        let pq_sig_bytes: Vec<u8> = match proof.get_data(DataType::PqSignature.to_id(&tm)) {
+            DataValue::Bytes(b) => b.clone(),
+            _ => vec![],
+        };
+
+        let (proof_payload, bundle) = match &flow {
+            Flow::Login { id, bundle } => (
+                auth::login_proof_payload(&version_str, *id, server_challenge, client_nonce),
+                bundle,
+            ),
+            Flow::Register {
+                bundle, pk_bytes, ..
+            } => (
+                auth::register_proof_payload(
+                    &version_str,
+                    pk_bytes,
+                    server_challenge,
+                    client_nonce,
+                ),
+                bundle,
+            ),
+        };
+
+        let proof_ok = verify_ed25519(&bundle.sig_cl_public_key, &proof_payload, &sig_bytes)
+            .is_ok()
+            && (pq_sig_bytes.is_empty()
+                || verify_ml_dsa(&bundle.sig_pq_public_key, &proof_payload, &pq_sig_bytes).is_ok());
+
+        if !proof_ok {
+            let rejection = CommunicationValue::new(response_type)
+                .add_typed_default(DataType::Connected, DataValue::BoolFalse);
+            let _ = sender.send(&rejection).await;
+            sender.close();
+            return None;
+        }
+
+        // Proof verified: resolve the assigned id and retain the client's bundle.
+        let (assigned_id, client_bundle) = match flow {
+            Flow::Login { id, bundle } => (id, bundle),
+            Flow::Register { bundle, .. } => {
+                let new_id = (self.config.complete_register)(bundle.clone());
+                (new_id, bundle)
+            }
+        };
+
+        // ===== Step 4: send the host's final confirmation =====
+        let (host_sig, host_pq_sig) = host_sign(&auth::host_final_payload(
+            assigned_id,
+            client_nonce,
+            server_challenge,
+        ))?;
 
         let mut response = CommunicationValue::new(response_type)
             .add_typed_default(DataType::Connected, DataValue::BoolTrue)
+            .add_typed_default(DataType::Id, DataValue::UnsignedNumber(assigned_id as u128))
             .add_typed_default(
                 DataType::ClientNonce,
                 DataValue::UnsignedNumber(client_nonce),
             )
-            .add_typed_default(DataType::Timestamp, DataValue::UnsignedNumber(new_nonce))
-            .add_typed_default(DataType::Signature, DataValue::Bytes(host_sig))
-            .add_typed_default(DataType::Id, DataValue::UnsignedNumber(assigned_id as u128));
-
-        if !self
-            .config
-            .host_keyring
-            .sig_pq_secret_key
-            .as_bytes()
-            .is_empty()
-        {
-            use mtp_crypto::MlDsaSigner;
-            let host_pq_signer = MlDsaSigner::new(
-                &self.config.host_keyring.sig_pq_secret_key,
-                &self.config.host_keyring.sig_pq_public_key,
-            )
-            .ok()?;
-            let host_pq_sig = host_pq_signer.sign(&host_sig_payload).ok()?;
+            .add_typed_default(DataType::Signature, DataValue::Bytes(host_sig));
+        if pq_enabled {
             response =
                 response.add_typed_default(DataType::PqSignature, DataValue::Bytes(host_pq_sig));
         }
-        /* ===== End Signature ===== */
 
         sender.send(&response).await.ok()?;
         sender.finish_stream().await.ok()?;
 
-        // 3. Version negotiation
+        // ===== Version negotiation =====
         let negotiated = self.registry.negotiate(&[client_version])?;
         let codec = VersionedCodec::new(self.registry.clone());
 
@@ -369,7 +410,8 @@ impl MTPHost {
  * (reserved ID 3) mapping to `DataValue::Str("major.minor")`.
  */
 fn extract_version(msg: &CommunicationValue) -> Option<Version> {
-    let value = msg.get_data(DataTypeId(3));
+    let tm = TypeMap::latest();
+    let value = msg.get_data(DataType::Version.to_id(&tm));
     match value {
         DataValue::Str(s) => Version::parse(s.as_str()),
         _ => None,
@@ -388,7 +430,10 @@ mod tests {
             mtp_codec::CommunicationType::Identification,
             &tm,
         )
-        .add_data(DataTypeId(3), DataValue::Str("2.0".to_string()));
+        .add_data(
+            DataType::Version.to_id(&tm),
+            DataValue::Str("2.0".to_string()),
+        );
         let version = extract_version(&msg);
         assert_eq!(version, Some(Version(2, 0)));
     }
@@ -410,7 +455,7 @@ mod tests {
             mtp_codec::CommunicationType::Identification,
             &tm,
         )
-        .add_data(DataTypeId(3), DataValue::UnsignedNumber(42));
+        .add_data(DataType::Version.to_id(&tm), DataValue::UnsignedNumber(42));
         assert!(extract_version(&msg).is_none());
     }
 }

rustfmt.toml

@@ -0,0 +1 @@
+edition = "2024"

src/lib.rs

@@ -1,7 +1,7 @@
-pub use mtp_common as common;
-pub use mtp_type_map as type_map;
 pub use mtp_codec as codec;
+pub use mtp_common as common;
 pub use mtp_transport as transport;
+pub use mtp_type_map as type_map;
 
 #[cfg(feature = "crypto")]
 pub use mtp_crypto as crypto;

transport/Cargo.toml

@@ -24,6 +24,10 @@ log = "0.4"
 [dev-dependencies]
 rcgen = "0.14"
 
+[[test]]
+name = "integration"
+required-features = ["host"]
+
 [features]
 default = []
 # Enables hosting a MTP server

transport/src/connection.rs

@@ -83,7 +83,7 @@ impl Sender {
         data: &CommunicationValue,
         policy: &Policy,
     ) -> Result<(), CommunicationError> {
-        let bytes = data.to_bytes();
+        let bytes = data.to_bytes().map_err(|_| CommunicationError::Encode)?;
         if bytes.len() as u64 > policy.max_message_size
             || bytes.len() as u64 >= policy.close_frame_len as u64
         {
@@ -344,6 +344,18 @@ pub struct Receiver {
     handle: Arc<ConnectionHandle>,
 }
 
+impl Drop for Receiver {
+    fn drop(&mut self) {
+        // The accept loop holds clones of the connection and the shared
+        // ConnectionHandle. Without this, dropping a Receiver without first
+        // closing the connection would leave that task running forever. Abort
+        // it directly rather than closing the shared handle, so a still-live
+        // Sender on the same connection is unaffected. abort() is a no-op if
+        // the task already finished (e.g. the connection was closed).
+        self._accept_task.abort();
+    }
+}
+
 impl Receiver {
     pub fn new(connection: Connection, handle: Arc<ConnectionHandle>, policy: Arc<Policy>) -> Self {
         let (tx, rx) = mpsc::channel::<Result<CommunicationValue, CommunicationError>>(

transport/src/connection_handle.rs

@@ -53,7 +53,7 @@ impl ConnectionHandle {
         if self.is_closed() {
             return rx.borrow().clone();
         }
-        let _ = rx.changed().await.ok()?;
+        rx.changed().await.ok()?;
         rx.borrow().clone()
     }
 }

transport/src/host.rs

@@ -26,6 +26,22 @@ impl Host {
     pub fn local_addr(&self) -> std::net::SocketAddr {
         self.local_addr
     }
+
+    /// Stop accepting new connections. Already-accepted connections run on their
+    /// own spawned tasks and are not affected.
+    pub fn shutdown(&mut self) {
+        self._task.abort();
+    }
+}
+
+impl Drop for Host {
+    fn drop(&mut self) {
+        // The accept loop runs forever on its own task; dropping the Host
+        // JoinHandle would only detach it. Abort it so dropping the Host
+        // actually stops accepting new connections. Per-connection handler
+        // tasks are spawned independently and keep running.
+        self._task.abort();
+    }
 }
 
 pub async fn host(

transport/tests/integration.rs

@@ -1,6 +1,7 @@
 use std::net::{IpAddr, Ipv4Addr};
 
-use mtp_transport::{Policy, host, connect};
+use mtp_codec::{CommunicationType, DataType, TypeMap};
+use mtp_transport::{Policy, connect, host};
 
 fn generate_self_signed_cert() -> (Vec<u8>, Vec<u8>) {
     let key_pair = rcgen::KeyPair::generate().unwrap();
@@ -44,38 +45,44 @@ async fn test_send_receive_roundtrip() {
     let addr = h.local_addr();
 
     let url = format!("https://127.0.0.1:{}", addr.port());
-    let (client_tx, client_rx) =
-        connect(&url, Some(cert_pem), Policy::default()).await.unwrap();
+    let (client_tx, client_rx) = connect(&url, Some(cert_pem), Policy::default())
+        .await
+        .unwrap();
 
     // Accept on host side
     let (host_tx, host_rx) = h.next().await.unwrap();
 
+    let tm = TypeMap::latest();
+
     // Client sends a simple message
-    let msg = mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::Ping)
-        .add_data(
-            mtp_codec::DataTypeId(6),
+    let msg = mtp_codec::CommunicationValue::new(CommunicationType::Ping).add_data(
+        DataType::PqSignature.to_id(&tm),
         mtp_codec::DataValue::UnsignedNumber(42),
     );
     client_tx.send(&msg).await.unwrap();
 
     // Host receives it
     let received = host_rx.receive().await.unwrap();
-    assert_eq!(received.get_type(), mtp_codec::CommunicationTypeId(19)); // Ping
-    let val = received.get_data(mtp_codec::DataTypeId(6)).clone();
+    assert_eq!(received.get_type(), CommunicationType::Ping.to_id(&tm));
+    let val = received.get_data(DataType::PqSignature.to_id(&tm)).clone();
     assert_eq!(val, mtp_codec::DataValue::UnsignedNumber(42));
 
     // Host sends a response
-    let resp = mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::Pong)
-        .add_data(
-            mtp_codec::DataTypeId(6),
+    let resp = mtp_codec::CommunicationValue::new(CommunicationType::Pong).add_data(
+        DataType::PqSignature.to_id(&tm),
         mtp_codec::DataValue::UnsignedNumber(99),
     );
     host_tx.send(&resp).await.unwrap();
 
     // Client receives it
     let client_received = client_rx.receive().await.unwrap();
-    assert_eq!(client_received.get_type(), mtp_codec::CommunicationTypeId(20)); // Pong
-    let client_val = client_received.get_data(mtp_codec::DataTypeId(6)).clone();
+    assert_eq!(
+        client_received.get_type(),
+        CommunicationType::Pong.to_id(&tm)
+    );
+    let client_val = client_received
+        .get_data(DataType::PqSignature.to_id(&tm))
+        .clone();
     assert_eq!(client_val, mtp_codec::DataValue::UnsignedNumber(99));
 
     // Close both sides
@@ -98,16 +105,18 @@ async fn test_concurrent_messages() {
     let addr = h.local_addr();
 
     let url = format!("https://127.0.0.1:{}", addr.port());
-    let (client_tx, _client_rx) =
-        connect(&url, Some(cert_pem), Policy::default()).await.unwrap();
+    let (client_tx, _client_rx) = connect(&url, Some(cert_pem), Policy::default())
+        .await
+        .unwrap();
 
     let (_host_tx, host_rx) = h.next().await.unwrap();
 
+    let tm = TypeMap::latest();
+
     // Send 5 messages in sequence
     for i in 0..5u128 {
-        let msg = mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::Ping)
-            .add_data(
-                mtp_codec::DataTypeId(6),
+        let msg = mtp_codec::CommunicationValue::new(CommunicationType::Ping).add_data(
+            DataType::PqSignature.to_id(&tm),
             mtp_codec::DataValue::UnsignedNumber(i),
         );
         client_tx.send(&msg).await.unwrap();
@@ -116,15 +125,14 @@ async fn test_concurrent_messages() {
     // Receive all 5 in order
     for i in 0..5u128 {
         let received = host_rx.receive().await.unwrap();
-        let val = received.get_data(mtp_codec::DataTypeId(6)).clone();
+        let val = received.get_data(DataType::PqSignature.to_id(&tm)).clone();
         assert_eq!(val, mtp_codec::DataValue::UnsignedNumber(i));
     }
 
     // Send 3 responses back
     for i in 0..3u128 {
-        let msg = mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::Pong)
-            .add_data(
-                mtp_codec::DataTypeId(6),
+        let msg = mtp_codec::CommunicationValue::new(CommunicationType::Pong).add_data(
+            DataType::PqSignature.to_id(&tm),
             mtp_codec::DataValue::UnsignedNumber(i * 10),
         );
         client_tx.send(&msg).await.unwrap();
@@ -132,7 +140,7 @@ async fn test_concurrent_messages() {
 
     for i in 0..3u128 {
         let received = host_rx.receive().await.unwrap();
-        let val = received.get_data(mtp_codec::DataTypeId(6)).clone();
+        let val = received.get_data(DataType::PqSignature.to_id(&tm)).clone();
         assert_eq!(val, mtp_codec::DataValue::UnsignedNumber(i * 10));
     }
 
@@ -154,21 +162,109 @@ async fn test_close_detection() {
     let addr = h.local_addr();
 
     let url = format!("https://127.0.0.1:{}", addr.port());
-    let (client_tx, _client_rx) =
-        connect(&url, Some(cert_pem), Policy::default()).await.unwrap();
+    let (client_tx, _client_rx) = connect(&url, Some(cert_pem), Policy::default())
+        .await
+        .unwrap();
 
     let (_host_tx, host_rx) = h.next().await.unwrap();
 
     // Send a message then close
-    let msg = mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::Ping);
+    let msg = mtp_codec::CommunicationValue::new(CommunicationType::Ping);
     client_tx.send(&msg).await.unwrap();
     client_tx.close();
 
     // Host should still receive the message
+    let tm = TypeMap::latest();
     let received = host_rx.receive().await.unwrap();
-    assert_eq!(received.get_type(), mtp_codec::CommunicationTypeId(19)); // Ping
+    assert_eq!(received.get_type(), CommunicationType::Ping.to_id(&tm));
 
     // Host should get an error or closed signal on next receive
     let result = host_rx.receive().await;
     assert!(result.is_err());
 }
+
+#[tokio::test]
+async fn test_host_shutdown_stops_accepting() {
+    let (cert_pem, key_pem) = generate_self_signed_cert();
+    let mut h = host(
+        IpAddr::V4(Ipv4Addr::LOCALHOST),
+        0,
+        cert_pem.clone(),
+        key_pem,
+        Policy::default(),
+    )
+    .await
+    .unwrap();
+    let addr = h.local_addr();
+    let url = format!("https://127.0.0.1:{}", addr.port());
+
+    // A connection succeeds while the host is accepting.
+    let (_c_tx, _c_rx) = connect(&url, Some(cert_pem.clone()), Policy::default())
+        .await
+        .unwrap();
+    let _accepted = h.next().await.unwrap();
+
+    // After shutdown the accept task is aborted and its endpoint is dropped, so
+    // new connections no longer succeed. Guard with a timeout so a hung connect
+    // still fails the assertion rather than blocking the test.
+    h.shutdown();
+
+    let result = tokio::time::timeout(
+        std::time::Duration::from_secs(5),
+        connect(&url, Some(cert_pem), Policy::default()),
+    )
+    .await;
+    assert!(
+        matches!(result, Err(_) | Ok(Err(_))),
+        "connect should not succeed after host shutdown"
+    );
+}
+
+#[tokio::test]
+async fn test_drop_receiver_keeps_sender_alive() {
+    let (cert_pem, key_pem) = generate_self_signed_cert();
+    let mut h = host(
+        IpAddr::V4(Ipv4Addr::LOCALHOST),
+        0,
+        cert_pem.clone(),
+        key_pem,
+        Policy::default(),
+    )
+    .await
+    .unwrap();
+    let addr = h.local_addr();
+
+    let url = format!("https://127.0.0.1:{}", addr.port());
+    let (client_tx, client_rx) = connect(&url, Some(cert_pem), Policy::default())
+        .await
+        .unwrap();
+
+    let (host_tx, host_rx) = h.next().await.unwrap();
+
+    // Client sends a message the host receives.
+    let msg = mtp_codec::CommunicationValue::new(CommunicationType::Ping);
+    client_tx.send(&msg).await.unwrap();
+    let _ = host_rx.receive().await.unwrap();
+
+    // Dropping the host Receiver aborts only its accept task; the Sender shares
+    // the same connection and must keep working.
+    drop(host_rx);
+
+    let tm = TypeMap::latest();
+
+    let resp = mtp_codec::CommunicationValue::new(CommunicationType::Pong).add_data(
+        DataType::PqSignature.to_id(&tm),
+        mtp_codec::DataValue::UnsignedNumber(7),
+    );
+    host_tx.send(&resp).await.unwrap();
+
+    let got = client_rx.receive().await.unwrap();
+    assert_eq!(got.get_type(), CommunicationType::Pong.to_id(&tm));
+    assert_eq!(
+        got.get_data(DataType::PqSignature.to_id(&tm)).clone(),
+        mtp_codec::DataValue::UnsignedNumber(7)
+    );
+
+    client_tx.close();
+    host_tx.close();
+}

type-map/build.rs

@@ -25,139 +25,150 @@ struct ReservedEntry {
 
 const RESERVED_COMM_TYPES: &[ReservedEntry] = &[
     ReservedEntry {
-        name: "Error",
+        name: "Identification",
         id: 0,
     },
     ReservedEntry {
-        name: "ErrorParsing",
+        name: "IdentificationResponse",
         id: 1,
     },
     ReservedEntry {
-        name: "ErrorBadVersion",
+        name: "Register",
         id: 2,
     },
     ReservedEntry {
-        name: "Disconnect",
+        name: "RegisterResponse",
         id: 3,
     },
     ReservedEntry {
-        name: "Redirect",
+        name: "Challenge",
         id: 4,
     },
     ReservedEntry {
-        name: "Shutdown",
+        name: "ChallengeResponse",
         id: 5,
     },
     ReservedEntry {
-        name: "BadRequest",
+        name: "Ping",
         id: 6,
     },
     ReservedEntry {
-        name: "Unauthorized",
+        name: "Pong",
         id: 7,
     },
     ReservedEntry {
-        name: "Forbidden",
+        name: "Disconnect",
         id: 8,
     },
     ReservedEntry {
-        name: "NotFound",
+        name: "Redirect",
         id: 9,
     },
     ReservedEntry {
-        name: "TooManyRequests",
+        name: "Shutdown",
         id: 10,
     },
     ReservedEntry {
-        name: "InternalServerError",
+        name: "Error",
         id: 11,
     },
     ReservedEntry {
-        name: "BadGateway",
+        name: "ErrorParsing",
         id: 12,
     },
     ReservedEntry {
-        name: "ServiceUnavailable",
+        name: "ErrorBadVersion",
         id: 13,
     },
     ReservedEntry {
-        name: "GatewayTimeout",
+        name: "BadRequest",
         id: 14,
     },
     ReservedEntry {
-        name: "Identification",
+        name: "Unauthorized",
         id: 15,
     },
     ReservedEntry {
-        name: "IdentificationResponse",
+        name: "Forbidden",
         id: 16,
     },
     ReservedEntry {
-        name: "Register",
+        name: "NotFound",
         id: 17,
     },
     ReservedEntry {
-        name: "RegisterResponse",
+        name: "TooManyRequests",
         id: 18,
     },
     ReservedEntry {
-        name: "Ping",
+        name: "InternalServerError",
         id: 19,
     },
     ReservedEntry {
-        name: "Pong",
+        name: "BadGateway",
         id: 20,
     },
+    ReservedEntry {
+        name: "ServiceUnavailable",
+        id: 21,
+    },
+    ReservedEntry {
+        name: "GatewayTimeout",
+        id: 22,
+    },
 ];
 
 const RESERVED_DATA_TYPES: &[ReservedEntry] = &[
     ReservedEntry {
-        name: "Error",
+        name: "Version",
         id: 0,
     },
     ReservedEntry {
-        name: "ErrorParsing",
+        name: "Id",
         id: 1,
     },
     ReservedEntry {
-        name: "ErrorMessage",
+        name: "ClientNonce",
         id: 2,
     },
     ReservedEntry {
-        name: "Version",
+        name: "ServerNonce",
         id: 3,
     },
     ReservedEntry {
-        name: "Description",
+        name: "PublicKeys",
         id: 4,
     },
-    ReservedEntry {
-        name: "Timestamp",
-        id: 5,
-    },
-    ReservedEntry { name: "Id", id: 6 },
-    ReservedEntry {
-        name: "ClientNonce",
-        id: 7,
-    },
-    ReservedEntry {
-        name: "ServerNonce",
-        id: 8,
-    },
-    ReservedEntry {
-        name: "PublicKeys",
-        id: 9,
-    },
     ReservedEntry {
         name: "Signature",
-        id: 10,
-    },
-    ReservedEntry {
-        name: "Connected",
-        id: 11,
+        id: 5,
     },
     ReservedEntry {
         name: "PqSignature",
+        id: 6,
+    },
+    ReservedEntry {
+        name: "Description",
+        id: 7,
+    },
+    ReservedEntry {
+        name: "Connected",
+        id: 8,
+    },
+    ReservedEntry {
+        name: "Timestamp",
+        id: 9,
+    },
+    ReservedEntry {
+        name: "Error",
+        id: 10,
+    },
+    ReservedEntry {
+        name: "ErrorParsing",
+        id: 11,
+    },
+    ReservedEntry {
+        name: "ErrorMessage",
         id: 12,
     },
 ];
@@ -220,7 +231,7 @@ fn generate(config: &Config, multi_version: bool) -> String {
     let (all_comm_names, all_data_names, sorted) = if multi_version {
         let mut all_comm_names: BTreeSet<&str> = BTreeSet::new();
         let mut all_data_names: BTreeSet<&str> = BTreeSet::new();
-        for (_version_key, tm) in &config.type_maps {
+        for tm in config.type_maps.values() {
             for name in tm.communication_types.keys() {
                 all_comm_names.insert(name.as_str());
             }

type-map/src/lib.rs

@@ -1,3 +1,10 @@
+/*
+ * try_new returns Result<Self, ()> deliberately: the only failure mode is "id
+ * is in the reserved range", which carries no extra information worth an error
+ * type. The unit error is the intended API.
+ */
+#![allow(clippy::result_unit_err)]
+
 pub const INTERNAL_COMM_RESERVED: std::ops::Range<u16> = 0..32;
 pub const INTERNAL_DATA_RESERVED: std::ops::Range<u16> = 0..32;
 

wasm/.cargo/config.toml

@@ -1,5 +1,6 @@
+# Default to the wasm32 target when running cargo from inside this crate dir.
+# The web_sys_unstable_apis cfg lives in the workspace-root .cargo/config.toml,
+# scoped to [target.wasm32-unknown-unknown], so it applies here too (the root
+# config is an ancestor) and to `-p mtp-wasm` builds invoked from the root.
 [build]
 target = "wasm32-unknown-unknown"
-
-[target.wasm32-unknown-unknown]
-rustflags = ["--cfg=web_sys_unstable_apis"]

wasm/src/client.rs

@@ -4,7 +4,7 @@ use std::rc::Rc;
 use wasm_bindgen::prelude::*;
 
 use mtp_codec::{CommunicationType, CommunicationValue, DataType, DataValue, PROTOCOL_VERSION};
-use mtp_type_map::{CommunicationTypeId, DataTypeId};
+use mtp_type_map::CommunicationTypeId;
 
 use mtp_crypto::SignatureScheme;
 
@@ -36,6 +36,73 @@ fn unexpected_response_type_error(
     ))
 }
 
+/*
+ * Verify the host's signature over the challenge it issued (step 2), mirroring
+ * the native client (`client/src/lib.rs`). `id` is the client id for a login or
+ * `0` for a registration. The Ed25519 signature is mandatory; the ML-DSA
+ * signature is verified only when the host included one.
+ */
+fn verify_host_challenge(
+    challenge: &CommunicationValue,
+    tm: &mtp_codec::TypeMap,
+    host_pk: &mtp_crypto::PublicKeyBundle,
+    id: u64,
+    server_challenge: u128,
+) -> Result<(), JsValue> {
+    let sig = match challenge.get_data(DataType::Signature.to_id(tm)) {
+        DataValue::Bytes(b) => b.clone(),
+        _ => return Err(js_error("missing host challenge signature")),
+    };
+    let pq_sig = match challenge.get_data(DataType::PqSignature.to_id(tm)) {
+        DataValue::Bytes(b) => b.clone(),
+        _ => vec![],
+    };
+
+    let payload = mtp_crypto::auth::challenge_payload(id, server_challenge);
+    mtp_crypto::verify_ed25519(&host_pk.sig_cl_public_key, &payload, &sig)
+        .map_err(|_| js_error("host challenge signature invalid"))?;
+    if !pq_sig.is_empty() {
+        mtp_crypto::verify_ml_dsa(&host_pk.sig_pq_public_key, &payload, &pq_sig)
+            .map_err(|_| js_error("host challenge PQ signature invalid"))?;
+    }
+    Ok(())
+}
+
+/*
+ * Verify the host's final confirmation (step 4): the echoed `client_nonce` and
+ * the host signature over the handshake transcript. `id` is the client id for a
+ * login and the host-assigned id for a register.
+ */
+fn verify_host_final(
+    resp: &CommunicationValue,
+    tm: &mtp_codec::TypeMap,
+    host_pk: &mtp_crypto::PublicKeyBundle,
+    id: u64,
+    client_nonce: u128,
+    server_challenge: u128,
+) -> Result<(), JsValue> {
+    if *resp.get_data(DataType::ClientNonce.to_id(tm)) != DataValue::UnsignedNumber(client_nonce) {
+        return Err(js_error("nonce mismatch"));
+    }
+    let host_sig = match resp.get_data(DataType::Signature.to_id(tm)) {
+        DataValue::Bytes(b) => b.clone(),
+        _ => return Err(js_error("missing host signature")),
+    };
+    let host_pq_sig = match resp.get_data(DataType::PqSignature.to_id(tm)) {
+        DataValue::Bytes(b) => b.clone(),
+        _ => vec![],
+    };
+
+    let payload = mtp_crypto::auth::host_final_payload(id, client_nonce, server_challenge);
+    mtp_crypto::verify_ed25519(&host_pk.sig_cl_public_key, &payload, &host_sig)
+        .map_err(|_| js_error("host signature invalid"))?;
+    if !host_pq_sig.is_empty() {
+        mtp_crypto::verify_ml_dsa(&host_pk.sig_pq_public_key, &payload, &host_pq_sig)
+            .map_err(|_| js_error("host PQ signature invalid"))?;
+    }
+    Ok(())
+}
+
 #[wasm_bindgen]
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum ConnectionState {
@@ -126,7 +193,6 @@ impl WasmClient {
         self.set_state(ConnectionState::Connecting);
         let transport =
             WasmTransport::connect(&config.url, config.server_certificate_hashes.clone()).await?;
-        let inner = transport.inner().clone();
 
         let version_str = format!("{}", PROTOCOL_VERSION);
         let ident = CommunicationValue::new(CommunicationType::Identification)
@@ -135,8 +201,12 @@ impl WasmClient {
                 DataType::Id,
                 DataValue::UnsignedNumber(config.client_id as u128),
             );
-        transport.send_frame(&ident.to_bytes()).await?;
+        let ident_bytes = ident
+            .to_bytes()
+            .map_err(|e| js_error(&format!("encode failed: {}", e)))?;
+        transport.send_frame(&ident_bytes).await?;
 
+        let loop_transport = transport.clone();
         self.transport = Some(transport);
         self.set_state(ConnectionState::Connected);
 
@@ -144,9 +214,7 @@ impl WasmClient {
         let on_msg = self.on_message.clone();
         let on_err = self.on_error.clone();
         wasm_bindgen_futures::spawn_local(async move {
-            WasmTransport::from_inner(inner)
-                .receive_loop(on_msg, on_err)
-                .await;
+            loop_transport.receive_loop(on_msg, on_err).await;
             state.set(ConnectionState::Disconnected);
         });
         Ok(())
@@ -170,51 +238,89 @@ impl WasmClient {
     ) -> Result<u64, JsValue> {
         self.set_state(ConnectionState::Connecting);
 
-        let _host_pk = mtp_crypto::PublicKeyBundle::from_bytes(host_public_key_bytes)
+        let host_pk = mtp_crypto::PublicKeyBundle::from_bytes(host_public_key_bytes)
             .map_err(|e| js_error(&format!("invalid host public key: {}", e)))?;
         let keyring = mtp_crypto::Keyring::from_bytes(keyring_bytes)
             .map_err(|e| js_error(&format!("invalid keyring: {}", e)))?;
 
+        let tm = mtp_codec::TypeMap::latest();
         let version_str = format!("{}", PROTOCOL_VERSION);
+
+        let transport =
+            WasmTransport::connect(&config.url, config.server_certificate_hashes.clone()).await?;
+
+        // 1. Send the unsigned Identification hello.
+        let hello = CommunicationValue::new(CommunicationType::Identification)
+            .add_typed_default(DataType::Version, DataValue::Str(version_str.clone()))
+            .add_typed_default(DataType::Id, DataValue::UnsignedNumber(client_id as u128))
+            .to_bytes()
+            .map_err(|e| js_error(&format!("encode failed: {}", e)))?;
+        transport.send_frame(&hello).await?;
+
+        // 2. Receive and verify the host's challenge.
+        let challenge_bytes = transport.read_one_frame().await?;
+        let challenge = CommunicationValue::from_bytes(&challenge_bytes)
+            .map_err(|e| js_error(&format!("parse challenge: {}", e)))?;
+        let expected = CommunicationType::Challenge.to_id(&tm);
+        if challenge.get_type() != expected {
+            self.set_state(ConnectionState::Disconnected);
+            return Err(unexpected_response_type_error(
+                "auth_connect challenge",
+                expected,
+                challenge.get_type(),
+                &challenge_bytes,
+                &challenge,
+            ));
+        }
+        let server_challenge = match challenge.get_data(DataType::ServerNonce.to_id(&tm)) {
+            DataValue::UnsignedNumber(n) => *n,
+            _ => {
+                self.set_state(ConnectionState::Disconnected);
+                return Err(js_error("missing server challenge"));
+            }
+        };
+        if let Err(e) =
+            verify_host_challenge(&challenge, &tm, &host_pk, client_id, server_challenge)
+        {
+            self.set_state(ConnectionState::Disconnected);
+            return Err(e);
+        }
+
+        // 3. Sign the host's challenge and send the proof.
         let mut nonce_bytes = [0u8; 16];
         getrandom::fill(&mut nonce_bytes).map_err(|_| js_error("rng failed"))?;
         let client_nonce = u128::from_be_bytes(nonce_bytes);
 
-        // Build signature payload: version || client_id || client_nonce
-        let mut sig_payload = Vec::new();
-        sig_payload.extend_from_slice(version_str.as_bytes());
-        sig_payload.extend_from_slice(&client_id.to_be_bytes());
-        sig_payload.extend_from_slice(&client_nonce.to_be_bytes());
-
+        let proof_payload = mtp_crypto::auth::login_proof_payload(
+            &version_str,
+            client_id,
+            server_challenge,
+            client_nonce,
+        );
         let signer = mtp_crypto::Ed25519Signer::new(&keyring.sig_cl_secret_key)
             .map_err(|e| js_error(&format!("signer creation failed: {}", e)))?;
         let signature = signer
-            .sign(&sig_payload)
+            .sign(&proof_payload)
             .map_err(|e| js_error(&format!("signature failed: {}", e)))?;
 
-        let frame = CommunicationValue::new(CommunicationType::Identification)
-            .add_typed_default(DataType::Version, DataValue::Str(version_str))
-            .add_typed_default(DataType::Id, DataValue::UnsignedNumber(client_id as u128))
+        let proof = CommunicationValue::new(CommunicationType::ChallengeResponse)
             .add_typed_default(
                 DataType::ClientNonce,
                 DataValue::UnsignedNumber(client_nonce),
             )
             .add_typed_default(DataType::Signature, DataValue::Bytes(signature))
-            .to_bytes();
+            .to_bytes()
+            .map_err(|e| js_error(&format!("encode failed: {}", e)))?;
+        transport.send_frame(&proof).await?;
 
-        let transport =
-            WasmTransport::connect(&config.url, config.server_certificate_hashes.clone()).await?;
-        let inner = transport.inner().clone();
-        transport.send_frame(&frame).await?;
-
-        // Read and verify the host's IdentificationResponse
+        // 4. Receive and verify the host's final confirmation.
         let response = transport.read_one_frame().await?;
         let resp_comm = CommunicationValue::from_bytes(&response)
             .map_err(|e| js_error(&format!("parse response: {}", e)))?;
-
         let resp_type = resp_comm.get_type();
-        let expected_type = CommunicationTypeId(16); // IdentificationResponse
+        let expected_type = CommunicationType::IdentificationResponse.to_id(&tm);
         if resp_type != expected_type {
+            self.set_state(ConnectionState::Disconnected);
             return Err(unexpected_response_type_error(
                 "auth_connect",
                 expected_type,
@@ -224,22 +330,34 @@ impl WasmClient {
             ));
         }
 
-        if resp_comm.get_data(DataTypeId(11)) != &DataValue::BoolTrue {
+        if resp_comm.get_data(DataType::Connected.to_id(&tm)) != &DataValue::BoolTrue {
+            self.set_state(ConnectionState::Disconnected);
             return Err(js_error("host rejected authentication"));
         }
 
-        // Verify echoed nonce
-        let echo_nonce = resp_comm.get_data(DataTypeId(7));
-        if *echo_nonce != DataValue::UnsignedNumber(client_nonce) {
-            return Err(js_error("nonce mismatch"));
+        // Verify echoed nonce + host signature (login: id is client_id).
+        if let Err(e) = verify_host_final(
+            &resp_comm,
+            &tm,
+            &host_pk,
+            client_id,
+            client_nonce,
+            server_challenge,
+        ) {
+            self.set_state(ConnectionState::Disconnected);
+            return Err(e);
         }
 
         // Extract assigned ID
-        let assigned_id = match resp_comm.get_data(DataTypeId(6)) {
+        let assigned_id = match resp_comm.get_data(DataType::Id.to_id(&tm)) {
             DataValue::UnsignedNumber(n) => *n as u64,
-            _ => return Err(js_error("missing assigned ID")),
+            _ => {
+                self.set_state(ConnectionState::Disconnected);
+                return Err(js_error("missing assigned ID"));
+            }
         };
 
+        let loop_transport = transport.clone();
         self.transport = Some(transport);
         self.set_state(ConnectionState::Connected);
 
@@ -247,9 +365,7 @@ impl WasmClient {
         let on_msg = self.on_message.clone();
         let on_err = self.on_error.clone();
         wasm_bindgen_futures::spawn_local(async move {
-            WasmTransport::from_inner(inner)
-                .receive_loop(on_msg, on_err)
-                .await;
+            loop_transport.receive_loop(on_msg, on_err).await;
             state.set(ConnectionState::Disconnected);
         });
 
@@ -272,52 +388,88 @@ impl WasmClient {
     ) -> Result<u64, JsValue> {
         self.set_state(ConnectionState::Connecting);
 
-        let _host_pk = mtp_crypto::PublicKeyBundle::from_bytes(host_public_key_bytes)
+        let host_pk = mtp_crypto::PublicKeyBundle::from_bytes(host_public_key_bytes)
             .map_err(|e| js_error(&format!("invalid host public key: {}", e)))?;
         let keyring = mtp_crypto::Keyring::from_bytes(keyring_bytes)
             .map_err(|e| js_error(&format!("invalid keyring: {}", e)))?;
 
+        let tm = mtp_codec::TypeMap::latest();
         let version_str = format!("{}", PROTOCOL_VERSION);
+        let pk_bytes = keyring.public_key_bundle().as_bytes();
+
+        let transport =
+            WasmTransport::connect(&config.url, config.server_certificate_hashes.clone()).await?;
+
+        // 1. Send the unsigned Register hello (version + public-key bundle).
+        let hello = CommunicationValue::new(CommunicationType::Register)
+            .add_typed_default(DataType::Version, DataValue::Str(version_str.clone()))
+            .add_typed_default(DataType::PublicKeys, DataValue::Bytes(pk_bytes.clone()))
+            .to_bytes()
+            .map_err(|e| js_error(&format!("encode failed: {}", e)))?;
+        transport.send_frame(&hello).await?;
+
+        // 2. Receive and verify the host's challenge (register binds id = 0).
+        let challenge_bytes = transport.read_one_frame().await?;
+        let challenge = CommunicationValue::from_bytes(&challenge_bytes)
+            .map_err(|e| js_error(&format!("parse challenge: {}", e)))?;
+        let expected = CommunicationType::Challenge.to_id(&tm);
+        if challenge.get_type() != expected {
+            self.set_state(ConnectionState::Disconnected);
+            return Err(unexpected_response_type_error(
+                "auth_register challenge",
+                expected,
+                challenge.get_type(),
+                &challenge_bytes,
+                &challenge,
+            ));
+        }
+        let server_challenge = match challenge.get_data(DataType::ServerNonce.to_id(&tm)) {
+            DataValue::UnsignedNumber(n) => *n,
+            _ => {
+                self.set_state(ConnectionState::Disconnected);
+                return Err(js_error("missing server challenge"));
+            }
+        };
+        if let Err(e) = verify_host_challenge(&challenge, &tm, &host_pk, 0, server_challenge) {
+            self.set_state(ConnectionState::Disconnected);
+            return Err(e);
+        }
+
+        // 3. Sign the host's challenge over the bundle and send the proof.
         let mut nonce_bytes = [0u8; 16];
         getrandom::fill(&mut nonce_bytes).map_err(|_| js_error("rng failed"))?;
         let client_nonce = u128::from_be_bytes(nonce_bytes);
 
-        let pk_bytes = keyring.public_key_bundle().as_bytes();
-
-        // Build signature payload: version || client_nonce || pk_bytes
-        let mut sig_payload = Vec::new();
-        sig_payload.extend_from_slice(version_str.as_bytes());
-        sig_payload.extend_from_slice(&client_nonce.to_be_bytes());
-        sig_payload.extend_from_slice(&pk_bytes);
-
+        let proof_payload = mtp_crypto::auth::register_proof_payload(
+            &version_str,
+            &pk_bytes,
+            server_challenge,
+            client_nonce,
+        );
         let signer = mtp_crypto::Ed25519Signer::new(&keyring.sig_cl_secret_key)
             .map_err(|e| js_error(&format!("signer creation failed: {}", e)))?;
         let signature = signer
-            .sign(&sig_payload)
+            .sign(&proof_payload)
             .map_err(|e| js_error(&format!("signature failed: {}", e)))?;
 
-        let frame = CommunicationValue::new(CommunicationType::Register)
-            .add_typed_default(DataType::Version, DataValue::Str(version_str))
+        let proof = CommunicationValue::new(CommunicationType::ChallengeResponse)
             .add_typed_default(
                 DataType::ClientNonce,
                 DataValue::UnsignedNumber(client_nonce),
             )
-            .add_typed_default(DataType::PublicKeys, DataValue::Bytes(pk_bytes))
             .add_typed_default(DataType::Signature, DataValue::Bytes(signature))
-            .to_bytes();
-
-        let transport =
-            WasmTransport::connect(&config.url, config.server_certificate_hashes.clone()).await?;
-        let inner = transport.inner().clone();
-        transport.send_frame(&frame).await?;
+            .to_bytes()
+            .map_err(|e| js_error(&format!("encode failed: {}", e)))?;
+        transport.send_frame(&proof).await?;
 
+        // 4. Receive the host's final confirmation; extract + verify assigned id.
         let response = transport.read_one_frame().await?;
         let resp_comm = CommunicationValue::from_bytes(&response)
             .map_err(|e| js_error(&format!("parse response: {}", e)))?;
-
         let resp_type = resp_comm.get_type();
-        let expected_type = CommunicationTypeId(18); // RegisterResponse
+        let expected_type = CommunicationType::RegisterResponse.to_id(&tm);
         if resp_type != expected_type {
+            self.set_state(ConnectionState::Disconnected);
             return Err(unexpected_response_type_error(
                 "auth_register",
                 expected_type,
@@ -327,20 +479,33 @@ impl WasmClient {
             ));
         }
 
-        if resp_comm.get_data(DataTypeId(11)) != &DataValue::BoolTrue {
+        if resp_comm.get_data(DataType::Connected.to_id(&tm)) != &DataValue::BoolTrue {
+            self.set_state(ConnectionState::Disconnected);
             return Err(js_error("host rejected registration"));
         }
 
-        let echo = resp_comm.get_data(DataTypeId(7));
-        if *echo != DataValue::UnsignedNumber(client_nonce) {
-            return Err(js_error("nonce mismatch"));
-        }
-
-        let assigned_id = match resp_comm.get_data(DataTypeId(6)) {
+        let assigned_id = match resp_comm.get_data(DataType::Id.to_id(&tm)) {
             DataValue::UnsignedNumber(n) => *n as u64,
-            _ => return Err(js_error("missing assigned ID")),
+            _ => {
+                self.set_state(ConnectionState::Disconnected);
+                return Err(js_error("missing assigned ID"));
+            }
         };
 
+        // Verify echoed nonce + host signature (register: id is host-assigned).
+        if let Err(e) = verify_host_final(
+            &resp_comm,
+            &tm,
+            &host_pk,
+            assigned_id,
+            client_nonce,
+            server_challenge,
+        ) {
+            self.set_state(ConnectionState::Disconnected);
+            return Err(e);
+        }
+
+        let loop_transport = transport.clone();
         self.transport = Some(transport);
         self.set_state(ConnectionState::Connected);
 
@@ -348,9 +513,7 @@ impl WasmClient {
         let on_msg = self.on_message.clone();
         let on_err = self.on_error.clone();
         wasm_bindgen_futures::spawn_local(async move {
-            WasmTransport::from_inner(inner)
-                .receive_loop(on_msg, on_err)
-                .await;
+            loop_transport.receive_loop(on_msg, on_err).await;
             state.set(ConnectionState::Disconnected);
         });
 

wasm/src/crypto.rs

@@ -1,9 +1,9 @@
 use wasm_bindgen::prelude::*;
 
 use mtp_crypto::{
-    AeadDecrypt, AeadEncrypt, Ed25519Signer, KemPrivateKey, KemPublicKey, Keyring,
-    PublicKeyBundle, SignaturePqPrivateKey, SignaturePqPublicKey, SignaturePrivateKey,
-    SignaturePublicKey, SignatureScheme, ChaCha20Poly1305, sha256, sha256_double,
+    AeadDecrypt, AeadEncrypt, ChaCha20Poly1305, Ed25519Signer, KemPrivateKey, KemPublicKey,
+    Keyring, PublicKeyBundle, SignaturePqPrivateKey, SignaturePqPublicKey, SignaturePrivateKey,
+    SignaturePublicKey, SignatureScheme, sha256, sha256_double,
 };
 
 use crate::error::js_error;
@@ -28,8 +28,8 @@ impl WasmKeyring {
     /// Deserialise a keyring from bytes.
     #[wasm_bindgen]
     pub fn from_bytes(bytes: &[u8]) -> Result<WasmKeyring, JsValue> {
-        let inner =
-            Keyring::from_bytes(bytes).map_err(|e| js_error(&format!("Keyring::from_bytes: {}", e)))?;
+        let inner = Keyring::from_bytes(bytes)
+            .map_err(|e| js_error(&format!("Keyring::from_bytes: {}", e)))?;
         Ok(Self { inner })
     }
 
@@ -217,7 +217,11 @@ pub fn ed25519_generate() -> Result<JsValue, JsValue> {
 
 /// Standalone Ed25519 signature verification.
 #[wasm_bindgen]
-pub fn ed25519_verify(public_key: Vec<u8>, message: &[u8], signature: &[u8]) -> Result<(), JsValue> {
+pub fn ed25519_verify(
+    public_key: Vec<u8>,
+    message: &[u8],
+    signature: &[u8],
+) -> Result<(), JsValue> {
     let pk = SignaturePublicKey::new(public_key);
     mtp_crypto::verify_ed25519(&pk, message, signature)
         .map_err(|e| js_error(&format!("verify_ed25519 failed: {}", e)))
@@ -245,7 +249,12 @@ pub fn wasm_sha256_double(data: &[u8]) -> Vec<u8> {
 
 /// HKDF-expand: derive `len` bytes from `ikm` with `salt` and `info`.
 #[wasm_bindgen]
-pub fn wasm_hkdf_expand(ikm: &[u8], salt: &[u8], info: &[u8], len: usize) -> Result<Vec<u8>, JsValue> {
+pub fn wasm_hkdf_expand(
+    ikm: &[u8],
+    salt: &[u8],
+    info: &[u8],
+    len: usize,
+) -> Result<Vec<u8>, JsValue> {
     mtp_crypto::hkdf_expand(ikm, salt, info, len)
         .map_err(|e| js_error(&format!("hkdf_expand failed: {}", e)))
 }
@@ -416,7 +425,8 @@ mod tests {
     fn sha256_empty() {
         let result = wasm_sha256(b"");
         // SHA-256 of empty string
-        let expected = hex::decode("e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855")
+        let expected =
+            hex::decode("e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855")
                 .expect("hex decode");
         assert_eq!(result, expected);
     }
@@ -424,7 +434,8 @@ mod tests {
     #[wasm_bindgen_test]
     fn sha256_hello() {
         let result = wasm_sha256(b"hello");
-        let expected = hex::decode("2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824")
+        let expected =
+            hex::decode("2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824")
                 .expect("hex decode");
         assert_eq!(result, expected);
     }
@@ -443,8 +454,7 @@ mod tests {
 
     #[wasm_bindgen_test]
     fn hkdf_expand_produces_correct_length() {
-        let result = wasm_hkdf_expand(b"ikm", b"salt", b"info", 32)
-            .expect("hkdf_expand failed");
+        let result = wasm_hkdf_expand(b"ikm", b"salt", b"info", 32).expect("hkdf_expand failed");
         assert_eq!(result.len(), 32);
     }
 
@@ -457,22 +467,21 @@ mod tests {
 
     #[wasm_bindgen_test]
     fn derive_encryption_key_roundtrip() {
-        let key = wasm_derive_encryption_key(b"password", b"salt", b"context")
-            .expect("derive failed");
+        let key =
+            wasm_derive_encryption_key(b"password", b"salt", b"context").expect("derive failed");
         assert_eq!(key.len(), 32);
 
         // Deterministic: same inputs = same key
-        let key2 = wasm_derive_encryption_key(b"password", b"salt", b"context")
-            .expect("derive failed");
+        let key2 =
+            wasm_derive_encryption_key(b"password", b"salt", b"context").expect("derive failed");
         assert_eq!(key, key2);
     }
 
     #[wasm_bindgen_test]
     fn derive_encryption_key_different_inputs_different_key() {
-        let key = wasm_derive_encryption_key(b"pass1", b"salt", b"context")
-            .expect("derive failed");
-        let key2 = wasm_derive_encryption_key(b"pass2", b"salt", b"context")
-            .expect("derive failed");
+        let key = wasm_derive_encryption_key(b"pass1", b"salt", b"context").expect("derive failed");
+        let key2 =
+            wasm_derive_encryption_key(b"pass2", b"salt", b"context").expect("derive failed");
         assert_ne!(key, key2);
     }
 }

wasm/src/lib.rs

@@ -4,6 +4,7 @@ pub mod error;
 pub mod message;
 pub mod transport;
 
+#[cfg(not(test))]
 use wasm_bindgen::prelude::*;
 
 #[cfg(not(test))]

wasm/src/message.rs

@@ -1,14 +1,19 @@
 use wasm_bindgen::prelude::*;
 
 use mtp_codec::{CommunicationType, CommunicationValue, DataType, DataTypeId, DataValue};
-use mtp_crypto::{ChaCha20Poly1305, Ed25519Signer, Keyring, SigAlgorithm, derive_encryption_key};
-use mtp_type_map::communication_type_name;
+use mtp_crypto::{Ed25519Signer, EncryptionType, Keyring, PublicKeyBundle, SigAlgorithm};
+use mtp_type_map::{communication_type_name, TypeMap};
 
 use crate::error::js_error;
 
 /// Build a simple Ping frame with description, timestamp, and optional data.
 #[wasm_bindgen]
-pub fn build_ping_frame(client_id: u64, description: &str, timestamp: u64, data: &[u8]) -> Vec<u8> {
+pub fn build_ping_frame(
+    client_id: u64,
+    description: &str,
+    timestamp: u64,
+    data: &[u8],
+) -> Result<Vec<u8>, JsValue> {
     let mut msg = CommunicationValue::new(CommunicationType::Ping)
         .add_typed_default(
             DataType::Description,
@@ -25,40 +30,42 @@ pub fn build_ping_frame(client_id: u64, description: &str, timestamp: u64, data:
     }
 
     msg.to_bytes()
+        .map_err(|e| js_error(&format!("encode failed: {}", e)))
 }
 
 /// Build a demo Ping frame with encrypted and signed containers
 /// (mirrors the Rust client example but uses only reserved data types).
 #[wasm_bindgen]
-pub fn build_demo_message(client_id: u64, keyring_bytes: &[u8]) -> Result<Vec<u8>, JsValue> {
+pub fn build_demo_message(
+    client_id: u64,
+    keyring_bytes: &[u8],
+    host_bundle_bytes: &[u8],
+) -> Result<Vec<u8>, JsValue> {
     let keyring = Keyring::from_bytes(keyring_bytes)
         .map_err(|e| js_error(&format!("invalid keyring: {}", e)))?;
 
-    let enc_key = derive_encryption_key(
-        b"MTP-demo-shared-secret",
-        b"MTP-demo-salt",
-        b"encrypted-container-demo",
-    )
-    .map_err(|e| js_error(&format!("key derivation failed: {}", e)))?;
-
-    let cipher = ChaCha20Poly1305::new(enc_key);
+    // Encrypt to the server's KEM public key; the server decrypts with its keyring.
+    // (The client keyring only needs the Ed25519 signing key for this demo.)
+    let recipient = PublicKeyBundle::from_bytes(host_bundle_bytes)
+        .map_err(|e| js_error(&format!("invalid host bundle: {}", e)))?;
+    let enc_type = EncryptionType::MlKemChaCha20Poly1305;
     let signer = Ed25519Signer::new(&keyring.sig_cl_secret_key)
         .map_err(|e| js_error(&format!("signer creation failed: {}", e)))?;
 
-    // Encrypted container (DataTypeId 1 = arbitrary custom)
+    // Encrypted container
     let inner_enc = DataValue::Container(vec![
-        (DataTypeId(1), DataValue::Str("secret inner data".into())),
-        (DataTypeId(2), DataValue::UnsignedNumber(42)),
+        (DataType::Version.to_id(&TypeMap::latest()), DataValue::Str("secret inner data".into())),
+        (DataType::Id.to_id(&TypeMap::latest()), DataValue::UnsignedNumber(42)),
     ]);
     let mut dv_enc = inner_enc;
     dv_enc
-        .encrypt_container(&cipher, b"demo-aad")
+        .encrypt_container(enc_type, &recipient, b"demo-aad")
         .ok_or_else(|| js_error("encryption failed"))?;
 
     // Signed container
     let inner_sig = DataValue::Container(vec![
-        (DataTypeId(1), DataValue::Str("signed by client".into())),
-        (DataTypeId(2), DataValue::UnsignedNumber(99)),
+        (DataType::Version.to_id(&TypeMap::latest()), DataValue::Str("signed by client".into())),
+        (DataType::Id.to_id(&TypeMap::latest()), DataValue::UnsignedNumber(99)),
     ]);
     let mut dv_sig = inner_sig;
     dv_sig
@@ -68,14 +75,20 @@ pub fn build_demo_message(client_id: u64, keyring_bytes: &[u8]) -> Result<Vec<u8
     // Signed + encrypted container
     let inner_sec = DataValue::Container(vec![
         (
-            DataTypeId(1),
+            DataType::Version.to_id(&TypeMap::latest()),
             DataValue::Str("signed+encrypted payload".into()),
         ),
-        (DataTypeId(2), DataValue::UnsignedNumber(7)),
+        (DataType::Id.to_id(&TypeMap::latest()), DataValue::UnsignedNumber(7)),
     ]);
     let mut dv_sec = inner_sec;
     dv_sec
-        .sign_and_encrypt_container(SigAlgorithm::ED25519, &signer, &cipher, b"demo-aad")
+        .sign_and_encrypt_container(
+            SigAlgorithm::ED25519,
+            &signer,
+            enc_type,
+            &recipient,
+            b"demo-aad",
+        )
         .ok_or_else(|| js_error("sign+encrypt failed"))?;
 
     let timestamp = js_sys::Date::now() as u64;
@@ -92,7 +105,8 @@ pub fn build_demo_message(client_id: u64, keyring_bytes: &[u8]) -> Result<Vec<u8
         .add_typed_default(DataType::Version, DataValue::Str("demo-wasm".into()))
         .with_sender(client_id);
 
-    Ok(msg.to_bytes())
+    msg.to_bytes()
+        .map_err(|e| js_error(&format!("encode failed: {}", e)))
 }
 
 /// Parse an auth response frame into a JS object.
@@ -101,24 +115,24 @@ pub fn parse_auth_response(response: &[u8]) -> Result<JsValue, JsValue> {
     let comm = CommunicationValue::from_bytes(response)
         .map_err(|e| js_error(&format!("parse failed: {}", e)))?;
 
-    let connected = matches!(comm.get_data(DataTypeId(11)), DataValue::BoolTrue);
+    let connected = matches!(comm.get_data(DataType::Connected.to_id(&TypeMap::latest())), DataValue::BoolTrue);
 
-    let client_nonce = match comm.get_data(DataTypeId(7)) {
+    let client_nonce = match comm.get_data(DataType::ClientNonce.to_id(&TypeMap::latest())) {
         DataValue::UnsignedNumber(n) => Some(*n),
         _ => None,
     };
 
-    let assigned_id = match comm.get_data(DataTypeId(6)) {
+    let assigned_id = match comm.get_data(DataType::Id.to_id(&TypeMap::latest())) {
         DataValue::UnsignedNumber(n) => Some(*n as u64),
         _ => None,
     };
 
-    let timestamp = match comm.get_data(DataTypeId(5)) {
+    let timestamp = match comm.get_data(DataType::Timestamp.to_id(&TypeMap::latest())) {
         DataValue::UnsignedNumber(n) => Some(*n),
         _ => None,
     };
 
-    let signature = match comm.get_data(DataTypeId(10)) {
+    let signature = match comm.get_data(DataType::Signature.to_id(&TypeMap::latest())) {
         DataValue::Bytes(b) => Some(b.clone()),
         _ => None,
     };
@@ -169,7 +183,8 @@ pub fn build_request_frame(comm_type: &str, id: u32, json_data: &str) -> Result<
     let frame = CommunicationValue::new(comm_type_enum)
         .with_id(id)
         .add_data(DataTypeId(32), DataValue::Str(json_data.to_string()))
-        .to_bytes();
+        .to_bytes()
+        .map_err(|e| js_error(&format!("encode failed: {}", e)))?;
 
     Ok(frame)
 }
@@ -235,17 +250,18 @@ mod tests {
 
     #[wasm_bindgen_test]
     fn build_ping_frame_roundtrip() {
-        let bytes = build_ping_frame(42, "test-ping", 1234567890, &[]);
+        let bytes = build_ping_frame(42, "test-ping", 1234567890, &[]).expect("encode failed");
         let cv = CommunicationValue::from_bytes(&bytes).expect("decode failed");
+        let tm = TypeMap::latest();
 
-        assert_eq!(cv.get_type(), CommunicationTypeId(19)); // Ping
+        assert_eq!(cv.get_type(), CommunicationType::Ping.to_id(&tm));
         assert_eq!(cv.get_sender(), 42);
         assert_eq!(
-            cv.get_data(DataTypeId(4)),
+            cv.get_data(DataType::Description.to_id(&tm)),
             &DataValue::Str("test-ping".into())
         );
         assert_eq!(
-            cv.get_data(DataTypeId(5)),
+            cv.get_data(DataType::Timestamp.to_id(&tm)),
             &DataValue::UnsignedNumber(1234567890)
         );
     }
@@ -253,59 +269,57 @@ mod tests {
     #[wasm_bindgen_test]
     fn build_ping_frame_with_data() {
         let payload = b"attachment-data";
-        let bytes = build_ping_frame(99, "with-data", 555, payload);
+        let bytes = build_ping_frame(99, "with-data", 555, payload).expect("encode failed");
         let cv = CommunicationValue::from_bytes(&bytes).expect("decode failed");
+        let tm = TypeMap::latest();
 
-        assert_eq!(cv.get_type(), CommunicationTypeId(19));
+        assert_eq!(cv.get_type(), CommunicationType::Ping.to_id(&tm));
         assert_eq!(cv.get_sender(), 99);
         assert_eq!(
-            cv.get_data(DataTypeId(4)),
+            cv.get_data(DataType::Description.to_id(&tm)),
             &DataValue::Str("with-data".into())
         );
-        assert_eq!(cv.get_data(DataTypeId(5)), &DataValue::UnsignedNumber(555));
+        assert_eq!(cv.get_data(DataType::Timestamp.to_id(&tm)), &DataValue::UnsignedNumber(555));
         assert_eq!(
-            cv.get_data(DataTypeId(6)),
+            cv.get_data(DataType::Id.to_id(&tm)),
             &DataValue::Bytes(payload.to_vec())
         );
     }
 
     #[wasm_bindgen_test]
     fn build_ping_frame_client_id_zero() {
-        let bytes = build_ping_frame(0, "zero-id", 0, &[]);
+        let bytes = build_ping_frame(0, "zero-id", 0, &[]).expect("encode failed");
         let cv = CommunicationValue::from_bytes(&bytes).expect("decode failed");
         assert_eq!(cv.get_sender(), 0);
     }
 
     #[wasm_bindgen_test]
     fn build_demo_message_roundtrip() {
-        let (_signer, sk, pk) = Ed25519Signer::generate();
-        let keyring = Keyring::new(
-            mtp_crypto::KemPublicKey::new(vec![]),
-            mtp_crypto::KemPrivateKey::new(vec![]),
-            mtp_crypto::SignaturePqPublicKey::new(vec![]),
-            mtp_crypto::SignaturePqPrivateKey::new(vec![]),
-            pk,
-            sk,
-        );
+        // The demo KEM-encrypts to the host's bundle, so a real host keypair is
+        // required; the client keyring only needs its Ed25519 signing key.
+        let keyring = Keyring::generate();
         let keyring_bytes = keyring.to_bytes();
+        let host_bundle = Keyring::generate().public_key_bundle().as_bytes();
 
-        let result = build_demo_message(7, &keyring_bytes);
+        let result = build_demo_message(7, &keyring_bytes, &host_bundle);
         assert!(result.is_ok());
 
         let bytes = result.unwrap();
         let cv = CommunicationValue::from_bytes(&bytes).expect("decode failed");
+        let tm = TypeMap::latest();
 
-        assert_eq!(cv.get_type(), CommunicationTypeId(19)); // Ping
+        assert_eq!(cv.get_type(), CommunicationType::Ping.to_id(&tm));
         assert_eq!(cv.get_sender(), 7);
         assert_eq!(
-            cv.get_data(DataTypeId(4)),
+            cv.get_data(DataType::Description.to_id(&tm)),
             &DataValue::Str("MTP WASM Demo".into())
         );
     }
 
     #[wasm_bindgen_test]
     fn build_demo_message_invalid_keyring() {
-        let result = build_demo_message(1, b"not-a-valid-keyring");
+        let host_bundle = Keyring::generate().public_key_bundle().as_bytes();
+        let result = build_demo_message(1, b"not-a-valid-keyring", &host_bundle);
         assert!(result.is_err());
         let err = result.unwrap_err();
         assert!(err.as_string().unwrap().contains("invalid keyring"));
@@ -318,7 +332,8 @@ mod tests {
             .add_typed_default(DataType::ClientNonce, DataValue::UnsignedNumber(999))
             .add_typed_default(DataType::Id, DataValue::UnsignedNumber(42))
             .add_typed_default(DataType::Timestamp, DataValue::UnsignedNumber(12345))
-            .to_bytes();
+            .to_bytes()
+            .expect("encode failed");
 
         let result = parse_auth_response(&resp).expect("parse failed");
 
@@ -337,7 +352,8 @@ mod tests {
     fn parse_auth_response_rejected() {
         let resp = CommunicationValue::new(CommunicationType::IdentificationResponse)
             .add_typed_default(DataType::Connected, DataValue::BoolFalse)
-            .to_bytes();
+            .to_bytes()
+            .expect("encode failed");
 
         let result = parse_auth_response(&resp).expect("parse failed");
 
@@ -357,7 +373,8 @@ mod tests {
         let resp = CommunicationValue::new(CommunicationType::IdentificationResponse)
             .add_typed_default(DataType::Connected, DataValue::BoolTrue)
             .add_typed_default(DataType::Signature, DataValue::Bytes(sig_bytes.clone()))
-            .to_bytes();
+            .to_bytes()
+            .expect("encode failed");
 
         let result = parse_auth_response(&resp).expect("parse failed");
 

wasm/src/transport.rs

@@ -1,3 +1,6 @@
+use std::cell::RefCell;
+use std::rc::Rc;
+
 use wasm_bindgen::JsCast;
 use wasm_bindgen::prelude::*;
 use wasm_bindgen_futures::JsFuture;
@@ -27,9 +30,65 @@ fn resolve_stream_readable(recv_stream: &JsValue) -> Result<JsValue, JsValue> {
     }
 }
 
+/// Outcome of reading the next framed message from the incoming stream(s).
+enum FrameOutcome {
+    /// A complete application frame.
+    Frame(Vec<u8>),
+    /// The peer sent an explicit close frame (length == `u32::MAX`).
+    Closed,
+    /// The incoming-streams readable ended (transport gone), no more frames.
+    Ended,
+}
+
+/*
+ * Debug-log the exact bytes about to be written to the WebTransport stream.
+ *
+ * Wire layout (note the DOUBLE length prefix):
+ *   [0..4]   outer_len  u32 BE  - added by send_frame (= inner frame length)
+ *   [4..8]   inner_len  u32 BE  - added by CommunicationValue::to_bytes
+ *   [8..10]  comm_type  u16 BE  - e.g. Identification
+ *   [10]     flags      u8
+ *   [11..]   id/sender/receiver/signature/data, gated by `flags`
+ */
+fn log_frame_bytes(wire: &[u8]) {
+    let hex: String = wire.iter().map(|b| format!("{b:02x}")).collect::<Vec<_>>().join(" ");
+
+    let outer_len = wire.get(0..4).map(|b| u32::from_be_bytes([b[0], b[1], b[2], b[3]]));
+    let inner_len = wire.get(4..8).map(|b| u32::from_be_bytes([b[0], b[1], b[2], b[3]]));
+    let comm_type = wire.get(8..10).map(|b| u16::from_be_bytes([b[0], b[1]]));
+    let flags = wire.get(10).copied();
+
+    web_sys::console::log_1(
+        &format!(
+            "mtp-wasm send_frame: {len} bytes | outer_len={outer_len:?} inner_len={inner_len:?} \
+             comm_type={comm_type:?} flags={flags:?}\n{hex}",
+            len = wire.len(),
+        )
+        .into(),
+    );
+}
+
+/*
+ * WebTransport client transport.
+ *
+ * The native host sends with a *persistent* uni-directional stream: the auth
+ * `Challenge` and the final `IdentificationResponse`/`RegisterResponse` arrive
+ * as two length-prefixed frames on the *same* QUIC stream, and later
+ * application messages may arrive on subsequent streams. The reader state
+ * (`streams_reader`, `stream_reader`, `buffer`) is therefore shared via `Rc`
+ * between the handshake (`read_one_frame`) and the background `receive_loop`,
+ * so frames are never lost across the boundary and multiple frames can be read
+ * from one stream.
+ */
 #[derive(Clone)]
 pub struct WasmTransport {
     inner: WebTransport,
+    /// Reader over `incoming_unidirectional_streams()` (a singleton stream of streams).
+    streams_reader: Rc<RefCell<Option<JsValue>>>,
+    /// Reader over the host's current uni-directional stream, if one is open.
+    stream_reader: Rc<RefCell<Option<JsValue>>>,
+    /// Bytes already read from the current stream but not yet consumed as a frame.
+    buffer: Rc<RefCell<Vec<u8>>>,
 }
 
 impl WasmTransport {
@@ -58,17 +117,18 @@ impl WasmTransport {
         JsFuture::from(transport.ready())
             .await
             .map_err(|e| js_error(&format!("WebTransport ready failed: {:?}", e)))?;
-        Ok(Self { inner: transport })
+        Ok(Self {
+            inner: transport,
+            streams_reader: Rc::new(RefCell::new(None)),
+            stream_reader: Rc::new(RefCell::new(None)),
+            buffer: Rc::new(RefCell::new(Vec::new())),
+        })
     }
 
     pub fn inner(&self) -> &WebTransport {
         &self.inner
     }
 
-    pub fn from_inner(inner: WebTransport) -> Self {
-        Self { inner }
-    }
-
     pub async fn send_frame(&self, frame: &[u8]) -> Result<(), JsValue> {
         let stream_promise = self.inner.create_unidirectional_stream();
         let stream = JsFuture::from(stream_promise).await?;
@@ -87,6 +147,8 @@ impl WasmTransport {
         wire.extend_from_slice(&len.to_be_bytes());
         wire.extend_from_slice(frame);
 
+        log_frame_bytes(&wire);
+
         let chunk = js_sys::Uint8Array::from(&wire[..]);
 
         let write_fn = js_sys::Reflect::get(&writer_val, &JsValue::from_str("write"))
@@ -110,241 +172,178 @@ impl WasmTransport {
         Ok(())
     }
 
-    /// Read exactly one frame from incoming uni streams, then release the reader
-    /// so `receive_loop` can pick up from where we left off.
-    pub async fn read_one_frame(&self) -> Result<Vec<u8>, JsValue> {
+    /// Get (creating once) the reader over `incoming_unidirectional_streams()`.
+    fn ensure_streams_reader(&self) -> Result<JsValue, JsValue> {
+        if let Some(reader) = self.streams_reader.borrow().clone() {
+            return Ok(reader);
+        }
         let incoming = self.inner.incoming_unidirectional_streams();
-
-        let reader_fn = js_sys::Reflect::get(&incoming, &JsValue::from_str("getReader"))
+        let reader = js_sys::Reflect::get(&incoming, &JsValue::from_str("getReader"))
             .map_err(|_| js_error("missing getReader"))?
             .dyn_into::<js_sys::Function>()
-            .map_err(|_| js_error("getReader not a function"))?;
-        let reader_val = reader_fn
+            .map_err(|_| js_error("getReader not a function"))?
             .call0(&incoming)
             .map_err(|_| js_error("getReader call failed"))?;
+        *self.streams_reader.borrow_mut() = Some(reader.clone());
+        Ok(reader)
+    }
 
-        let read_fn = js_sys::Reflect::get(&reader_val, &JsValue::from_str("read"))
+    /// Accept the next incoming uni-directional stream and make it current.
+    /// Returns `false` if the incoming-streams readable has ended.
+    async fn open_next_stream(&self) -> Result<bool, JsValue> {
+        let streams_reader = self.ensure_streams_reader()?;
+
+        let read_fn = js_sys::Reflect::get(&streams_reader, &JsValue::from_str("read"))
             .map_err(|_| js_error("missing read"))?
             .dyn_into::<js_sys::Function>()
             .map_err(|_| js_error("read not a function"))?;
-        let result_promise = read_fn
-            .call0(&reader_val)
-            .map_err(|_| js_error("read call failed"))?;
-        let result = JsFuture::from(result_promise.unchecked_into::<js_sys::Promise>())
+        let result = JsFuture::from(
+            read_fn
+                .call0(&streams_reader)
+                .map_err(|_| js_error("read call failed"))?
+                .unchecked_into::<js_sys::Promise>(),
+        )
         .await
-            .map_err(|e| js_error(&format!("read failed: {:?}", e)))?;
-
-        // Release the reader lock so receive_loop can create its own reader
-        if let Some(release_fn) =
-            js_sys::Reflect::get(&reader_val, &JsValue::from_str("releaseLock"))
-                .ok()
-                .and_then(|f| f.dyn_into::<js_sys::Function>().ok())
-        {
-            let _ = release_fn.call0(&reader_val);
-        }
+        .map_err(|e| js_error(&format!("accept stream failed: {:?}", e)))?;
 
         let done = js_sys::Reflect::get(&result, &JsValue::from_str("done"))
             .ok()
             .and_then(|v| v.as_bool())
             .unwrap_or(false);
         if done {
-            return Err(js_error("stream ended before frame"));
+            return Ok(false);
         }
 
         let recv_stream = js_sys::Reflect::get(&result, &JsValue::from_str("value"))
             .map_err(|_| js_error("missing value"))?;
-
-        let readable_or_stream = resolve_stream_readable(&recv_stream)?;
-        let stream_reader_fn =
-            js_sys::Reflect::get(&readable_or_stream, &JsValue::from_str("getReader"))
+        let readable = resolve_stream_readable(&recv_stream)?;
+        let reader = js_sys::Reflect::get(&readable, &JsValue::from_str("getReader"))
             .map_err(|_| js_error("missing stream getReader"))?
             .dyn_into::<js_sys::Function>()
-                .map_err(|_| js_error("stream getReader not a function"))?;
-        let stream_reader = stream_reader_fn
-            .call0(&readable_or_stream)
+            .map_err(|_| js_error("stream getReader not a function"))?
+            .call0(&readable)
             .map_err(|_| js_error("stream getReader call failed"))?;
 
-        let mut buffer: Vec<u8> = Vec::new();
-        loop {
-            let stream_read_fn =
-                match js_sys::Reflect::get(&stream_reader, &JsValue::from_str("read"))
-                    .ok()
-                    .and_then(|f| f.dyn_into::<js_sys::Function>().ok())
-                {
-                    Some(f) => f,
-                    None => break,
-                };
-            let chunk_promise = match stream_read_fn.call0(&stream_reader) {
-                Ok(p) => p,
-                Err(_) => break,
-            };
-            let chunk_result =
-                match JsFuture::from(chunk_promise.unchecked_into::<js_sys::Promise>()).await {
-                    Ok(v) => v,
-                    Err(_) => break,
-                };
-
-            let chunk_done = js_sys::Reflect::get(&chunk_result, &JsValue::from_str("done"))
-                .ok()
-                .and_then(|v| v.as_bool())
-                .unwrap_or(true);
-            if chunk_done {
-                return Err(js_error("stream closed before complete frame"));
+        *self.stream_reader.borrow_mut() = Some(reader);
+        Ok(true)
     }
 
-            if let Ok(chunk_val) = js_sys::Reflect::get(&chunk_result, &JsValue::from_str("value"))
-            {
-                let arr = js_sys::Uint8Array::new(&chunk_val).to_vec();
-                if !arr.is_empty() {
-                    buffer.extend_from_slice(&arr);
-                }
-            }
-
-            if buffer.len() >= 4 {
-                let frame_len = u32::from_be_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
-                if frame_len == CLOSE_FRAME_LEN {
-                    return Err(js_error("connection closed before frame"));
-                }
-
-                let frame_len = frame_len as usize;
-                let Some(frame_end) = 4usize.checked_add(frame_len) else {
-                    return Err(js_error("invalid frame length"));
-                };
-                if frame_end <= buffer.len() {
-                    return Ok(buffer[4..frame_end].to_vec());
-                }
-            }
-        }
-
-        Err(js_error("stream ended before frame complete"))
-    }
-
-    pub async fn receive_loop(&self, on_message: js_sys::Function, on_error: js_sys::Function) {
-        let incoming = self.inner.incoming_unidirectional_streams();
-
-        let reader_fn = match js_sys::Reflect::get(&incoming, &JsValue::from_str("getReader")) {
-            Ok(f) => f.dyn_into::<js_sys::Function>().unwrap(),
-            Err(_) => return,
-        };
-        let reader_val = match reader_fn.call0(&incoming) {
-            Ok(v) => v,
-            Err(_) => return,
+    /// Read one chunk from the current stream. `Ok(None)` means the stream ended.
+    async fn read_chunk(&self) -> Result<Option<Vec<u8>>, JsValue> {
+        let reader = match self.stream_reader.borrow().clone() {
+            Some(r) => r,
+            None => return Ok(None),
         };
 
-        loop {
-            let read_fn = match js_sys::Reflect::get(&reader_val, &JsValue::from_str("read")) {
-                Ok(f) => f.dyn_into::<js_sys::Function>().unwrap(),
-                Err(_) => break,
-            };
-            let result = match read_fn.call0(&reader_val) {
-                Ok(p) => match JsFuture::from(p.unchecked_into::<js_sys::Promise>()).await {
-                    Ok(v) => v,
-                    Err(_) => break,
-                },
-                Err(_) => break,
-            };
+        let read_fn = js_sys::Reflect::get(&reader, &JsValue::from_str("read"))
+            .map_err(|_| js_error("missing read"))?
+            .dyn_into::<js_sys::Function>()
+            .map_err(|_| js_error("read not a function"))?;
+        let result = JsFuture::from(
+            read_fn
+                .call0(&reader)
+                .map_err(|_| js_error("read call failed"))?
+                .unchecked_into::<js_sys::Promise>(),
+        )
+        .await
+        .map_err(|e| js_error(&format!("read failed: {:?}", e)))?;
 
         let done = js_sys::Reflect::get(&result, &JsValue::from_str("done"))
-                .ok()
-                .and_then(|v| v.as_bool())
-                .unwrap_or(false);
-            if done {
-                break;
-            }
-
-            let recv_stream = match js_sys::Reflect::get(&result, &JsValue::from_str("value")) {
-                Ok(v) => v,
-                Err(_) => continue,
-            };
-
-            let this = self.clone();
-            let on_msg = on_message.clone();
-            let on_err = on_error.clone();
-            wasm_bindgen_futures::spawn_local(async move {
-                let _ = this.handle_stream(recv_stream, on_msg, on_err).await;
-            });
-        }
-    }
-
-    async fn handle_stream(
-        &self,
-        recv_stream: JsValue,
-        on_message: js_sys::Function,
-        _on_error: js_sys::Function,
-    ) -> Result<(), JsValue> {
-        let readable_or_stream = resolve_stream_readable(&recv_stream)?;
-        let stream_reader_fn =
-            js_sys::Reflect::get(&readable_or_stream, &JsValue::from_str("getReader"))?
-                .dyn_into::<js_sys::Function>()?;
-        let stream_reader = stream_reader_fn.call0(&readable_or_stream)?;
-
-        let mut buffer: Vec<u8> = Vec::new();
-
-        loop {
-            let stream_read_fn = js_sys::Reflect::get(&stream_reader, &JsValue::from_str("read"))?
-                .dyn_into::<js_sys::Function>()?;
-            let chunk_promise = stream_read_fn.call0(&stream_reader)?;
-            let chunk_result =
-                JsFuture::from(chunk_promise.unchecked_into::<js_sys::Promise>()).await?;
-
-            let chunk_done = js_sys::Reflect::get(&chunk_result, &JsValue::from_str("done"))
             .ok()
             .and_then(|v| v.as_bool())
             .unwrap_or(true);
-            if chunk_done {
-                break;
+        if done {
+            return Ok(None);
         }
 
-            if let Ok(chunk_val) = js_sys::Reflect::get(&chunk_result, &JsValue::from_str("value"))
-            {
-                let arr = js_sys::Uint8Array::new(&chunk_val).to_vec();
-                if !arr.is_empty() {
-                    buffer.extend_from_slice(&arr);
-                }
+        let value = js_sys::Reflect::get(&result, &JsValue::from_str("value"))
+            .map_err(|_| js_error("missing value"))?;
+        Ok(Some(js_sys::Uint8Array::new(&value).to_vec()))
     }
 
-            // Extract all complete frames from the buffer
-            while buffer.len() >= 4 {
-                let frame_len = u32::from_be_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
+    /// Try to pull one complete frame out of the buffer without reading more.
+    fn parse_buffer(&self) -> Result<Option<FrameOutcome>, JsValue> {
+        let buf = self.buffer.borrow();
+        if buf.len() < 4 {
+            return Ok(None);
+        }
+        let frame_len = u32::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]);
         if frame_len == CLOSE_FRAME_LEN {
-                    return Ok(());
+            return Ok(Some(FrameOutcome::Closed));
         }
-
         let frame_len = frame_len as usize;
         let Some(frame_end) = 4usize.checked_add(frame_len) else {
             return Err(js_error("invalid frame length"));
         };
-                if frame_end > buffer.len() {
-                    break;
+        if frame_end > buf.len() {
+            return Ok(None);
         }
-                let frame = buffer[4..frame_end].to_vec();
+        let frame = buf[4..frame_end].to_vec();
+        drop(buf);
+        self.buffer.borrow_mut().drain(..frame_end);
+        Ok(Some(FrameOutcome::Frame(frame)))
+    }
+
+    /*
+     * Read the next framed message from the host. Frames are length-prefixed
+     * (u32 big-endian) and may be packed several-per-stream (the host reuses a
+     * persistent uni stream) or one-per-stream; both are handled by buffering
+     * across reads and advancing to the next stream when the current one ends.
+     */
+    async fn next_frame(&self) -> Result<FrameOutcome, JsValue> {
+        loop {
+            if let Some(outcome) = self.parse_buffer()? {
+                return Ok(outcome);
+            }
+
+            let have_stream = self.stream_reader.borrow().is_some();
+            if !have_stream && !self.open_next_stream().await? {
+                return Ok(FrameOutcome::Ended);
+            }
+
+            match self.read_chunk().await? {
+                Some(chunk) => {
+                    if !chunk.is_empty() {
+                        self.buffer.borrow_mut().extend_from_slice(&chunk);
+                    }
+                }
+                None => {
+                    // Current stream finished; the next frame (if any) is on a
+                    // subsequent stream. Any trailing partial bytes are dropped
+                    // since the host never splits a frame across streams.
+                    *self.stream_reader.borrow_mut() = None;
+                    self.buffer.borrow_mut().clear();
+                }
+            }
+        }
+    }
+
+    /// Read exactly one application frame (used during the auth handshake).
+    pub async fn read_one_frame(&self) -> Result<Vec<u8>, JsValue> {
+        match self.next_frame().await? {
+            FrameOutcome::Frame(frame) => Ok(frame),
+            FrameOutcome::Closed => Err(js_error("connection closed before frame")),
+            FrameOutcome::Ended => Err(js_error("stream ended before frame")),
+        }
+    }
+
+    /// Background loop: deliver every incoming frame to `on_message` until the
+    /// connection closes. Shares reader state with `read_one_frame`, so frames
+    /// buffered during the handshake are not lost.
+    pub async fn receive_loop(&self, on_message: js_sys::Function, on_error: js_sys::Function) {
+        loop {
+            match self.next_frame().await {
+                Ok(FrameOutcome::Frame(frame)) => {
                     let arr = js_sys::Uint8Array::from(&frame[..]);
                     let _ = on_message.call1(&JsValue::NULL, &arr);
-                buffer.drain(..frame_end);
                 }
-        }
-
-        // Process any remaining complete frames after stream closes
-        while buffer.len() >= 4 {
-            let frame_len = u32::from_be_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
-            if frame_len == CLOSE_FRAME_LEN {
-                return Ok(());
-            }
-
-            let frame_len = frame_len as usize;
-            let Some(frame_end) = 4usize.checked_add(frame_len) else {
-                return Err(js_error("invalid frame length"));
-            };
-            if frame_end > buffer.len() {
+                Ok(FrameOutcome::Closed) | Ok(FrameOutcome::Ended) => break,
+                Err(e) => {
+                    let _ = on_error.call1(&JsValue::NULL, &e);
                     break;
                 }
-            let frame = buffer[4..frame_end].to_vec();
-            let arr = js_sys::Uint8Array::from(&frame[..]);
-            let _ = on_message.call1(&JsValue::NULL, &arr);
-            buffer.drain(..frame_end);
             }
-
-        Ok(())
+        }
     }
 
     pub fn close(&self) {