A lot

CONNECTOR.md

@@ -1,13 +1,13 @@
 # Connector
 
-This file documents the Connection and Version Negotiation logic.
+This file documents the connection and version negotiation logic.
 
 ## Registry
 
-The `registry` crate provides a multi-version `Registry` used by the host for version negotiation. A `Registry` holds one `TypeMap` per protocol version and supports `negotiate()`:
+The `registry` module provides a multi-version `Registry` used by the host for version negotiation. Accessed through the `mtp` facade (requires the `host` feature):
 
 ```rust
-use registry::Registry;
+use mtp::codec::registry::Registry;
 
 let registry = Registry::builtin(); // loads all TypeMaps from config
 
@@ -33,11 +33,12 @@ The host creates a QUIC server, manages the registry, and handles version negoti
 
 ### Initialization
 
+The host binds to the address from the `mtp_BIND` environment variable (defaults to `::`) on the specified port:
+
 ```rust
-use mtp_host::{MTPHost, HostConfig};
+use mtp::host::{MTPHost, HostConfig};
 
 let config = HostConfig {
-    ip: "::".into(),
     port: 4433,
     tls_fullchain: std::fs::read("cert.pem")?,
     tls_key: std::fs::read("key.pem")?,
@@ -50,35 +51,31 @@ let mut host = MTPHost::new(config).await?;
 
 ```rust
 while let Some(conn) = host.accept().await {
-    /*
+    // conn.version is the negotiated version
-     *  conn.version is the negotiated version
+    // conn.codec is a VersionedCodec scoped to that version
-     * conn.codec is a VersionedCodec scoped to that version
+    // conn.sender / conn.receiver for raw CommunicationValue I/O
-     * conn.sender / conn.receiver for raw CommunicationValue I/O
-     */
 
     let msg = conn.receiver.receive().await?;
-    // ...
 }
 ```
 
 The host's `accept()` method:
 1. Accepts a QUIC connection
-2. Reads the first `CommunicationValue` (always encoded with reserved type IDs)
+2. If authentication is required (crypto feature): performs login/register handshake
-3. Extracts the client's protocol version from `DataType::Version` (wire ID 3)
+3. Reads the first `CommunicationValue` (always encoded with reserved type IDs)
-4. Calls `registry.negotiate(&[client_version])`
+4. Extracts the client's protocol version from `DataType::Version` (wire ID 3)
-5. Returns `None` if the version is unsupported (caller sends `ErrorBadVersion` and disconnects)
+5. Calls `registry.negotiate(&[client_version])`
-6. Returns an `MTPConnection` with the negotiated version otherwise
+6. Returns `None` if the version is unsupported
+7. Returns an `MTPConnection` with the negotiated version otherwise
 
-### Login/Register Flow
+### Login/Register Handshake (crypto feature)
 
-The complete login/register handshake (see design docs) builds on top of `MTPConnection`:
+When `require_authentication` is set, the host sends a **greeting** first (host ID, public keys, nonce). The client then responds with either:
 
-1. Client sends `Identification` with version, ID, nonce, signature
+- **Login** (`CommunicationType::Identification`, ID 15): client ID, nonce, signature
-2. Host verifies signature via `get_key` callback
+- **Register** (`CommunicationType::Register`, ID 17): public keys, nonce, signature
-3. Host responds with approval + nonces + signature
-4. Client verifies response
 
-New clients use the `Register` variant instead, presenting their public key for registration.
+The host verifies the client's signature, sends a signed response, and the client verifies the host's signature.
 
 ---
 
@@ -87,23 +84,24 @@ New clients use the `Register` variant instead, presenting their public key for
 The client connects to a host and uses a single compiled-in protocol version.
 
 ```rust
-use mtp_client::{MTPClient, ClientConfig};
+use mtp::client::{MTPClient, ClientConfig};
 
 let config = ClientConfig {
     url: "https://host.example.com:4433".into(),
     server_cert: None, // or Some(cert_pem_bytes)
 };
 
-// Connect (existing client)
+// Connect (unauthenticated, existing client)
 let conn = MTPClient::connect(config, 8765).await?;
 
-/*
+// Authenticated login
- *  conn.version is the compiled-in PROTOCOL_VERSION
+let conn = MTPClient::auth_connect(config, 8765, keys, host_pk).await?;
- *  conn.sender / conn.receiver for I/O
+
- */
+// Registration (new client)
+let conn = MTPClient::auth_register(config, keys, host_pk).await?;
 ```
 
-The client's `PROTOCOL_VERSION` constant is set by `protocol_version` in `type-maps.yaml` and baked in at compile time. The client never imports the `registry` crate; it only uses `type-map` for enum types and `codec` for encoding.
+The client's `PROTOCOL_VERSION` constant is set by `protocol_version` in `type-maps.yaml` and baked in at compile time. The client never imports the `registry` crate; it only uses `mtp::type_map` for enum types and `mtp::codec` for encoding.
 
 ---
 
@@ -111,26 +109,26 @@ The client's `PROTOCOL_VERSION` constant is set by `protocol_version` in `type-m
 
 ```
 Client (v2.0)         Host (v0.0, v1.0, v2.0)
-    │                        │
+    |                        |
-    │  QUIC connect          │
+    |  QUIC connect          |
-    │───────────────────────→│
+    |----------------------->|
-    │                        │
+    |                        |
-    │  CommValue{ Ident. }   │
+    |  CommValue{ Ident. }   |
-    │    Version → "2.0"     │
+    |    Version -> "2.0"    |
-    │    Id → 8765           │
+    |    Id -> 8765          |
-    │    Nonce → ...         │
+    |    Nonce -> ...        |
-    │    Signature → ...     │
+    |    Signature -> ...    |
-    │───────────────────────→│
+    |----------------------->|
-    │                        │  registry.negotiate(&[Version(2,0)])
+    |                        |  registry.negotiate(&[Version(2,0)])
-    │                        │  → Some(Version(2,0))
+    |                        |  -> Some(Version(2,0))
-    │                        │
+    |                        |
-    │  Response              │
+    |  Response              |
-    │←───────────────────────│  (uses v2.0 TypeMap for encoding)
+    |<-----------------------|  (uses v2.0 TypeMap for encoding)
-    │    Status, Nonces,     │
+    |    Status, Nonces,     |
-    │    Signature           │
+    |    Signature           |
-    │                        │
+    |                        |
-    │  (subsequent messages  │
+    |  (subsequent messages  |
-    │   use v2.0 TypeMap)    │
+    |   use v2.0 TypeMap)    |
 ```
 
-If the client sends an unsupported version (e.g. v3.0 when the host only knows up to v2.0), `negotiate` returns `None` and the host sends `ErrorBadVersion` using reserved types.
+If the client sends an unsupported version (e.g. v3.0 when the host only knows up to v2.0), `negotiate` returns `None` and the connection is closed.

Cargo.toml

@@ -2,7 +2,6 @@
 members = [
     "common",
     "crypto",
-    "registry",
     "type-map",
     "codec",
     "transport",
@@ -10,3 +9,54 @@ members = [
     "client",
 ]
 resolver = "3"
+
+# =============================================================================
+# mtp - umbrella crate
+#
+# Re-exports all MTP sub-crates behind feature flags so consumers write:
+#
+#   [dependencies]
+#   mtp = { version = "0.1", features = ["client", "crypto"] }
+#
+# instead of listing every sub-crate individually.
+# =============================================================================
+[package]
+name = "mtp"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+# --- always-on core ---
+mtp-common = { path = "common" }
+mtp-type-map = { path = "type-map" }
+mtp-codec = { path = "codec" }
+mtp-transport = { path = "transport" }
+
+# --- optional, behind features ---
+mtp-crypto = { path = "crypto", optional = true, features = [
+    "serde",
+    "mlkem-tls",
+] }
+mtp-host = { path = "host", optional = true }
+mtp-client = { path = "client", optional = true }
+
+[features]
+default = []
+
+# Serialization
+serde = ["mtp-crypto/serde"]
+
+# Message encryption (AEAD, signatures, KEM, KDF, hashing).
+# When combined with `host` or `client`, also enables connection authentication.
+crypto = [
+    "dep:mtp-crypto",
+    "mtp-codec/crypto",
+    "mtp-host?/crypto",
+    "mtp-client?/crypto",
+]
+
+# MTP server host - version negotiation, Registry, incoming QUIC connections.
+host = ["dep:mtp-host", "mtp-codec/registry", "mtp-transport/host"]
+
+# MTP client - outgoing QUIC connections to a host.
+client = ["dep:mtp-client"]

README.md

@@ -1,47 +1,80 @@
-# **M**ethanium **T**ransport **P**rotocol
+# Methanium Transport Protocol
 
-**MTP** is a **m**odular **t**ransport **p**rotocol by Methanium.
+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.
 
-## Crates
+## Getting Started
-- Codec
-- Transport
-- Common
-- Crypto
-- TypeMap
-- Registry
-- Host
-- Client
 
-## Codec
+Add the `mtp` crate with your desired features:
 
-The Codec crate handles encoding and decoding of MTP packets using Communication Types and Data Types from the Registry.
+```toml
+[dependencies]
+mtp = { path = "..", features = ["client", "crypto"] }
+```
 
-**Data Values:**
+The `mtp` umbrella crate re-exports all sub-crates behind feature flags:
-- Container
+
-- Encrypted Container (requires `crypto` feature)
+| Feature  | Pulls in                    | Enables                                   |
-- Signed Integer
+| -------- | --------------------------- | ----------------------------------------- |
-- Unsigned Integer
+| `crypto` | `mtp::crypto`               | AEAD, signatures, KEM, KDF, hashing       |
+| `host`   | `mtp::host`, `mtp::codec::registry` | QUIC server, version negotiation  |
+| `client` | `mtp::client`               | QUIC client connections                   |
+
+Core crates (`codec`, `transport`, `common`, `type_map`) are always available.
+
+```rust
+use mtp::codec::{CommunicationValue, DataValue};
+use mtp::type_map::{CommunicationType, DataType, TypeMap};
+use mtp::transport::{Sender, Receiver};
+
+#[cfg(feature = "crypto")]
+use mtp::crypto::ChaCha20Poly1305;
+```
+
+## Sub-crates
+
+All sub-crates are re-exported through the `mtp` facade and can be referenced as `mtp::codec`, `mtp::transport`, `mtp::common`, `mtp::type_map`, `mtp::crypto`, `mtp::host`, `mtp::client`.
+
+### Codec
+
+The codec crate handles binary encoding and decoding of MTP packets using Communication Types and Data Types resolved through the type-map registry.
+
+**Data Value types:**
+- Container (key-value map of typed entries)
+- Encrypted Container (requires `crypto`)
+- Signed Container (requires `crypto`)
+- SignedEncrypted Container (requires `crypto`)
+- Signed Integer (i128)
+- Unsigned Integer (u128)
 - Boolean
-- Signed Float
+- Float (exponent + mantissa)
 - String
 - Array
-- Binary (List of Bytes)
+- Bytes
+- Null
 
-Encoding/decoding uses a `TypeMap` to resolve type names to wire IDs.
+Encoding and decoding use a `TypeMap` to resolve type names to wire IDs. The `CommunicationValue` struct provides the frame format (type, flags, optional id/sender/receiver, data payload, optional signature).
 
-The Codec uses the Crypto crate to encrypt and decrypt Encrypted Containers.
+### Transport
 
----
+The transport crate wraps QUIC using `wtransport`. It provides `Sender`/`Receiver` for bidirectional message passing over uni-directional QUIC streams. Supports two send modes: persistent stream and single-stream-per-message.
 
-## Transport
+### Host
 
-The Transport crate wraps QUIC using `wtransport`. It provides `Sender`/`Receiver` for bidirectional message passing over QUIC streams.
+The host crate provides `MTPHost` with built-in version negotiation and optional authenticated login/registration (requires `crypto`). Accepts connections, negotiates protocol version, and returns `MTPConnection` handles.
 
-## Common
+### Client
 
-Common handles logging and error handling. Personal information reported to Common will be anonymized.
+The client crate provides `MTPClient` that connects to an MTP host. Supports `connect` (unauthenticated), `auth_connect` (login), and `auth_register` (registration) when built with `crypto`.
 
-## Crypto Stack
+### Common
+
+Common defines shared error types (`CodecError`, `CommunicationError`) used across all crates.
+
+### Type Map
+
+The type-map build script reads a YAML configuration to generate `CommunicationType` and `DataType` enums at compile time. The runtime crate provides `TypeMap`, `Version`, `CommunicationTypeId`, `DataTypeId`, and the multi-version `Registry` (requires `registry` feature).
+
+### Crypto Stack
 
 | Crate            | Audited? | Notes                                            |
 | ---------------- | -------- | ------------------------------------------------ |
@@ -53,3 +86,16 @@ Common handles logging and error handling. Personal information reported to Comm
 | sha2             | No       | Standard construction; widely reviewed           |
 | zeroize          | No       | Simple; widely used                              |
 | mlkem-tls        | No       | mlkem-rs backend unaudited                       |
+
+## Test Coverage
+
+| Crate          | Tests | Notes                                      |
+| -------------- | ----- | ------------------------------------------ |
+| common         | 7     | Error type Display, Clone, format          |
+| type-map       | 3     | Registry builtin, negotiate, latest        |
+| codec          | 32    | DataValue roundtrip, accessors, Display, Hash, base64; CommunicationValue frame encode/decode |
+| crypto         | 20    | AEAD, signatures, KEM, KDF, hash, key types, multi-encrypt |
+| transport      | 12    | ConnectionHandle state transitions; Policy defaults; SendMode |
+| host           | 3     | Version extraction from messages           |
+| client         | 7     | ClientConfig; AuthState derives            |
+| **Total**      | **84** |                                            |

TYPE-MAP.md

@@ -17,7 +17,7 @@ An example `type-maps.yaml` is provided in the [`example-type-maps.yaml`](./exam
 After editing the config and rebuilding, `CommunicationType` and `DataType` enums are generated automatically. Use them in code:
 
 ```rust
-use mtp_type_map::{CommunicationType, DataType, TypeMap};
+use mtp::type_map::{CommunicationType, DataType, TypeMap};
 
 let tm = TypeMap::v2_0();
 let id = tm.data_id_enum(DataType::SomeType).unwrap();
@@ -28,8 +28,8 @@ The enums are a **union across all versions**; every type name from every versio
 Encoding/decoding uses a `TypeMap` to resolve type names to wire IDs:
 
 ```rust
-use mtp_codec::{encode, decode, DataValue};
+use mtp::codec::{encode, decode, DataValue};
-use mtp_type_map::TypeMap;
+use mtp::type_map::TypeMap;
 
 let tm = TypeMap::v2_0();
 let value = DataValue::Str("hello".into());
@@ -62,8 +62,15 @@ This is by design: the host maps unknown types to `Error`, and the client should
 
 The `registry` feature of the Codec crate adds `VersionedCodec` for version-aware encoding:
 
+Requires the `host` feature (which enables `mtp-codec`'s `registry` feature):
+
+```toml
+[dependencies]
+mtp = { path = "..", features = ["host"] }
+```
+
 ```rust
-use mtp_codec::registry::{Registry, VersionedCodec};
+use mtp::codec::registry::{Registry, VersionedCodec};
 
 let registry = Registry::builtin();
 let codec = VersionedCodec::new(registry);

client/Cargo.toml

@@ -7,4 +7,9 @@ edition = "2024"
 mtp-common = { path = "../common" }
 mtp-codec = { path = "../codec" }
 mtp-transport = { path = "../transport" }
+mtp-crypto = { path = "../crypto", optional = true }
 tokio = { version = "1", features = ["full"] }
+rand = "0.8"
+
+[features]
+crypto = ["dep:mtp-crypto", "mtp-codec/crypto"]

client/src/lib.rs

@@ -1,10 +1,13 @@
-use mtp_codec::{CommunicationValue, DataTypeId, DataValue, Version, PROTOCOL_VERSION};
+#[cfg(feature = "crypto")]
+use mtp_codec::DataType;
+use mtp_codec::{CommunicationValue, DataTypeId, DataValue, PROTOCOL_VERSION, Version};
 use mtp_common::CommunicationError;
 use mtp_transport::{Policy, Receiver, Sender};
 
 pub struct ClientConfig {
     pub url: String,
     pub server_cert: Option<Vec<u8>>,
+    pub client_id: u64,
 }
 
 // Established MTP connection with a single negotiated version.
@@ -12,6 +15,18 @@ pub struct MTPConnection {
     pub version: Version,
     pub sender: Sender,
     pub receiver: Receiver,
+    #[cfg(feature = "crypto")]
+    pub auth_state: AuthState,
+    #[cfg(feature = "crypto")]
+    pub client_id: u64,
+}
+
+#[cfg(feature = "crypto")]
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum AuthState {
+    Pending,
+    Authenticated,
+    Failed,
 }
 
 pub struct MTPClient;
@@ -23,10 +38,7 @@ impl MTPClient {
      * The first message includes the client's protocol version
      * (a reserved `Version` data entry) so the host can negotiate.
      */
-    pub async fn connect(
+    pub async fn connect(config: ClientConfig) -> Result<MTPConnection, CommunicationError> {
-        config: ClientConfig,
-        client_id: u64,
-    ) -> Result<MTPConnection, CommunicationError> {
         let (sender, receiver) =
             mtp_transport::connect(&config.url, config.server_cert, Policy::default()).await?;
 
@@ -34,7 +46,10 @@ impl MTPClient {
         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), DataValue::UnsignedNumber(client_id.into()));
+            .add_data(
+                DataTypeId(6),
+                DataValue::UnsignedNumber(config.client_id.into()),
+            );
 
         sender.send(&ident).await?;
 
@@ -42,6 +57,338 @@ impl MTPClient {
             version: PROTOCOL_VERSION,
             sender,
             receiver,
+            #[cfg(feature = "crypto")]
+            auth_state: AuthState::Authenticated,
+            #[cfg(feature = "crypto")]
+            client_id: config.client_id,
         })
     }
 }
+
+/* ===== Authentication ===== */
+#[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,
+        };
+
+        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 connected = response.get_data(DataTypeId(11));
+        match connected {
+            DataValue::BoolTrue => {}
+            DataValue::BoolFalse => {
+                return Err(CommunicationError::AuthenticationFailed(
+                    "Server rejected authentication".into(),
+                ));
+            }
+            _ => {
+                return Err(CommunicationError::AuthenticationFailed(
+                    "Invalid response".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(&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,
+        )
+        .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 ===== */
+
+        Ok(MTPConnection {
+            version: PROTOCOL_VERSION,
+            sender,
+            receiver,
+            auth_state: AuthState::Authenticated,
+            client_id: config.client_id,
+        })
+    }
+
+    pub async fn auth_register(
+        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,
+        };
+
+        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 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);
+
+        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 register = CommunicationValue::new(mtp_codec::CommunicationType::Register)
+            .add_typed_default(DataType::Version, DataValue::Str(version_str))
+            .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)
+                .map_err(|e| CommunicationError::Other(e.to_string()))?;
+            register =
+                register.add_typed_default(DataType::PqSignature, DataValue::Bytes(pq_signature));
+        }
+        /* ===== End Signature ===== */
+
+        sender.send(&register).await?;
+
+        // 2. Receive host response (single message)
+        let response = receiver.receive().await?;
+
+        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,
+            _ => {
+                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 ===== */
+
+        Ok(MTPConnection {
+            version: PROTOCOL_VERSION,
+            sender,
+            receiver,
+            auth_state: AuthState::Authenticated,
+            client_id: assigned_id as u64,
+        })
+    }
+}
+
+/* ================================ TESTS ================================ */
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_client_config_url() {
+        let config = ClientConfig {
+            url: "https://example.com:4433".into(),
+            server_cert: None,
+            client_id: 0,
+        };
+        assert_eq!(config.url, "https://example.com:4433");
+        assert!(config.server_cert.is_none());
+    }
+
+    #[test]
+    fn test_client_config_with_cert() {
+        let config = ClientConfig {
+            url: "https://localhost:4433".into(),
+            server_cert: Some(vec![0x01, 0x02, 0x03]),
+            client_id: 42,
+        };
+        assert_eq!(config.server_cert, Some(vec![0x01, 0x02, 0x03]));
+        assert_eq!(config.client_id, 42);
+    }
+
+    #[cfg(feature = "crypto")]
+    #[test]
+    fn test_auth_state_derive() {
+        assert_eq!(AuthState::Pending, AuthState::Pending);
+        assert_ne!(AuthState::Authenticated, AuthState::Failed);
+    }
+}

codec/Cargo.toml

@@ -7,12 +7,11 @@ edition = "2024"
 mtp-type-map = { path = "../type-map" }
 mtp-common = { path = "../common" }
 mtp-crypto = { path = "../crypto", optional = true }
-registry = { path = "../registry", optional = true }
+base64 = "0.22"
-base64 = "*"
+byteorder = "1.5"
-byteorder = "*"
+rand = { version = "0.8", features = ["std", "std_rng"] }
-rand = { version = "*", features = ["std", "std_rng"] }
 
 [features]
 default = []
-registry = ["dep:registry", "mtp-type-map/multi-version"]
+registry = ["mtp-type-map/registry"]
 crypto = ["dep:mtp-crypto"]

codec/src/communication_value.rs

@@ -5,11 +5,23 @@ use std::io::{Cursor, Read};
 
 use crate::data_value::DataValue;
 use crate::rand_u32;
+use mtp_common::CodecError;
 use mtp_type_map::{
-    communication_type_name, data_type_name, CommunicationType, CommunicationTypeId, DataType,
+    CommunicationType, CommunicationTypeId, DataType, DataTypeId, PROTOCOL_VERSION, TypeMap,
-    DataTypeId, TypeMap, PROTOCOL_VERSION,
+    communication_type_name, data_type_name,
 };
 
+#[cfg(feature = "crypto")]
+use mtp_crypto::SigAlgorithm;
+#[cfg(feature = "crypto")]
+use mtp_crypto::SignatureScheme;
+
+const FLAG_HAS_SENDER: u8 = 0b0000_0001;
+const FLAG_HAS_RECEIVER: u8 = 0b0000_0010;
+const FLAG_HAS_ID: u8 = 0b0000_0100;
+const FLAG_ENCRYPTED: u8 = 0b0000_1000;
+const FLAG_SIGNED: u8 = 0b0001_0000;
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct CommunicationValue {
     id: u32,
@@ -18,10 +30,12 @@ pub struct CommunicationValue {
     receiver: u64,
     data: BTreeMap<DataTypeId, DataValue>,
     type_map: Option<TypeMap>,
+    #[cfg(feature = "crypto")]
+    frame_signature: Option<(u8, Vec<u8>)>,
 }
 
-#[allow(dead_code)]
 impl CommunicationValue {
+    #[must_use]
     pub fn new(comm_type: CommunicationType) -> Self {
         let tm = TypeMap::new(PROTOCOL_VERSION);
         let id = comm_type.to_id(&tm);
@@ -32,10 +46,13 @@ impl CommunicationValue {
             receiver: 0,
             data: BTreeMap::new(),
             type_map: Some(tm),
+            #[cfg(feature = "crypto")]
+            frame_signature: None,
         }
     }
 
     #[cfg(feature = "registry")]
+    #[must_use]
     pub fn from_comm(comm_type: CommunicationType, tm: &TypeMap) -> Self {
         let id = comm_type.to_id(tm);
         Self {
@@ -45,22 +62,12 @@ impl CommunicationValue {
             receiver: 0,
             data: BTreeMap::new(),
             type_map: Some(tm.clone()),
+            #[cfg(feature = "crypto")]
+            frame_signature: None,
         }
     }
 
-    pub fn from_comm_default(comm_type: CommunicationType) -> Self {
+    #[must_use]
-        let tm = TypeMap::new(PROTOCOL_VERSION);
-        let id = comm_type.to_id(&tm);
-        Self {
-            id: rand_u32(),
-            comm_type: id,
-            sender: 0,
-            receiver: 0,
-            data: BTreeMap::new(),
-            type_map: Some(tm),
-        }
-    }
-
     pub fn with_id(mut self, p0: u32) -> Self {
         self.id = p0;
         self
@@ -70,6 +77,7 @@ impl CommunicationValue {
         self.id
     }
 
+    #[must_use]
     pub fn with_sender(mut self, sender: u64) -> Self {
         self.sender = sender;
         self
@@ -79,6 +87,7 @@ impl CommunicationValue {
         self.sender
     }
 
+    #[must_use]
     pub fn with_receiver(mut self, receiver: u64) -> Self {
         self.receiver = receiver;
         self
@@ -92,21 +101,20 @@ impl CommunicationValue {
         self.comm_type
     }
 
-    pub fn is_type(&self, p0: CommunicationTypeId) -> bool {
+    #[must_use]
-        self.comm_type == p0
-    }
-
     pub fn add_data(mut self, data: DataTypeId, value: DataValue) -> Self {
         self.data.insert(data, value);
         self
     }
 
     #[cfg(feature = "registry")]
+    #[must_use]
     pub fn add_typed(mut self, data: DataType, tm: &TypeMap, value: DataValue) -> Self {
         self.data.insert(data.to_id(tm), value);
         self
     }
 
+    #[must_use]
     pub fn add_typed_default(mut self, data: DataType, value: DataValue) -> Self {
         let tm = self.type_map.clone().unwrap_or_else(TypeMap::latest);
         self.data.insert(data.to_id(&tm), value);
@@ -116,10 +124,6 @@ impl CommunicationValue {
     pub fn get_data(&self, data_type: DataTypeId) -> &DataValue {
         self.data.get(&data_type).unwrap_or(&DataValue::Null)
     }
-
-    pub fn get_data_container(&self) -> &BTreeMap<DataTypeId, DataValue> {
-        &self.data
-    }
 }
 
 impl CommunicationValue {
@@ -131,51 +135,105 @@ impl CommunicationValue {
      * [optional 4 bytes id]                // if flags bit2 set
      * [optional 6 bytes sender]            // if flags bit0 set
      * [optional 6 bytes receiver]          // if flags bit1 set
+     * [optional 1 byte signature type]     // if flags bit4 set; Type defines length of signature
+     * [optional signature]                 // if flags bit4 set
      * [data container bytes...]
      *
      * Flags:
      * bit0 => has sender
      * bit1 => has receiver
      * bit2 => has id
+     * 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> {
-        let mut payload = Vec::new();
-
         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| {
+            matches!(
+                v,
+                DataValue::EncryptedContainer(_) | DataValue::SignedEncryptedContainer(_)
+            )
+        });
+        #[cfg(not(feature = "crypto"))]
+        let is_encrypted = false;
+
+        #[cfg(feature = "crypto")]
+        let has_frame_sig = self.frame_signature.is_some();
+        #[cfg(not(feature = "crypto"))]
+        let has_frame_sig = false;
+
         let mut flags: u8 = 0;
         if has_sender {
-            flags |= 0b0000_0001;
+            flags |= FLAG_HAS_SENDER;
         }
         if has_receiver {
-            flags |= 0b0000_0010;
+            flags |= FLAG_HAS_RECEIVER;
         }
         if has_id {
-            flags |= 0b0000_0100;
+            flags |= FLAG_HAS_ID;
+        }
+        if is_encrypted {
+            flags |= FLAG_ENCRYPTED;
+        }
+        if has_frame_sig {
+            flags |= FLAG_SIGNED;
         }
 
-        let _ = payload.write_u16::<BigEndian>(self.comm_type.0);
+        let mut metadata = Vec::new();
-        payload.push(flags);
+        let _ = metadata.write_u16::<BigEndian>(self.comm_type.0);
+        metadata.push(flags);
 
         if has_id {
-            let _ = payload.write_u32::<BigEndian>(self.id);
+            let _ = metadata.write_u32::<BigEndian>(self.id);
         }
 
         if has_sender {
             let sender_be = self.sender.to_be_bytes();
-            payload.extend_from_slice(&sender_be[2..]); // 6 bytes
+            metadata.extend_from_slice(&sender_be[2..]);
         }
 
         if has_receiver {
             let receiver_be = self.receiver.to_be_bytes();
-            payload.extend_from_slice(&receiver_be[2..]); // 6 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);
-        let container_bytes = container_value.to_bytes();
+            container_value.to_bytes()
-        payload.extend_from_slice(&container_bytes);
+        };
+
+        #[cfg(not(feature = "crypto"))]
+        let data_bytes = {
+            let container_value = DataValue::container_from_map(&self.data);
+            container_value.to_bytes()
+        };
+
+        let mut payload = Vec::new();
+        payload.extend_from_slice(&metadata);
+
+        #[cfg(feature = "crypto")]
+        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.extend_from_slice(&data_bytes);
 
         let mut frame = Vec::with_capacity(4 + payload.len());
         let _ = frame.write_u32::<BigEndian>(payload.len() as u32);
@@ -184,33 +242,48 @@ impl CommunicationValue {
         frame
     }
 
-    pub fn from_bytes(bytes: &[u8]) -> Option<Self> {
+    pub fn from_bytes(bytes: &[u8]) -> Result<Self, CodecError> {
         let mut cursor = Cursor::new(bytes);
 
-        let total_len = cursor.read_u32::<BigEndian>().ok()? as usize;
+        let total_len = cursor
+            .read_u32::<BigEndian>()
+            .map_err(|_| CodecError::InvalidEncoding)? as usize;
         if bytes.len() < 4 + total_len {
-            return None;
+            return Err(CodecError::InvalidEncoding);
         }
 
         let frame_end = 4 + total_len;
 
-        let comm_type_num = cursor.read_u16::<BigEndian>().ok()?;
+        let comm_type_num = cursor
+            .read_u16::<BigEndian>()
+            .map_err(|_| CodecError::InvalidEncoding)?;
         let comm_type = CommunicationTypeId(comm_type_num);
 
-        let flags = cursor.read_u8().ok()?;
+        let flags = cursor.read_u8().map_err(|_| CodecError::InvalidEncoding)?;
-        let has_sender = (flags & 0b0000_0001) != 0;
+        let has_sender = (flags & FLAG_HAS_SENDER) != 0;
-        let has_receiver = (flags & 0b0000_0010) != 0;
+        let has_receiver = (flags & FLAG_HAS_RECEIVER) != 0;
-        let has_id = (flags & 0b0000_0100) != 0;
+        let has_id = (flags & FLAG_HAS_ID) != 0;
+        let is_encrypted = (flags & FLAG_ENCRYPTED) != 0;
+        let is_signed = (flags & FLAG_SIGNED) != 0;
+
+        #[cfg(not(feature = "crypto"))]
+        if is_signed || is_encrypted {
+            return Err(CodecError::InvalidEncoding);
+        }
 
         let id = if has_id {
-            cursor.read_u32::<BigEndian>().ok()?
+            cursor
+                .read_u32::<BigEndian>()
+                .map_err(|_| CodecError::InvalidEncoding)?
         } else {
             0
         };
 
         let sender = if has_sender {
             let mut buf = [0u8; 8];
-            cursor.read_exact(&mut buf[2..]).ok()?;
+            cursor
+                .read_exact(&mut buf[2..])
+                .map_err(|_| CodecError::InvalidEncoding)?;
             u64::from_be_bytes(buf)
         } else {
             0
@@ -218,28 +291,226 @@ impl CommunicationValue {
 
         let receiver = if has_receiver {
             let mut buf = [0u8; 8];
-            cursor.read_exact(&mut buf[2..]).ok()?;
+            cursor
+                .read_exact(&mut buf[2..])
+                .map_err(|_| CodecError::InvalidEncoding)?;
             u64::from_be_bytes(buf)
         } else {
             0
         };
 
+        #[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 mut sig = vec![0u8; sig_len];
+            cursor
+                .read_exact(&mut sig)
+                .map_err(|_| CodecError::InvalidEncoding)?;
+            Some((alg, sig))
+        } else {
+            None
+        };
+
         let pos = cursor.position() as usize;
         if pos > frame_end {
-            return None;
+            return Err(CodecError::InvalidEncoding);
         }
 
         let data_bytes = &bytes[pos..frame_end];
-        let data_value = DataValue::from_bytes(data_bytes)?;
-        let data = data_value.as_map()?;
 
-        Some(Self {
+        #[cfg(feature = "crypto")]
+        let data = if is_encrypted {
+            let mut map = BTreeMap::new();
+            map.insert(
+                DataTypeId(0),
+                DataValue::EncryptedContainer(data_bytes.to_vec()),
+            );
+            map
+        } else {
+            let data_value =
+                DataValue::from_bytes(data_bytes).ok_or(CodecError::InvalidEncoding)?;
+            data_value.as_map().ok_or(CodecError::InvalidEncoding)?
+        };
+
+        #[cfg(not(feature = "crypto"))]
+        let data = {
+            let data_value =
+                DataValue::from_bytes(data_bytes).ok_or(CodecError::InvalidEncoding)?;
+            data_value.as_map().ok_or(CodecError::InvalidEncoding)?
+        };
+
+        Ok(Self {
             id,
             comm_type,
             sender,
             receiver,
             data,
             type_map: None,
+            #[cfg(feature = "crypto")]
+            frame_signature,
+        })
+    }
+
+    pub fn from_bytes_with(bytes: &[u8], tm: &TypeMap) -> Result<Self, CodecError> {
+        let mut val = Self::from_bytes(bytes)?;
+        val.type_map = Some(tm.clone());
+        Ok(val)
+    }
+
+    /*
+     * Sign the frame. Computes a signature over the canonical form:
+     *   comm_type || flags || id? || sender? || receiver? || data_bytes
+     *
+     * After calling this, `to_bytes()` will embed the algorithm and
+     * 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();
+        let sig = signer.sign(&signed_payload).ok()?;
+        self.frame_signature = Some((algorithm, sig));
+        Some(())
+    }
+
+    /*
+     * Verify the frame signature. Reconstructs the signed payload from
+     * current state and checks it against the stored signature.
+     */
+    #[cfg(feature = "crypto")]
+    pub fn verify_frame(&self, verifier: &impl SignatureScheme) -> Result<(), CodecError> {
+        let (_algorithm, sig) = self
+            .frame_signature
+            .as_ref()
+            .ok_or(CodecError::InvalidEncoding)?;
+
+        let signed_payload = self.build_signed_payload();
+        verifier
+            .verify(&signed_payload, sig)
+            .map_err(|_| CodecError::InvalidEncoding)
+    }
+
+    /*
+     * Reconstruct the signed payload that the frame signature covers:
+     *   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()
+    }
+
+    #[cfg(feature = "crypto")]
+    pub fn get_frame_signature(&self) -> Option<&(u8, Vec<u8>)> {
+        self.frame_signature.as_ref()
+    }
+
+    #[cfg(feature = "registry")]
+    pub fn migrate(&self, target_tm: &TypeMap) -> Result<Self, CodecError> {
+        let comm_name = communication_type_name(self.comm_type.0)
+            .ok_or_else(|| CodecError::UnknownCommunicationType(self.comm_type.0.to_string()))?;
+        let comm_variant = CommunicationType::from_name(comm_name)
+            .ok_or_else(|| CodecError::UnknownCommunicationType(comm_name.to_string()))?;
+        let new_comm_id = CommunicationTypeId(
+            target_tm
+                .comm_id_enum(comm_variant)
+                .ok_or_else(|| CodecError::UnknownCommunicationType(comm_name.to_string()))?,
+        );
+
+        let mut new_data = BTreeMap::new();
+        for (&old_id, value) in &self.data {
+            let name = data_type_name(old_id.0)
+                .ok_or_else(|| CodecError::UnknownDataType(old_id.0.to_string()))?;
+            let variant = DataType::from_name(name)
+                .ok_or_else(|| CodecError::UnknownDataType(name.to_string()))?;
+            let new_id = DataTypeId(
+                target_tm
+                    .data_id_enum(variant)
+                    .ok_or_else(|| CodecError::UnknownDataType(name.to_string()))?,
+            );
+            new_data.insert(new_id, value.clone());
+        }
+
+        Ok(Self {
+            id: self.id,
+            comm_type: new_comm_id,
+            sender: self.sender,
+            receiver: self.receiver,
+            data: new_data,
+            type_map: Some(target_tm.clone()),
+            #[cfg(feature = "crypto")]
+            frame_signature: self.frame_signature.clone(),
         })
     }
 }
@@ -275,11 +546,26 @@ fn fmt_data_value(val: &DataValue, f: &mut fmt::Formatter<'_>) -> fmt::Result {
     }
 }
 
+#[cfg(debug_assertions)]
 const BOLD_BLUE: &str = "\x1b[1;34m";
+#[cfg(not(debug_assertions))]
+const BOLD_BLUE: &str = "";
+#[cfg(debug_assertions)]
 const GREEN: &str = "\x1b[32m";
+#[cfg(not(debug_assertions))]
+const GREEN: &str = "";
+#[cfg(debug_assertions)]
 const YELLOW: &str = "\x1b[33m";
+#[cfg(not(debug_assertions))]
+const YELLOW: &str = "";
+#[cfg(debug_assertions)]
 const ORANGE: &str = "\x1b[38;5;208m";
+#[cfg(not(debug_assertions))]
+const ORANGE: &str = "";
+#[cfg(debug_assertions)]
 const RESET: &str = "\x1b[0m";
+#[cfg(not(debug_assertions))]
+const RESET: &str = "";
 
 impl fmt::Display for CommunicationValue {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
@@ -333,7 +619,7 @@ mod tests {
 
     #[test]
     fn test_flags_and_order_without_optional() {
-        let cv = CommunicationValue::from_comm_default(CommunicationType::ErrorParsing).with_id(0);
+        let cv = CommunicationValue::new(CommunicationType::ErrorParsing).with_id(0);
         let bytes = cv.to_bytes();
 
         // [u32 len][u16 type][flags]...
@@ -351,7 +637,7 @@ mod tests {
 
     #[test]
     fn test_flags_and_order_with_all_optional() {
-        let cv = CommunicationValue::from_comm_default(CommunicationType::ErrorBadVersion)
+        let cv = CommunicationValue::new(CommunicationType::ErrorBadVersion)
             .with_id(0xAABBCCDD)
             .with_sender(0x0000_1122_3344_5566)
             .with_receiver(0x0000_6677_8899_AABB);
@@ -382,7 +668,7 @@ mod tests {
 
     #[test]
     fn test_roundtrip_complex() {
-        let cv = CommunicationValue::from_comm_default(CommunicationType::Disconnect)
+        let cv = CommunicationValue::new(CommunicationType::Disconnect)
             .with_id(1234)
             .with_sender(111)
             .with_receiver(222)
@@ -415,6 +701,6 @@ mod tests {
         let mut bad = vec![0u8; 8];
         // total_length claims more than available
         bad[0..4].copy_from_slice(&(1000u32.to_be_bytes()));
-        assert!(CommunicationValue::from_bytes(&bad).is_none());
+        assert!(CommunicationValue::from_bytes(&bad).is_err());
     }
 }

codec/src/data_value.rs

@@ -9,7 +9,7 @@ use std::io::Cursor;
 use mtp_type_map::DataTypeId;
 
 #[cfg(feature = "crypto")]
-use mtp_crypto::{AeadDecrypt, AeadEncrypt};
+use mtp_crypto::{AeadDecrypt, AeadEncrypt, SigAlgorithm, SignatureScheme};
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum DataKind {
@@ -27,6 +27,10 @@ pub enum DataKind {
 
     #[cfg(feature = "crypto")]
     EncryptedContainer,
+    #[cfg(feature = "crypto")]
+    SignedContainer,
+    #[cfg(feature = "crypto")]
+    SignedEncryptedContainer,
 
     Null,
 }
@@ -44,12 +48,47 @@ pub enum DataValue {
     Str(String),
     Bytes(Vec<u8>),
     Array(Vec<DataValue>),
-
+    /*
+     * Container format:
+     * [2 bytes u16 entry_count]       // number of entries
+     * [1 byte kind]                   // DataValue kind marker
+     * [if kind == BOOL_TRUE or BOOL_FALSE:]
+     *     [2 bytes u16 key]           // DataTypeId discriminant
+     * [else:]
+     *     [4 bytes u32 payload_len]   // length of the value payload
+     *     [2 bytes u16 key]           // DataTypeId discriminant
+     *     [payload_len bytes payload] // value data (interpreted based on kind)
+     */
     Container(Vec<(DataTypeId, DataValue)>),
 
+    /*
+     * Container format:
+     * [4 bytes u32 entry_count]       // length of the container
+     * [binary data]
+     * -> After decryption, the container is parsed as a regular container
+     */
     #[cfg(feature = "crypto")]
     EncryptedContainer(Vec<u8>),
 
+    /*
+     * Container format:
+     * [4 bytes u32 entry_count]       // length of the container
+     * [binary data]
+     * -> Can be turned into Container
+     * -> Can be used with a public key to verify integrity
+     */
+    #[cfg(feature = "crypto")]
+    SignedContainer(Vec<u8>),
+
+    /*
+     * Container format:
+     * [4 bytes u32 entry_count]       // length of the container
+     * [binary data]
+     * -> After decryption, the container is parsed as a signed container
+     */
+    #[cfg(feature = "crypto")]
+    SignedEncryptedContainer(Vec<u8>),
+
     Null,
 }
 
@@ -86,7 +125,9 @@ impl DataValue {
      * 0x08 => Array
      * 0x09 => Container
      * 0x0A => EncryptedContainer (4 bytes u32 len + encrypted bytes)
-     * 0x0B => Null
+     * 0x0B => SignedContainer (4 bytes u32 len + 3373 bytes signature)
+     * 0x0C => SignedEncryptedContainer (4 bytes u32 len + 3373 bytes signature + encrypted bytes)
+     * 0xFF => Null
      */
     const KIND_BOOL_TRUE: u8 = 0x01;
     const KIND_BOOL_FALSE: u8 = 0x02;
@@ -102,8 +143,12 @@ impl DataValue {
     const KIND_CONTAINER: u8 = 0x09;
     #[cfg(feature = "crypto")]
     const KIND_ENCRYPTED_CONTAINER: u8 = 0x0A;
+    #[cfg(feature = "crypto")]
+    const KIND_SIGNED_CONTAINER: u8 = 0x0B;
+    #[cfg(feature = "crypto")]
+    const KIND_SIGNED_ENCRYPTED_CONTAINER: u8 = 0x0C;
 
-    const KIND_NULL: u8 = 0x0B;
+    const KIND_NULL: u8 = 0xFF;
 
     pub fn container_from_map(map: &BTreeMap<DataTypeId, DataValue>) -> DataValue {
         let mut container = Vec::new();
@@ -131,6 +176,10 @@ impl DataValue {
             DataValue::Container(_) => DataKind::Container,
             #[cfg(feature = "crypto")]
             DataValue::EncryptedContainer(_) => DataKind::EncryptedContainer,
+            #[cfg(feature = "crypto")]
+            DataValue::SignedContainer(_) => DataKind::SignedContainer,
+            #[cfg(feature = "crypto")]
+            DataValue::SignedEncryptedContainer(_) => DataKind::SignedEncryptedContainer,
             DataValue::Null => DataKind::Null,
         }
     }
@@ -205,17 +254,29 @@ impl DataValue {
         }
     }
 
+    #[cfg(feature = "crypto")]
+    pub fn as_signed_container(&self) -> Option<Vec<u8>> {
+        match self {
+            DataValue::SignedContainer(b) => Some(b.clone()),
+            _ => None,
+        }
+    }
+
+    #[cfg(feature = "crypto")]
+    pub fn as_signed_encrypted_container(&self) -> Option<Vec<u8>> {
+        match self {
+            DataValue::SignedEncryptedContainer(c) => Some(c.clone()),
+            _ => None,
+        }
+    }
+
     /*
      * Decrypt an `EncryptedContainer` in-place, replacing it with the
      * deserialized `Container`. Returns `None` if decryption or
      * deserialization fails.
      */
     #[cfg(feature = "crypto")]
-    pub fn decrypt_into_container(
+    pub fn decrypt_into_container(&mut self, cipher: &impl AeadDecrypt, aad: &[u8]) -> Option<()> {
-        &mut self,
-        cipher: &impl AeadDecrypt,
-        aad: &[u8],
-    ) -> Option<()> {
         let data = self.as_encrypted_container()?;
         let plaintext = cipher.decrypt(&data, aad).ok()?;
         let dv = DataValue::from_bytes(&plaintext)?;
@@ -233,11 +294,7 @@ impl DataValue {
      * Returns `None` if the value is not a `Container` or encryption fails.
      */
     #[cfg(feature = "crypto")]
-    pub fn encrypt_container(
+    pub fn encrypt_container(&mut self, cipher: &impl AeadEncrypt, aad: &[u8]) -> Option<()> {
-        &mut self,
-        cipher: &impl AeadEncrypt,
-        aad: &[u8],
-    ) -> Option<()> {
         let entries = self.as_container()?;
         let plaintext = DataValue::Container(entries).to_bytes();
         let ct = cipher.encrypt(&plaintext, aad).ok()?;
@@ -245,6 +302,98 @@ impl DataValue {
         Some(())
     }
 
+    /*
+     * Sign a `Container` in-place, replacing it with a `SignedContainer`.
+     * The wire blob is: [1 byte alg] [N bytes sig] [serialized container bytes].
+     * The signature covers only the serialized container bytes (not the alg byte).
+     * Returns `None` if the value is not a `Container` or signing fails.
+     */
+    #[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 sig = signer.sign(&container_bytes).ok()?;
+
+        let mut blob = Vec::with_capacity(1 + sig.len() + container_bytes.len());
+        blob.push(algorithm);
+        blob.extend_from_slice(&sig);
+        blob.extend_from_slice(&container_bytes);
+
+        *self = DataValue::SignedContainer(blob);
+        Some(())
+    }
+
+    /*
+     * Verify a `SignedContainer` in-place, replacing it with the deserialized
+     * `Container` on success. Returns `None` if verification fails or the
+     * blob is malformed.
+     */
+    #[cfg(feature = "crypto")]
+    pub fn verify_into_container(
+        &mut self,
+        verifier: &impl SignatureScheme,
+    ) -> Option<()> {
+        let blob = self.as_signed_container()?;
+        if blob.len() < 1 + 64 + 2 {
+            return None;
+        }
+
+        let algorithm = blob[0];
+        let sig_len = SigAlgorithm::length(algorithm)?;
+        if blob.len() < 1 + sig_len + 2 {
+            return None;
+        }
+
+        let signature = &blob[1..1 + sig_len];
+        let container_bytes = &blob[1 + sig_len..];
+
+        verifier.verify(container_bytes, signature).ok()?;
+
+        let entries = DataValue::from_bytes(container_bytes)?.as_container()?;
+        *self = DataValue::Container(entries);
+        Some(())
+    }
+
+    /*
+     * Encrypt a `Container` into a `SignedEncryptedContainer` in-place.
+     * 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,
+        aad: &[u8],
+    ) -> Option<()> {
+        self.sign_container(algorithm, signer)?;
+        let blob = self.as_signed_container()?;
+        let ct = cipher.encrypt(&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.
+     */
+    #[cfg(feature = "crypto")]
+    pub fn decrypt_signed_encrypted_container(
+        &mut self,
+        cipher: &impl AeadDecrypt,
+        aad: &[u8],
+    ) -> Option<()> {
+        let data = self.as_signed_encrypted_container()?;
+        let plaintext = cipher.decrypt(&data, aad).ok()?;
+        *self = DataValue::SignedContainer(plaintext);
+        Some(())
+    }
+
     pub fn as_map(&self) -> Option<BTreeMap<DataTypeId, DataValue>> {
         match self {
             DataValue::Container(c) => {
@@ -311,7 +460,8 @@ impl DataValue {
         let kind = Self::kind_marker(value);
         buf.push(kind);
 
-        if kind == Self::KIND_BOOL_TRUE || kind == Self::KIND_BOOL_FALSE {
+        if kind == Self::KIND_BOOL_TRUE || kind == Self::KIND_BOOL_FALSE || kind == Self::KIND_NULL
+        {
             let _ = buf.write_u16::<BigEndian>(key.0);
             return true;
         }
@@ -351,7 +501,8 @@ impl DataValue {
         let kind = Self::kind_marker(value);
         buf.push(kind);
 
-        if kind == Self::KIND_BOOL_TRUE || kind == Self::KIND_BOOL_FALSE {
+        if kind == Self::KIND_BOOL_TRUE || kind == Self::KIND_BOOL_FALSE || kind == Self::KIND_NULL
+        {
             return true;
         }
 
@@ -414,6 +565,16 @@ impl DataValue {
                 buf.extend_from_slice(data);
                 Some(())
             }
+            #[cfg(feature = "crypto")]
+            DataValue::SignedContainer(data) => {
+                buf.extend_from_slice(data);
+                Some(())
+            }
+            #[cfg(feature = "crypto")]
+            DataValue::SignedEncryptedContainer(data) => {
+                buf.extend_from_slice(data);
+                Some(())
+            }
 
             DataValue::Null => Some(()),
         }
@@ -448,12 +609,17 @@ impl DataValue {
         for _ in 0..count {
             let kind = cursor.read_u8().ok()?;
 
-            if kind == Self::KIND_BOOL_TRUE || kind == Self::KIND_BOOL_FALSE {
+            if kind == Self::KIND_BOOL_TRUE
+                || kind == Self::KIND_BOOL_FALSE
+                || kind == Self::KIND_NULL
+            {
                 let key = DataTypeId(cursor.read_u16::<BigEndian>().ok()?);
                 let value = if kind == Self::KIND_BOOL_TRUE {
                     DataValue::BoolTrue
-                } else {
+                } else if kind == Self::KIND_BOOL_FALSE {
                     DataValue::BoolFalse
+                } else {
+                    DataValue::Null
                 };
                 entries.push((key, value));
                 continue;
@@ -489,12 +655,18 @@ impl DataValue {
         for _ in 0..count {
             let kind = cursor.read_u8().ok()?;
 
-            if kind == Self::KIND_BOOL_TRUE {
+            if kind == Self::KIND_BOOL_TRUE
-                out.push(DataValue::BoolTrue);
+                || kind == Self::KIND_BOOL_FALSE
-                continue;
+                || kind == Self::KIND_NULL
-            }
+            {
-            if kind == Self::KIND_BOOL_FALSE {
+                let value = if kind == Self::KIND_BOOL_TRUE {
-                out.push(DataValue::BoolFalse);
+                    DataValue::BoolTrue
+                } else if kind == Self::KIND_BOOL_FALSE {
+                    DataValue::BoolFalse
+                } else {
+                    DataValue::Null
+                };
+                out.push(value);
                 continue;
             }
 
@@ -604,9 +776,33 @@ impl DataValue {
                 cursor.set_position(end as u64);
                 Some(DataValue::EncryptedContainer(data))
             }
+            #[cfg(feature = "crypto")]
+            Self::KIND_SIGNED_CONTAINER => {
+                let len = payload_len?;
+                let start = cursor.position() as usize;
+                let end = start.checked_add(len)?;
+                if end > cursor.get_ref().len() {
+                    return None;
+                }
+                let data = cursor.get_ref()[start..end].to_vec();
+                cursor.set_position(end as u64);
+                Some(DataValue::SignedContainer(data))
+            }
+            #[cfg(feature = "crypto")]
+            Self::KIND_SIGNED_ENCRYPTED_CONTAINER => {
+                let len = payload_len?;
+                let start = cursor.position() as usize;
+                let end = start.checked_add(len)?;
+                if end > cursor.get_ref().len() {
+                    return None;
+                }
+                let data = cursor.get_ref()[start..end].to_vec();
+                cursor.set_position(end as u64);
+                Some(DataValue::SignedEncryptedContainer(data))
+            }
             Self::KIND_NULL => Some(DataValue::Null),
             #[cfg(not(feature = "crypto"))]
-            0x0A => None,
+            0x0A | 0x0B | 0x0C => None,
             _ => None,
         }
     }
@@ -631,6 +827,10 @@ impl DataValue {
             DataValue::Container(_) => Self::KIND_CONTAINER,
             #[cfg(feature = "crypto")]
             DataValue::EncryptedContainer(_) => Self::KIND_ENCRYPTED_CONTAINER,
+            #[cfg(feature = "crypto")]
+            DataValue::SignedContainer(_) => Self::KIND_SIGNED_CONTAINER,
+            #[cfg(feature = "crypto")]
+            DataValue::SignedEncryptedContainer(_) => Self::KIND_SIGNED_ENCRYPTED_CONTAINER,
             DataValue::Null => Self::KIND_NULL,
         }
     }
@@ -669,6 +869,10 @@ impl fmt::Display for DataValue {
             DataValue::Bytes(_) => write!(f, "(Binary)"),
             #[cfg(feature = "crypto")]
             DataValue::EncryptedContainer(_) => write!(f, "(Secure)"),
+            #[cfg(feature = "crypto")]
+            DataValue::SignedContainer(_) => write!(f, "(Signed)"),
+            #[cfg(feature = "crypto")]
+            DataValue::SignedEncryptedContainer(_) => write!(f, "(SignedSecure)"),
             DataValue::Null => write!(f, "null"),
         }
     }
@@ -691,6 +895,10 @@ impl PartialEq for DataValue {
             (Container(a), Container(b)) => a == b,
             #[cfg(feature = "crypto")]
             (EncryptedContainer(a), EncryptedContainer(b)) => a == b,
+            #[cfg(feature = "crypto")]
+            (SignedContainer(a), SignedContainer(b)) => a == b,
+            #[cfg(feature = "crypto")]
+            (SignedEncryptedContainer(a), SignedEncryptedContainer(b)) => a == b,
             (Null, Null) => true,
             _ => false,
         }
@@ -743,9 +951,408 @@ impl Hash for DataValue {
                 6u8.hash(state);
                 c.hash(state);
             }
-            Null => {
+            #[cfg(feature = "crypto")]
+            SignedContainer(c) => {
                 7u8.hash(state);
+                c.hash(state);
+            }
+            #[cfg(feature = "crypto")]
+            SignedEncryptedContainer(c) => {
+                8u8.hash(state);
+                c.hash(state);
+            }
+            Null => {
+                9u8.hash(state);
             }
         }
     }
 }
+
+/* ================================ TESTS ================================ */
+#[cfg(test)]
+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 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 decoded = DataValue::from_bytes(&bytes).expect("roundtrip failed");
+        assert_eq!(dv, decoded, "array roundtrip mismatch");
+    }
+
+    #[test]
+    fn test_bool_in_container() {
+        container_roundtrip(vec![
+            (DataTypeId(1), DataValue::BoolTrue),
+            (DataTypeId(2), DataValue::BoolFalse),
+        ]);
+    }
+
+    #[test]
+    fn test_bool_true_eq() {
+        assert_eq!(DataValue::BoolTrue, DataValue::Bool(true));
+        assert_eq!(DataValue::BoolFalse, DataValue::Bool(false));
+        assert_ne!(DataValue::BoolTrue, DataValue::Bool(false));
+    }
+
+    #[test]
+    fn test_bool_as_bool() {
+        assert_eq!(DataValue::BoolTrue.as_bool(), Some(true));
+        assert_eq!(DataValue::BoolFalse.as_bool(), Some(false));
+        assert_eq!(DataValue::Bool(true).as_bool(), Some(true));
+        assert_eq!(DataValue::Null.as_bool(), None);
+    }
+
+    #[test]
+    fn test_signed_number_in_container() {
+        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)),
+        ]);
+    }
+
+    #[test]
+    fn test_unsigned_number_in_container() {
+        container_roundtrip(vec![
+            (DataTypeId(1), DataValue::UnsignedNumber(0)),
+            (DataTypeId(2), DataValue::UnsignedNumber(42)),
+            (DataTypeId(3), DataValue::UnsignedNumber(u128::MAX)),
+        ]);
+    }
+
+    #[test]
+    fn test_float_in_container() {
+        container_roundtrip(vec![
+            (DataTypeId(1), DataValue::Float(0, 0)),
+            (DataTypeId(2), DataValue::Float(2, 12345)),
+            (DataTypeId(3), DataValue::Float(255, 4294967295)),
+        ]);
+    }
+
+    #[test]
+    fn test_str_in_container() {
+        container_roundtrip(vec![
+            (DataTypeId(1), DataValue::Str(String::new())),
+            (DataTypeId(2), DataValue::Str("hello".to_string())),
+            (DataTypeId(3), DataValue::Str("a".repeat(1000))),
+        ]);
+    }
+
+    #[test]
+    fn test_bytes_in_container() {
+        container_roundtrip(vec![
+            (DataTypeId(1), DataValue::Bytes(vec![])),
+            (DataTypeId(2), DataValue::Bytes(vec![0x00, 0xFF, 0xAB])),
+            (DataTypeId(3), DataValue::Bytes(vec![0x42; 100])),
+        ]);
+    }
+
+    #[test]
+    fn test_null_in_container() {
+        container_roundtrip(vec![
+            (DataTypeId(1), DataValue::Null),
+        ]);
+    }
+
+    #[test]
+    fn test_array_non_empty_roundtrip() {
+        array_roundtrip(vec![
+            DataValue::BoolTrue,
+            DataValue::SignedNumber(42),
+            DataValue::Str("hello".to_string()),
+            DataValue::Null,
+        ]);
+    }
+
+    #[test]
+    fn test_array_nested_roundtrip() {
+        array_roundtrip(vec![
+            DataValue::Array(vec![DataValue::BoolTrue, DataValue::BoolFalse]),
+            DataValue::Array(vec![DataValue::SignedNumber(1), DataValue::SignedNumber(2)]),
+        ]);
+    }
+
+    #[test]
+    fn test_container_empty_roundtrip() {
+        container_roundtrip(vec![]);
+    }
+
+    #[test]
+    fn test_container_mixed_roundtrip() {
+        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),
+        ]);
+    }
+
+    #[test]
+    fn test_container_nested_roundtrip() {
+        container_roundtrip(vec![
+            (DataTypeId(1), DataValue::Container(vec![
+                (DataTypeId(10), DataValue::BoolTrue),
+            ])),
+            (DataTypeId(2), DataValue::Array(vec![
+                DataValue::SignedNumber(1),
+                DataValue::SignedNumber(2),
+            ])),
+        ]);
+    }
+
+    #[test]
+    fn test_container_base64_roundtrip() {
+        let dv = DataValue::Container(vec![
+            (DataTypeId(7), DataValue::Bytes(vec![0xDE, 0xAD, 0xBE, 0xEF])),
+        ]);
+        let b64 = dv.to_base64();
+        let decoded = DataValue::from_base64(&b64).expect("base64 roundtrip failed");
+        assert_eq!(dv, decoded);
+    }
+
+    #[test]
+    fn test_kind_classification() {
+        assert_eq!(DataValue::BoolTrue.kind(), DataKind::Bool);
+        assert_eq!(DataValue::Bool(false).kind(), DataKind::Bool);
+        assert_eq!(DataValue::SignedNumber(0).kind(), DataKind::SignedNumber);
+        assert_eq!(DataValue::UnsignedNumber(0).kind(), DataKind::UnsignedNumber);
+        assert_eq!(DataValue::Float(0, 0).kind(), DataKind::Float);
+        assert_eq!(DataValue::Str(String::new()).kind(), DataKind::Str);
+        assert_eq!(DataValue::Bytes(vec![]).kind(), DataKind::Bytes);
+        assert_eq!(DataValue::Array(vec![]).kind(), DataKind::Array(Box::new(DataKind::Null)));
+        assert_eq!(DataValue::Container(vec![]).kind(), DataKind::Container);
+        assert_eq!(DataValue::Null.kind(), DataKind::Null);
+    }
+
+    #[test]
+    fn test_as_accessors() {
+        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])),
+        ]);
+
+        let map = dv.as_map().expect("should be a container");
+        assert_eq!(map.get(&DataTypeId(1)).and_then(|v| v.as_str()), Some("alice"));
+        assert_eq!(map.get(&DataTypeId(2)).and_then(|v| v.as_signed_number()), Some(42));
+        assert_eq!(map.get(&DataTypeId(3)).and_then(|v| v.as_bytes()), Some(vec![0x01, 0x02]));
+        assert_eq!(map.get(&DataTypeId(4)).and_then(|v| v.as_array()), Some(vec![DataValue::BoolTrue]));
+    }
+
+    #[test]
+    fn test_as_string() {
+        let dv = DataValue::Str("hello".to_string());
+        assert_eq!(dv.as_string(), Some("hello".to_string()));
+        assert_eq!(dv.as_str(), Some("hello"));
+        assert_eq!(DataValue::Null.as_string(), None);
+    }
+
+    #[test]
+    fn test_as_float() {
+        assert_eq!(DataValue::Float(3, 14).as_float(), Some((3, 14)));
+        assert_eq!(DataValue::Null.as_float(), None);
+    }
+
+    #[test]
+    fn test_container_from_map() {
+        let mut map = BTreeMap::new();
+        map.insert(DataTypeId(1), DataValue::BoolTrue);
+        map.insert(DataTypeId(2), DataValue::SignedNumber(99));
+        let dv = DataValue::container_from_map(&map);
+        let container = dv.as_container().expect("should be container");
+        assert_eq!(container.len(), 2);
+    }
+
+    #[test]
+    fn test_invalid_short_input() {
+        assert!(DataValue::from_bytes(&[]).is_none());
+        assert!(DataValue::from_bytes(&[0x01]).is_none());
+    }
+
+    #[test]
+    fn test_invalid_kind_rejected() {
+        let bytes = vec![0x00, 0x01, 0x0D, 0x00, 0x00, 0x00, 0x01, 0x00, 0x01, 0x41];
+        assert!(DataValue::from_bytes(&bytes).is_none());
+    }
+
+    #[test]
+    fn test_truncated_container_rejected() {
+        let dv = DataValue::Container(vec![
+            (DataTypeId(1), DataValue::Str("hello".to_string())),
+        ]);
+        let bytes = dv.to_bytes();
+        // 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_display_basic() {
+        assert_eq!(format!("{}", DataValue::BoolTrue), "true");
+        assert_eq!(format!("{}", DataValue::BoolFalse), "false");
+        assert_eq!(format!("{}", DataValue::Null), "null");
+        assert_eq!(format!("{}", DataValue::SignedNumber(42)), "42");
+        assert_eq!(format!("{}", DataValue::UnsignedNumber(42)), "42");
+        assert_eq!(format!("{}", DataValue::Str("hi".to_string())), "\"hi\"");
+        assert_eq!(format!("{}", DataValue::Bytes(vec![])), "(Binary)");
+    }
+
+    #[test]
+    fn test_hash_consistency() {
+        use std::collections::HashSet;
+        let mut set = HashSet::new();
+        set.insert(DataValue::BoolTrue);
+        set.insert(DataValue::BoolFalse);
+        set.insert(DataValue::Null);
+        set.insert(DataValue::SignedNumber(1));
+        set.insert(DataValue::UnsignedNumber(1));
+        assert_eq!(set.len(), 5);
+        set.insert(DataValue::Bool(true));
+        assert_eq!(set.len(), 5);
+    }
+
+    #[test]
+    fn test_float_display() {
+        let s = format!("{}", DataValue::Float(2, 12345));
+        assert_eq!(s, "12345e2");
+    }
+
+    #[test]
+    fn test_container_display() {
+        let dv = DataValue::Container(vec![
+            (DataTypeId(3), DataValue::Str("v2.0".to_string())),
+            (DataTypeId(6), DataValue::UnsignedNumber(42)),
+        ]);
+        let s = format!("{}", dv);
+        assert!(s.contains("3:"));
+        assert!(s.contains("6:"));
+    }
+
+    #[test]
+    fn test_array_display() {
+        let dv = DataValue::Array(vec![DataValue::SignedNumber(1), DataValue::SignedNumber(2)]);
+        let s = format!("{}", dv);
+        assert_eq!(s, "[1, 2]");
+    }
+
+    /* ===== Crypto container tests ===== */
+
+    #[cfg(feature = "crypto")]
+    #[test]
+    fn test_encrypt_decrypt_container_roundtrip() {
+        use mtp_crypto::ChaCha20Poly1305;
+        let key = [0xAB; 32];
+        let cipher = ChaCha20Poly1305::new(key);
+
+        let mut dv = DataValue::Container(vec![
+            (DataTypeId(1), DataValue::Str("secret".to_string())),
+            (DataTypeId(2), DataValue::UnsignedNumber(42)),
+        ]);
+
+        assert!(dv.encrypt_container(&cipher, b"aad").is_some());
+        assert!(matches!(dv, DataValue::EncryptedContainer(_)));
+
+        assert!(dv.decrypt_into_container(&cipher, b"aad").is_some());
+        assert!(matches!(dv, DataValue::Container(_)));
+
+        let entries = dv.as_container().unwrap();
+        assert_eq!(entries.len(), 2);
+    }
+
+    #[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]);
+
+        let mut dv = DataValue::Container(vec![
+            (DataTypeId(1), 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());
+    }
+
+    #[cfg(feature = "crypto")]
+    #[test]
+    fn test_encrypt_container_wrong_aad_fails() {
+        use mtp_crypto::ChaCha20Poly1305;
+        let cipher = ChaCha20Poly1305::new([0xAB; 32]);
+
+        let mut dv = DataValue::Container(vec![
+            (DataTypeId(1), 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());
+    }
+
+    #[cfg(feature = "crypto")]
+    #[test]
+    fn test_encrypt_non_container_fails() {
+        let cipher = mtp_crypto::ChaCha20Poly1305::new([0xAB; 32]);
+
+        let mut dv = DataValue::Str("not a container".to_string());
+        assert!(dv.encrypt_container(&cipher, b"aad").is_none());
+    }
+
+    #[cfg(feature = "crypto")]
+    #[test]
+    fn test_sign_verify_container_roundtrip() {
+        use mtp_crypto::{ChaCha20Poly1305, Ed25519Signer, SigAlgorithm};
+
+        let (signer, sk, _pk) = Ed25519Signer::generate();
+        let cipher = ChaCha20Poly1305::new([0xAB; 32]);
+
+        let mut dv = DataValue::Container(vec![
+            (DataTypeId(1), DataValue::Str("signed data".to_string())),
+        ]);
+
+        assert!(dv
+            .sign_and_encrypt_container(SigAlgorithm::ED25519, &signer, &cipher, b"aad")
+            .is_some());
+        assert!(matches!(dv, DataValue::SignedEncryptedContainer(_)));
+
+        assert!(dv.decrypt_signed_encrypted_container(&cipher, b"aad").is_some());
+        assert!(matches!(dv, DataValue::SignedContainer(_)));
+
+        let verifier = Ed25519Signer::new(&sk).unwrap();
+        assert!(dv.verify_into_container(&verifier).is_some());
+        assert!(matches!(dv, DataValue::Container(_)));
+
+        let entries = dv.as_container().unwrap();
+        assert_eq!(entries.len(), 1);
+    }
+
+    #[cfg(feature = "crypto")]
+    #[test]
+    fn test_sign_container_wrong_key_fails() {
+        use mtp_crypto::{Ed25519Signer, SigAlgorithm};
+
+        let (signer, _, _) = Ed25519Signer::generate();
+        let (_, sk2, _) = Ed25519Signer::generate();
+        let wrong_verifier = Ed25519Signer::new(&sk2).unwrap();
+
+        let mut dv = DataValue::Container(vec![
+            (DataTypeId(1), DataValue::Str("signed data".to_string())),
+        ]);
+
+        assert!(dv.sign_container(SigAlgorithm::ED25519, &signer).is_some());
+        assert!(dv.verify_into_container(&wrong_verifier).is_none());
+    }
+}

codec/src/lib.rs

@@ -1,25 +1,17 @@
 pub mod communication_value;
 pub mod data_value;
-pub mod util;
 
 pub use communication_value::CommunicationValue;
 pub use data_value::{DataKind, DataValue};
-
+pub use mtp_common::CodecError;
-pub use util::rand_u32;
 
 pub use mtp_type_map::{
     communication_type_name, data_type_name, CommunicationType, CommunicationTypeId, DataType,
-    DataTypeId, PROTOCOL_VERSION, TypeMap, Version,
+    DataTypeId, TypeMap, Version, PROTOCOL_VERSION,
 };
 
-use mtp_common::CodecError;
+pub(crate) fn rand_u32() -> u32 {
-
+    rand::random()
-pub fn encode(_value: &DataValue, _typemap: &TypeMap) -> Result<Vec<u8>, CodecError> {
-    todo!()
-}
-
-pub fn decode(_bytes: &[u8], _typemap: &TypeMap) -> Result<DataValue, CodecError> {
-    todo!()
 }
 
 #[cfg(feature = "registry")]

codec/src/registry.rs

@@ -1,8 +1,7 @@
-pub use registry::Registry;
-
-use crate::{CodecError, DataValue, decode, encode};
 use mtp_type_map::Version;
 
+pub use mtp_type_map::Registry;
+
 /*
  * A version-aware codec that uses a multi-version registry to resolve
  * the correct TypeMap for encoding and decoding operations.
@@ -17,22 +16,6 @@ impl VersionedCodec {
         Self { registry }
     }
 
-    pub fn encode(&self, value: &DataValue, version: Version) -> Result<Vec<u8>, CodecError> {
-        let typemap = self
-            .registry
-            .get(&version)
-            .ok_or(CodecError::UnknownVersion)?;
-        encode(value, typemap)
-    }
-
-    pub fn decode(&self, bytes: &[u8], version: Version) -> Result<DataValue, CodecError> {
-        let typemap = self
-            .registry
-            .get(&version)
-            .ok_or(CodecError::UnknownVersion)?;
-        decode(bytes, typemap)
-    }
-
     pub fn negotiate(&self, client_versions: &[Version]) -> Option<Version> {
         self.registry.negotiate(client_versions)
     }

codec/src/util.rs

@@ -1,9 +0,0 @@
-use rand::random;
-
-pub fn rand_u32() -> u32 {
-    random::<u32>()
-}
-
-pub fn rand_u64() -> u64 {
-    random::<u64>()
-}

common/Cargo.toml

@@ -10,8 +10,8 @@ wtransport = { version = "0.7.1", default-features = false, features = [
     "quinn",
     "self-signed",
 ] }
-rustls = { version = "0.23.40" }
+rustls = { version = "0.23.41" }
-quinn = { version = "0.11.9", default-features = false, features = [
+quinn = { version = "0.11.11", default-features = false, features = [
     "rustls-aws-lc-rs",
     "rustls",
 ] }

common/src/lib.rs

@@ -1,20 +1,82 @@
 use thiserror::Error;
 
-pub enum RegistryError {
+
-    ReservedCommId(u16, String),
+
+#[derive(Clone, Debug, Error, PartialEq, Eq)]
+pub enum CodecError {
+    #[error("Unknown version")]
+    UnknownVersion,
+    #[error("Unknown communication type: {0}")]
+    UnknownCommunicationType(String),
+    #[error("Unknown data type: {0}")]
+    UnknownDataType(String),
+    #[error("Reserved communication type: {0}")]
+    ReservedCommunicationType(u16),
+    #[error("Invalid encoding")]
+    InvalidEncoding,
+    #[error("Crypto failed: {0}")]
+    CryptoFailed(String),
 }
 
-#[derive(Clone, Debug, PartialEq, Eq)]
+/* ================================ TESTS ================================ */
-pub enum CodecError {
+#[cfg(test)]
-    UnknownVersion,
+mod tests {
-    UnknownCommunicationType(String),
+    use super::*;
-    UnknownDataType(String),
+
-    ReservedCommunicationType(u16),
+    #[test]
-    InvalidEncoding,
+    fn test_codec_error_display() {
+        let e = CodecError::InvalidEncoding;
+        assert_eq!(format!("{}", e), "Invalid encoding");
+    }
+
+    #[test]
+    fn test_codec_error_unknown_version() {
+        assert_eq!(format!("{}", CodecError::UnknownVersion), "Unknown version");
+    }
+
+    #[test]
+    fn test_codec_error_clone_eq() {
+        let a = CodecError::InvalidEncoding;
+        let b = a.clone();
+        assert_eq!(a, b);
+    }
+
+    #[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");
+    }
+
+    #[test]
+    fn test_communication_error_clone() {
+        let a = CommunicationError::UseAfterClosed;
+        let b = a.clone();
+        assert_eq!(format!("{:?}", a), format!("{:?}", b));
+    }
+
+    #[test]
+    fn test_communication_error_authentication_failed() {
+        let e = CommunicationError::AuthenticationFailed("bad key".into());
+        assert!(format!("{}", e).contains("bad key"));
+    }
+
+    #[test]
+    fn test_communication_error_other() {
+        let e = CommunicationError::Other("custom error".into());
+        assert!(format!("{}", e).contains("custom error"));
+    }
 }
 
 #[derive(Debug, Error, Clone)]
 pub enum CommunicationError {
+    /* ====== NOTE on PartialEq ======
+     * Manual PartialEq implementation below compares some inner
+     * values (String, Vec<u8>) structurally and others only by
+     * variant discriminant (types from quinn / wtransport that
+     * don't themselves implement PartialEq).  Keep this impl in
+     * sync if variants are added or changed.
+     */
     #[error("Use after Closed")]
     UseAfterClosed,
 
@@ -39,17 +101,8 @@ pub enum CommunicationError {
     #[error("Loading Certificate error")]
     CertificateLoadFailed,
 
-    #[error("ParseBool error: {0}")]
+    #[error("Parse error: {0}")]
-    ParseBool(#[from] std::str::ParseBoolError),
+    ParseError(String),
-
-    #[error("ParseInt error: {0}")]
-    ParseInt(#[from] std::num::ParseIntError),
-
-    #[error("ParseFloat error: {0}")]
-    ParseFloat(#[from] std::num::ParseFloatError),
-
-    #[error("ParseAddr error: {0}")]
-    ParseAddr(#[from] std::net::AddrParseError),
 
     #[error("Connection error: {0}")]
     ConnectionError(#[from] wtransport::error::ConnectionError),
@@ -87,6 +140,48 @@ pub enum CommunicationError {
     #[error("Crypto Provider Install Error")]
     CryptoProviderInstallFailed,
 
+    #[error("Authentication failed: {0}")]
+    AuthenticationFailed(String),
+
     #[error("Other: {0}")]
     Other(String),
 }
+
+/* ===== Manual PartialEq =====
+ *
+ * Compare by variant, and structurally for String-holding variants.
+ * Third-party error types (quinn, wtransport) are compared only by
+ * discriminant since they don't implement PartialEq themselves.
+ */
+impl PartialEq for CommunicationError {
+    fn eq(&self, other: &Self) -> bool {
+        match (self, other) {
+            (Self::UseAfterClosed, Self::UseAfterClosed) => true,
+            (Self::ClosedLocally, Self::ClosedLocally) => true,
+            (Self::ClosedByPeer, Self::ClosedByPeer) => true,
+            (Self::ConnectionLost, Self::ConnectionLost) => true,
+            (Self::Quinn(_), Self::Quinn(_)) => true,
+            (Self::ParseCommunicationValue, Self::ParseCommunicationValue) => true,
+            (Self::CertificateParseFailed, Self::CertificateParseFailed) => true,
+            (Self::CertificateLoadFailed, Self::CertificateLoadFailed) => true,
+            (Self::ParseError(a), Self::ParseError(b)) => a == b,
+            (Self::ConnectionError(_), Self::ConnectionError(_)) => true,
+            (Self::ConnectingError(a), Self::ConnectingError(b)) => a == b,
+            (Self::ReadToEndError(_), Self::ReadToEndError(_)) => true,
+            (Self::WriteError(_), Self::WriteError(_)) => true,
+            (Self::ClosedError(_), Self::ClosedError(_)) => true,
+            (Self::MessageTooLarge, Self::MessageTooLarge) => true,
+            (Self::ReadExactError(_), Self::ReadExactError(_)) => true,
+            (Self::StreamClosed, Self::StreamClosed) => true,
+            (Self::StreamError, Self::StreamError) => true,
+            (Self::StreamWriteError(_), Self::StreamWriteError(_)) => true,
+            (Self::StreamReadExactError(_), Self::StreamReadExactError(_)) => true,
+            (Self::CryptoProviderInstallFailed, Self::CryptoProviderInstallFailed) => true,
+            (Self::AuthenticationFailed(a), Self::AuthenticationFailed(b)) => a == b,
+            (Self::Other(a), Self::Other(b)) => a == b,
+            _ => false,
+        }
+    }
+}
+
+impl Eq for CommunicationError {}

crypto/Cargo.toml

@@ -20,7 +20,9 @@ ml-dsa = { version = "0.1.1", optional = true }
 serde = { version = "1", optional = true, features = ["derive"] }
 
 [features]
-default = ["chacha20poly1305", "ed25519-dalek", "hkdf", "sha2"]
+default = ["chacha20poly1305", "ed25519-dalek", "hkdf", "sha2", "ml-dsa"]
-full = ["chacha20poly1305", "aes-gcm", "ed25519-dalek", "hkdf", "sha2"]
+# Enabling ml-dsa by default ensures dual-signature support in the handshake
+# without requiring a separate PQC feature flag in host/client crates.
+full = ["default", "aes-gcm"]
 pqc = ["mlkem-tls", "ml-dsa"]
 serde = ["dep:serde"]

crypto/src/helper.rs

@@ -57,15 +57,14 @@ impl MultiEncryptedMessage {
     pub fn from_bytes(bytes: &[u8]) -> Result<Self, CryptoError> {
         let mut offset = 0;
         let read_u16 = |off: &mut usize| -> Result<u16, CryptoError> {
-            let v = u16::from_be_bytes(
+            let slice = bytes
-                bytes
                 .get(*off..*off + 2)
-                    .ok_or(CryptoError::DecryptionFailed)?
+                .ok_or(CryptoError::DecryptionFailed)?;
+            let arr: [u8; 2] = slice
                 .try_into()
-                    .unwrap(),
+                .map_err(|_| CryptoError::DecryptionFailed)?;
-            );
             *off += 2;
-            Ok(v)
+            Ok(u16::from_be_bytes(arr))
         };
 
         let num = read_u16(&mut offset)? as usize;
@@ -198,36 +197,4 @@ pub fn decrypt_multi(
     Err(CryptoError::DecryptionFailed)
 }
 
-#[cfg(feature = "ed25519-dalek")]
-pub fn verify_ed25519_sig(
-    public_key: &crate::keypair::SignaturePublicKey,
-    msg: &[u8],
-    signature: &[u8],
-) -> Result<(), CryptoError> {
-    crate::sign::verify_ed25519(public_key, msg, signature)
-}
 
-#[cfg(feature = "ml-dsa")]
-pub fn verify_ml_dsa_sig(
-    public_key: &crate::keypair::SignaturePqPublicKey,
-    msg: &[u8],
-    signature: &[u8],
-) -> Result<(), CryptoError> {
-    crate::sign::verify_ml_dsa(public_key, msg, signature)
-}
-
-/*
- * Verify both an Ed25519 and ML-DSA signature (dual) against
- * the public keys in a `PublicKeyBundle`.
- */
-#[cfg(all(feature = "ed25519-dalek", feature = "ml-dsa"))]
-pub fn verify_dual_sig(
-    public_keys: &crate::keypair::PublicKeyBundle,
-    msg: &[u8],
-    ed25519_sig: &[u8],
-    mldsa_sig: &[u8],
-) -> Result<(), CryptoError> {
-    verify_ed25519_sig(&public_keys.sig_cl_public_key, msg, ed25519_sig)?;
-    verify_ml_dsa_sig(&public_keys.sig_pq_public_key, msg, mldsa_sig)?;
-    Ok(())
-}

crypto/src/kem.rs

@@ -12,7 +12,7 @@ pub struct HybridKem;
 #[cfg(feature = "mlkem-tls")]
 impl HybridKem {
     pub fn generate_keypair() -> (KemPrivateKey, KemPublicKey) {
-        let (dk, ek) =
+        let (ek, dk) =
             mlkem_tls::X25519MlKem768::keygen(&mut rand_core::OsRng);
         (
             KemPrivateKey::new(dk.as_bytes().to_vec()),

crypto/src/keypair.rs

@@ -84,32 +84,6 @@ impl From<Vec<u8>> for SignaturePublicKey {
     }
 }
 
-#[derive(ZeroizeOnDrop)]
-pub struct KeyGroup {
-    #[zeroize(skip)]
-    pub encryption_public_key: EncryptionPublicKey,
-    pub encryption_private_key: EncryptionPrivateKey,
-    #[zeroize(skip)]
-    pub signature_public_key: SignaturePublicKey,
-    pub signature_private_key: SignaturePrivateKey,
-}
-
-impl KeyGroup {
-    pub fn new(
-        encryption_public_key: EncryptionPublicKey,
-        encryption_private_key: EncryptionPrivateKey,
-        signature_public_key: SignaturePublicKey,
-        signature_private_key: SignaturePrivateKey,
-    ) -> Self {
-        Self {
-            encryption_public_key,
-            encryption_private_key,
-            signature_public_key,
-            signature_private_key,
-        }
-    }
-}
-
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
 #[cfg_attr(feature = "serde", serde(transparent))]
 #[derive(Zeroize, ZeroizeOnDrop)]
@@ -287,14 +261,18 @@ impl PublicKeyBundle {
         use crate::error::CryptoError;
         let mut offset = 0;
 
-        let kem_len = u16::from_be_bytes(
+        let read_u16 = |off: &mut usize| -> Result<u16, CryptoError> {
-            bytes
+            let slice = bytes
-                .get(offset..offset + 2)
+                .get(*off..*off + 2)
-                .ok_or(CryptoError::InvalidKeyLength)?
+                .ok_or(CryptoError::InvalidKeyLength)?;
+            let arr: [u8; 2] = slice
                 .try_into()
-                .unwrap(),
+                .map_err(|_| CryptoError::InvalidKeyLength)?;
-        ) as usize;
+            *off += 2;
-        offset += 2;
+            Ok(u16::from_be_bytes(arr))
+        };
+
+        let kem_len = read_u16(&mut offset)? as usize;
         let kem = KemPublicKey::new(
             bytes
                 .get(offset..offset + kem_len)
@@ -303,14 +281,7 @@ impl PublicKeyBundle {
         );
         offset += kem_len;
 
-        let pq_len = u16::from_be_bytes(
+        let pq_len = read_u16(&mut offset)? as usize;
-            bytes
-                .get(offset..offset + 2)
-                .ok_or(CryptoError::InvalidKeyLength)?
-                .try_into()
-                .unwrap(),
-        ) as usize;
-        offset += 2;
         let pq = SignaturePqPublicKey::new(
             bytes
                 .get(offset..offset + pq_len)
@@ -319,14 +290,7 @@ impl PublicKeyBundle {
         );
         offset += pq_len;
 
-        let cl_len = u16::from_be_bytes(
+        let cl_len = read_u16(&mut offset)? as usize;
-            bytes
-                .get(offset..offset + 2)
-                .ok_or(CryptoError::InvalidKeyLength)?
-                .try_into()
-                .unwrap(),
-        ) as usize;
-        offset += 2;
         let cl = SignaturePublicKey::new(
             bytes
                 .get(offset..offset + cl_len)

crypto/src/lib.rs

@@ -11,6 +11,9 @@ pub mod kdf;
 #[cfg(any(feature = "ed25519-dalek", feature = "ml-dsa"))]
 pub mod sign;
 
+#[cfg(any(feature = "ed25519-dalek", feature = "ml-dsa"))]
+pub use sign::SigAlgorithm;
+
 #[cfg(feature = "mlkem-tls")]
 pub mod kem;
 
@@ -19,7 +22,7 @@ pub mod helper;
 pub use aead::{AeadCipher, AeadDecrypt, AeadEncrypt};
 pub use error::CryptoError;
 pub use keypair::{
-    EncryptionPrivateKey, EncryptionPublicKey, KemPrivateKey, KemPublicKey, KeyGroup, Keyring,
+    EncryptionPrivateKey, EncryptionPublicKey, KemPrivateKey, KemPublicKey, Keyring,
     PublicKeyBundle, SignaturePqPrivateKey, SignaturePqPublicKey, SignaturePrivateKey,
     SignaturePublicKey,
 };
@@ -46,16 +49,248 @@ pub use hash::{sha256, sha256_double, Sha256Hasher};
 pub use kdf::{derive_encryption_key, hkdf_expand, hkdf_extract};
 
 #[cfg(feature = "mlkem-tls")]
-pub use kem::HybridKem;
+pub use kem::{Encapsulated, HybridKem};
 
 #[cfg(all(feature = "mlkem-tls", feature = "chacha20poly1305", feature = "hkdf"))]
 pub use helper::{decrypt_multi, encrypt_multi, MultiEncryptedMessage, RecipientEntry};
 
+/* ================================ TESTS ================================ */
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[cfg(feature = "chacha20poly1305")]
+    #[test]
+    fn aead_encrypt_decrypt() {
+        use crate::aead::{AeadDecrypt, AeadEncrypt, ChaCha20Poly1305};
+        let key = [0xAB; 32];
+        let cipher = ChaCha20Poly1305::new(key);
+        let ct = cipher.encrypt(b"hello world", b"aad").unwrap();
+        let pt = cipher.decrypt(&ct, b"aad").unwrap();
+        assert_eq!(pt, b"hello world");
+    }
+
+    #[cfg(feature = "chacha20poly1305")]
+    #[test]
+    fn aead_wrong_key_fails() {
+        use crate::aead::{AeadDecrypt, AeadEncrypt, ChaCha20Poly1305};
+        let cipher_a = ChaCha20Poly1305::new([0xAB; 32]);
+        let cipher_b = ChaCha20Poly1305::new([0xCD; 32]);
+        let ct = cipher_a.encrypt(b"hello", b"").unwrap();
+        assert!(cipher_b.decrypt(&ct, b"").is_err());
+    }
+
+    #[cfg(feature = "chacha20poly1305")]
+    #[test]
+    fn aead_wrong_aad_fails() {
+        use crate::aead::{AeadDecrypt, AeadEncrypt, ChaCha20Poly1305};
+        let cipher = ChaCha20Poly1305::new([0xAB; 32]);
+        let ct = cipher.encrypt(b"hello", b"correct-aad").unwrap();
+        assert!(cipher.decrypt(&ct, b"wrong-aad").is_err());
+    }
+
     #[cfg(feature = "ed25519-dalek")]
-pub use helper::verify_ed25519_sig;
+    #[test]
+    fn ed25519_sign_verify() {
+        let (signer, sk, pk) = Ed25519Signer::generate();
+        let msg = b"test message";
+        let sig = signer.sign(msg).unwrap();
+        signer.verify(msg, &sig).unwrap();
+        verify_ed25519(&pk, msg, &sig).unwrap();
+
+        let loaded = Ed25519Signer::new(&sk).unwrap();
+        loaded.verify(msg, &sig).unwrap();
+    }
+
+    #[cfg(feature = "ed25519-dalek")]
+    #[test]
+    fn ed25519_wrong_sig_fails() {
+        let (signer, _, pk) = Ed25519Signer::generate();
+        let msg = b"test message";
+        let sig = signer.sign(msg).unwrap();
+        assert!(verify_ed25519(&pk, b"wrong message", &sig).is_err());
+    }
 
     #[cfg(feature = "ml-dsa")]
-pub use helper::verify_ml_dsa_sig;
+    #[test]
+    fn mldsa_sign_verify() {
+        let (signer, sk, pk) = MlDsaSigner::generate();
+        let msg = b"test message";
+        let sig = signer.sign(msg).unwrap();
+        signer.verify(msg, &sig).unwrap();
+        verify_ml_dsa(&pk, msg, &sig).unwrap();
+
+        let loaded = MlDsaSigner::new(&sk, &pk).unwrap();
+        loaded.verify(msg, &sig).unwrap();
+    }
+
+    #[cfg(feature = "ml-dsa")]
+    #[test]
+    fn mldsa_wrong_sig_fails() {
+        let (signer, _, pk) = MlDsaSigner::generate();
+        let msg = b"test message";
+        let sig = signer.sign(msg).unwrap();
+        assert!(verify_ml_dsa(&pk, b"wrong message", &sig).is_err());
+    }
 
     #[cfg(all(feature = "ed25519-dalek", feature = "ml-dsa"))]
-pub use helper::verify_dual_sig;
+    #[test]
+    fn dual_sign_verify() {
+        use crate::sign::sign_dual;
+
+        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",
+            )
+            .unwrap();
+    }
+
+    #[cfg(all(feature = "ed25519-dalek", feature = "ml-dsa"))]
+    #[test]
+    fn dual_wrong_message_fails() {
+        use crate::sign::sign_dual;
+
+        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());
+    }
+
+    #[cfg(feature = "hkdf")]
+    #[test]
+    fn hkdf_expand_produces_key() {
+        let key = derive_encryption_key(b"ikm", b"salt", b"context").unwrap();
+        assert_eq!(key.len(), 32);
+
+        let expanded = hkdf_expand(b"ikm", b"salt", b"info", 64).unwrap();
+        assert_eq!(expanded.len(), 64);
+    }
+
+    #[cfg(feature = "hkdf")]
+    #[test]
+    fn hkdf_different_info_different_key() {
+        let a = derive_encryption_key(b"ikm", b"salt", b"info-a").unwrap();
+        let b = derive_encryption_key(b"ikm", b"salt", b"info-b").unwrap();
+        assert_ne!(a, b);
+    }
+
+    #[cfg(feature = "sha2")]
+    #[test]
+    fn sha256_hashes() {
+        let h = sha256(b"hello");
+        assert_eq!(h.len(), 32);
+
+        let h2 = sha256_double(b"hello");
+        assert_eq!(h2.len(), 32);
+        assert_ne!(h, h2);
+    }
+
+    #[cfg(feature = "sha2")]
+    #[test]
+    fn sha256_deterministic() {
+        assert_eq!(sha256(b"hello"), sha256(b"hello"));
+    }
+
+    #[cfg(feature = "sha2")]
+    #[test]
+    fn sha256_hasher_incremental() {
+        let mut hasher = Sha256Hasher::new();
+        hasher.update(b"hel");
+        hasher.update(b"lo");
+        let h = hasher.finalize();
+        assert_eq!(h, sha256(b"hello"));
+    }
+
+    #[test]
+    fn key_types_roundtrip() {
+        let data = vec![1u8, 2, 3, 4];
+        let enc_pk = EncryptionPublicKey::new(data.clone());
+        assert_eq!(enc_pk.as_bytes(), &data);
+        let enc_sk = EncryptionPrivateKey::new(data.clone());
+        assert_eq!(enc_sk.as_bytes(), &data);
+
+        let sig_pk = SignaturePublicKey::new(data.clone());
+        assert_eq!(sig_pk.as_bytes(), &data);
+        let sig_sk = SignaturePrivateKey::new(data.clone());
+        assert_eq!(sig_sk.as_bytes(), &data);
+
+        let kem_pk = KemPublicKey::new(data.clone());
+        assert_eq!(kem_pk.as_bytes(), &data);
+        let kem_sk = KemPrivateKey::new(data.clone());
+        assert_eq!(kem_sk.as_bytes(), &data);
+
+        let sig_pq_pk = SignaturePqPublicKey::new(data.clone());
+        assert_eq!(sig_pq_pk.as_bytes(), &data);
+        let sig_pq_sk = SignaturePqPrivateKey::new(data.clone());
+        assert_eq!(sig_pq_sk.as_bytes(), &data);
+    }
+
+    #[cfg(all(feature = "mlkem-tls", feature = "ml-dsa", feature = "ed25519-dalek"))]
+    #[test]
+    fn keyring_generate_and_bundle() {
+        let kr = Keyring::generate();
+        let bundle = kr.public_key_bundle();
+        assert!(!bundle.kem_public_key.as_bytes().is_empty());
+        assert!(!bundle.sig_pq_public_key.as_bytes().is_empty());
+        assert!(!bundle.sig_cl_public_key.as_bytes().is_empty());
+    }
+
+    #[cfg(all(feature = "mlkem-tls", feature = "ml-dsa", feature = "ed25519-dalek"))]
+    #[test]
+    fn keyring_serialize_roundtrip() {
+        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());
+    }
+
+    #[cfg(all(feature = "mlkem-tls", feature = "ml-dsa", feature = "ed25519-dalek"))]
+    #[test]
+    fn public_key_bundle_serialize_roundtrip() {
+        let kr = Keyring::generate();
+        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.sig_pq_public_key.as_bytes(),
+            loaded.sig_pq_public_key.as_bytes()
+        );
+        assert_eq!(
+            bundle.sig_cl_public_key.as_bytes(),
+            loaded.sig_cl_public_key.as_bytes()
+        );
+    }
+
+    #[cfg(feature = "mlkem-tls")]
+    #[test]
+    fn hybrid_kem_roundtrip() {
+        let (sk, pk) = HybridKem::generate_keypair();
+        let enc = HybridKem::encapsulate(&pk).unwrap();
+        let ss = HybridKem::decapsulate(&sk, &enc.ciphertext).unwrap();
+        assert_eq!(enc.shared_secret, ss);
+    }
+
+    #[cfg(all(feature = "mlkem-tls", feature = "chacha20poly1305", feature = "hkdf"))]
+    #[test]
+    fn encrypt_multi_roundtrip() {
+        use crate::keypair::Keyring;
+        use crate::helper::{encrypt_multi, decrypt_multi};
+
+        let kr = Keyring::generate();
+        let entities = vec![kr.public_key_bundle()];
+        let msg = b"secret data";
+        let ct = encrypt_multi(msg, b"aad", &entities).unwrap();
+        let pt = decrypt_multi(&ct, b"aad", &kr).unwrap();
+        assert_eq!(pt, msg);
+    }
+}

crypto/src/sign.rs

@@ -3,6 +3,24 @@ use crate::error::CryptoError;
 #[cfg(feature = "ed25519-dalek")]
 use crate::keypair::{SignaturePrivateKey, SignaturePublicKey};
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct SigAlgorithm;
+
+impl SigAlgorithm {
+    pub const ED25519: u8 = 0x01;
+    pub const ML_DSA_65: u8 = 0x02;
+    pub const DUAL: u8 = 0x03;
+
+    pub const fn length(alg: u8) -> Option<usize> {
+        match alg {
+            Self::ED25519 => Some(64),
+            Self::ML_DSA_65 => Some(3309),
+            Self::DUAL => Some(3373),
+            _ => None,
+        }
+    }
+}
+
 #[cfg(feature = "ed25519-dalek")]
 use rand_core::RngCore;
 
@@ -46,6 +64,14 @@ impl Ed25519Signer {
     pub fn public_key(&self) -> SignaturePublicKey {
         SignaturePublicKey::new(self.public.to_bytes().to_vec())
     }
+
+    pub fn signing_key(&self) -> &ed25519_dalek::SigningKey {
+        &self.secret
+    }
+
+    pub fn verifying_key(&self) -> &ed25519_dalek::VerifyingKey {
+        &self.public
+    }
 }
 
 #[cfg(feature = "ed25519-dalek")]
@@ -104,23 +130,26 @@ impl MlDsaSigner {
         secret_key: &SignaturePqPrivateKey,
         public_key: &SignaturePqPublicKey,
     ) -> Result<Self, CryptoError> {
-        let encoded_sk =
+        let seed_bytes: [u8; 32] = secret_key
-            ml_dsa::EncodedSigningKey::<ml_dsa::MlDsa65>::try_from(secret_key.as_bytes())
+            .as_bytes()
+            .try_into()
             .map_err(|_| CryptoError::KeyGenerationFailed)?;
-        let secret = ml_dsa::SigningKey::<ml_dsa::MlDsa65>::decode(&encoded_sk);
+        let seed = ml_dsa::Seed::from(seed_bytes);
+        let secret = ml_dsa::SigningKey::<ml_dsa::MlDsa65>::from_seed(&seed);
+
         let encoded_pk =
             ml_dsa::EncodedVerifyingKey::<ml_dsa::MlDsa65>::try_from(public_key.as_bytes())
                 .map_err(|_| CryptoError::KeyGenerationFailed)?;
         let public = ml_dsa::VerifyingKey::<ml_dsa::MlDsa65>::decode(&encoded_pk);
+
         Ok(Self { secret, public })
     }
 
     pub fn generate() -> (Self, SignaturePqPrivateKey, SignaturePqPublicKey) {
-        use ml_dsa::KeyGen;
+        use ml_dsa::{Generate, Keypair};
-        let kp = ml_dsa::MlDsa65::key_gen(&mut rand_core::OsRng);
+        let secret = ml_dsa::SigningKey::<ml_dsa::MlDsa65>::generate();
-        let secret = kp.signing_key().clone();
+        let public = secret.verifying_key();
-        let public = kp.verifying_key().clone();
+        let priv_key = SignaturePqPrivateKey::new(secret.to_seed().to_vec());
-        let priv_key = SignaturePqPrivateKey::new(secret.encode().to_vec());
         let pub_key = SignaturePqPublicKey::new(public.encode().to_vec());
         let signer = Self { secret, public };
         (signer, priv_key, pub_key)
@@ -142,20 +171,17 @@ impl MlDsaSigner {
 #[cfg(feature = "ml-dsa")]
 impl SignatureScheme for MlDsaSigner {
     fn sign(&self, msg: &[u8]) -> Result<Vec<u8>, CryptoError> {
-        use ml_dsa::signature::Signer;
+        use ml_dsa::Signer;
-        let signature = self.secret.sign(msg);
+        let signature = self.secret.try_sign(msg)
+            .map_err(|_| CryptoError::SigningFailed)?;
         Ok(signature.encode().to_vec())
     }
 
     fn verify(&self, msg: &[u8], signature: &[u8]) -> Result<(), CryptoError> {
-        use ml_dsa::signature::Verifier;
+        use ml_dsa::Verifier;
-        let encoded_sig =
+        let sig = ml_dsa::Signature::<ml_dsa::MlDsa65>::try_from(signature)
-            ml_dsa::EncodedSignature::<ml_dsa::MlDsa65>::try_from(signature)
             .map_err(|_| CryptoError::InvalidSignature)?;
-        let sig = ml_dsa::Signature::<ml_dsa::MlDsa65>::decode(&encoded_sig)
+        self.public.verify(msg, &sig)
-            .ok_or(CryptoError::InvalidSignature)?;
-        self.public
-            .verify(msg, &sig)
             .map_err(|_| CryptoError::VerificationFailed)
     }
 }
@@ -166,23 +192,22 @@ pub fn verify_ml_dsa(
     msg: &[u8],
     signature: &[u8],
 ) -> Result<(), CryptoError> {
-    use ml_dsa::signature::Verifier;
+    use ml_dsa::Verifier;
 
     let encoded_pk =
         ml_dsa::EncodedVerifyingKey::<ml_dsa::MlDsa65>::try_from(public_key.as_bytes())
             .map_err(|_| CryptoError::InvalidSignature)?;
     let public = ml_dsa::VerifyingKey::<ml_dsa::MlDsa65>::decode(&encoded_pk);
-    let encoded_sig =
+
-        ml_dsa::EncodedSignature::<ml_dsa::MlDsa65>::try_from(signature)
+    let sig = ml_dsa::Signature::<ml_dsa::MlDsa65>::try_from(signature)
         .map_err(|_| CryptoError::InvalidSignature)?;
-    let sig = ml_dsa::Signature::<ml_dsa::MlDsa65>::decode(&encoded_sig)
-        .ok_or(CryptoError::InvalidSignature)?;
 
     public
         .verify(msg, &sig)
         .map_err(|_| CryptoError::VerificationFailed)
 }
 
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub struct DualSignature {
     pub ed25519: Vec<u8>,
     pub mldsa: Vec<u8>,
@@ -194,12 +219,18 @@ pub fn sign_dual(
     mldsa_sk: &ml_dsa::SigningKey<ml_dsa::MlDsa65>,
     message: &[u8],
 ) -> DualSignature {
+    let ed25519 = {
         use ed25519_dalek::Signer;
-
+        ed25519_sk.sign(message).to_bytes().to_vec()
-    DualSignature {
+    };
-        ed25519: ed25519_sk.sign(message).to_bytes().to_vec(),
+    let mldsa = {
-        mldsa: mldsa_sk.sign(message).encode().to_vec(),
+        use ml_dsa::Signer;
-    }
+        mldsa_sk.try_sign(message)
+            .expect("ML-DSA signing failed")
+            .encode()
+            .to_vec()
+    };
+    DualSignature { ed25519, mldsa }
 }
 
 impl DualSignature {
@@ -210,22 +241,23 @@ impl DualSignature {
         mldsa_vk: &ml_dsa::VerifyingKey<ml_dsa::MlDsa65>,
         message: &[u8],
     ) -> Result<(), CryptoError> {
-        use ed25519_dalek::Verifier;
-
         let ed_sig = ed25519_dalek::Signature::from_slice(&self.ed25519)
             .map_err(|_| CryptoError::InvalidSignature)?;
+        {
+            use ed25519_dalek::Verifier;
             ed25519_vk
                 .verify(message, &ed_sig)
                 .map_err(|_| CryptoError::VerificationFailed)?;
+        }
 
-        let encoded_sig =
+        let ml_sig = ml_dsa::Signature::<ml_dsa::MlDsa65>::try_from(self.mldsa.as_slice())
-            ml_dsa::EncodedSignature::<ml_dsa::MlDsa65>::try_from(self.mldsa.as_slice())
             .map_err(|_| CryptoError::InvalidSignature)?;
-        let ml_sig = ml_dsa::Signature::<ml_dsa::MlDsa65>::decode(&encoded_sig)
+        {
-            .ok_or(CryptoError::InvalidSignature)?;
+            use ml_dsa::Verifier;
             mldsa_vk
                 .verify(message, &ml_sig)
                 .map_err(|_| CryptoError::VerificationFailed)?;
+        }
 
         Ok(())
     }

example-type-maps.yaml

@@ -33,9 +33,11 @@ protocol_version: "0.0"
 #   Description: 4
 #   Timestamp: 5
 #   Id: 6
-#   Nonce: 7
+#   ClientNonce: 7
-#   PublicKeys: 8
+#   ServerNonce: 8
-#   Signature: 9
+#   PublicKeys: 9
+#   Signature: 10
+#   Connected: 11
 #
 # If a Type can't be used it will be mapped to 0
 

example-usage/.gitignore

@@ -0,0 +1,6 @@
+client_keys.json
+host_keys.json
+host_sig_pk.bin
+host_sig_pq_pk.bin
+host_enc_kem_pk.bin
+clients.json

example-usage/client/Cargo.toml

@@ -1,13 +1,14 @@
 [package]
-name = "mtp-example-client"
+name = "client"
 version = "0.1.0"
 edition = "2024"
 
 [[bin]]
-name = "mtp-example-client"
+name = "client"
 path = "src/main.rs"
 
 [dependencies]
-mtp-client = { path = "../../client" }
+mtp = { path = "../../", features = ["client", "crypto"] }
-mtp-codec = { path = "../../codec" }
 tokio = { version = "1", features = ["full"] }
+serde_json = "1"
+hex = "0.4"

example-usage/client/src/main.rs

@@ -1,23 +1,70 @@
 use std::fs;
 
-use mtp_client::{ClientConfig, MTPClient};
+use mtp::client::{ClientConfig, MTPClient, MTPConnection};
-use mtp_codec::{CommunicationType, CommunicationValue, DataType, DataValue};
+use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
-
+use mtp::crypto::{
-#[tokio::main]
+    Ed25519Signer, KemPublicKey, Keyring, MlDsaSigner, PublicKeyBundle, SignaturePqPublicKey,
-async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    SignaturePublicKey,
-    let cert_pem = fs::read("server.pem")
-        .expect("Missing server.pem server generates it");
-
-    println!("Connecting to 127.0.0.1:8080 ...");
-
-    let config = ClientConfig {
-        url: "https://127.0.0.1:8080".into(),
-        server_cert: Some(cert_pem),
 };
 
-    let conn = MTPClient::connect(config, 1001).await?;
+/*
-    println!("Connected (version {})", conn.version);
+ * Either load a saved keyring from disk (login) or generate a fresh
+ * keypair and register with the server.  Returns the established
+ * connection and the local keyring.
+ */
+async fn connect_or_register(
+    mut config: ClientConfig,
+    host_public_key: PublicKeyBundle,
+    client_key_path: &str,
+) -> Result<(MTPConnection, Keyring), Box<dyn std::error::Error>> {
+    if let Ok(data) = fs::read_to_string(client_key_path) {
+        let json: serde_json::Value = serde_json::from_str(&data)?;
+        let client_id = json["client_id"].as_u64().expect("Invalid client_id");
+        let keyring = Keyring::from_bytes(&hex::decode(
+            json["keyring"].as_str().expect("Missing keyring"),
+        )?)?;
 
+        println!("Loaded client keys (ID: {})", client_id);
+
+        config.client_id = client_id;
+        let conn = MTPClient::auth_connect(config, &keyring, &host_public_key).await?;
+        println!("Authenticated (version {})", conn.version);
+        Ok((conn, keyring))
+    } else {
+        println!("No existing keys found: registering new client");
+
+        let (_ed_signer, sig_sk, sig_pk) = Ed25519Signer::generate();
+        let (_pq_signer, sig_pq_sk, sig_pq_pk) = MlDsaSigner::generate();
+        let keyring = Keyring::new(
+            mtp::crypto::KemPublicKey::new(vec![]),
+            mtp::crypto::KemPrivateKey::new(vec![]),
+            sig_pq_pk,
+            sig_pq_sk,
+            sig_pk,
+            sig_sk,
+        );
+
+        let keyring_bytes = keyring.to_bytes();
+        let conn = MTPClient::auth_register(config, &keyring, &host_public_key).await?;
+        println!("Registered with ID: {}", conn.client_id);
+
+        let json = serde_json::json!({
+            "client_id": conn.client_id,
+            "keyring": hex::encode(&keyring_bytes),
+        });
+        fs::write(client_key_path, serde_json::to_string_pretty(&json)?)?;
+        println!("Saved client keys -> {client_key_path}");
+
+        let keyring = Keyring::from_bytes(&keyring_bytes)?;
+        Ok((conn, keyring))
+    }
+}
+
+/*
+ * Send a Ping greeting with description, timestamp, and data fields,
+ * then print the server's Pong response.
+ */
+async fn send_ping_and_receive(conn: &MTPConnection) -> Result<(), Box<dyn std::error::Error>> {
     let timestamp = std::time::SystemTime::now()
         .duration_since(std::time::UNIX_EPOCH)
         .unwrap()
@@ -25,24 +72,51 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     let greeting = CommunicationValue::new(CommunicationType::Ping)
         .add_typed_default(DataType::Description, DataValue::Str("Hello MTP!".into()))
-        .add_typed_default(DataType::Timestamp, DataValue::UnsignedNumber(timestamp as u128))
+        .add_typed_default(
+            DataType::Timestamp,
+            DataValue::UnsignedNumber(timestamp as u128),
+        )
         .add_typed_default(DataType::Data, DataValue::UnsignedNumber(42))
-        .with_sender(1001);
+        .with_sender(conn.client_id);
 
     println!("Sending: {greeting}");
     conn.sender.send(&greeting).await?;
 
     match conn.receiver.receive().await {
-        Ok(msg) => {
+        Ok(msg) => println!("Received: {msg}"),
-            println!("Received: {msg}");
+        Err(e) => eprintln!("Receive error: {e}"),
-        }
-        Err(e) => {
-            eprintln!("Receive error: {e}");
-        }
     }
 
     conn.sender.close();
+    Ok(())
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let cert_pem = fs::read("server.pem").expect("Missing server.pem: run server first");
+    let host_public_key = PublicKeyBundle::new(
+        KemPublicKey::new(
+            fs::read("host_enc_kem_pk.bin").expect("Missing host_enc_kem_pk.bin: run server first"),
+        ),
+        SignaturePqPublicKey::new(
+            fs::read("host_sig_pq_pk.bin").expect("Missing host_sig_pq_pk.bin: run server first"),
+        ),
+        SignaturePublicKey::new(
+            fs::read("host_sig_pk.bin").expect("Missing host_sig_pk.bin: run server first"),
+        ),
+    );
+
+    println!("Connecting to 127.0.0.1:8080 ...");
+
+    let config = ClientConfig {
+        url: "https://127.0.0.1:8080".into(),
+        server_cert: Some(cert_pem),
+        client_id: 0,
+    };
+
+    let (conn, _keyring) = connect_or_register(config, host_public_key, "client_keys.json").await?;
+    send_ping_and_receive(&conn).await?;
+
     println!("\nDone");
-
     Ok(())
 }

example-usage/server/Cargo.toml

@@ -1,14 +1,16 @@
 [package]
-name = "mtp-example-server"
+name = "server"
 version = "0.1.0"
 edition = "2024"
 
 [[bin]]
-name = "mtp-example-server"
+name = "server"
 path = "src/main.rs"
 
 [dependencies]
-mtp-host = { path = "../../host" }
+mtp = { path = "../../", features = ["crypto", "host"] }
-mtp-codec = { path = "../../codec" }
 rcgen = "0.14"
 tokio = { version = "1", features = ["full"] }
+serde_json = { version = "1" }
+hex = "0.4"
+serde_core = "1.0.228"

example-usage/server/src/main.rs

@@ -1,24 +1,28 @@
+use std::collections::HashMap;
 use std::fs;
+use std::sync::{Arc, Mutex};
 
-use mtp_codec::{CommunicationType, CommunicationValue, DataType, DataValue, TypeMap};
+use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataTypeId, DataValue, TypeMap};
-use mtp_host::{HostConfig, MTPHost};
+use mtp::crypto::{Ed25519Signer, Keyring, MlDsaSigner, PublicKeyBundle, kem::HybridKem};
+use mtp::host::{HostConfig, MTPHost};
 
-#[tokio::main]
+/*
-async fn main() -> Result<(), Box<dyn std::error::Error>> {
+ * Load an existing TLS certificate and key pair from disk.  If neither
-    let cert_path = "server.pem";
+ * file exists, generate a self-signed certificate so the server is
-    let key_path = "server.key";
+ * immediately usable without external tooling.
-
+ */
-    let (cert_pem, key_pem) = if let (Ok(c), Ok(k)) = (fs::read(cert_path), fs::read(key_path)) {
+fn load_or_generate_tls(
+    cert_path: &str,
+    key_path: &str,
+) -> Result<(Vec<u8>, Vec<u8>), Box<dyn std::error::Error>> {
+    if let (Ok(c), Ok(k)) = (fs::read(cert_path), fs::read(key_path)) {
         println!("Using existing TLS cert from {cert_path}");
-        (c, k)
+        return Ok((c, k));
-    } else {
+    }
-        println!("Generating self-signed TLS certificate ...");
 
+    println!("Generating self-signed TLS certificate ...");
     let key_pair = rcgen::KeyPair::generate()?;
-        let params = rcgen::CertificateParams::new(vec![
+    let params = rcgen::CertificateParams::new(vec!["localhost".into(), "127.0.0.1".into()])?;
-            "localhost".into(),
-            "127.0.0.1".into(),
-        ])?;
     let cert = params.self_signed(&key_pair)?;
 
     let cert_str = cert.pem();
@@ -28,52 +32,152 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     fs::write(key_path, key_str.as_bytes())?;
     println!("Wrote {cert_path} and {key_path}");
 
-        (cert_str.into_bytes(), key_str.into_bytes())
+    Ok((cert_str.into_bytes(), key_str.into_bytes()))
-    };
+}
+
+/*
+ * Load the host's Ed25519 identity from a JSON file, or generate a
+ * fresh one and persist it.  Clients need the corresponding public key
+ * (exported separately as host_sig_pk.bin) to authenticate the host
+ * during the handshake.
+ */
+fn load_or_generate_host_keys(path: &str) -> Result<(u64, Keyring), Box<dyn std::error::Error>> {
+    if let Ok(data) = fs::read_to_string(path) {
+        let json: serde_json::Value = serde_json::from_str(&data)?;
+        let hid = json["host_id"].as_u64().unwrap_or(1);
+        let keyring = Keyring::from_bytes(&hex::decode(json["keyring"].as_str().unwrap())?)?;
+        println!("Loaded host keys (ID: {})", hid);
+        return Ok((hid, keyring));
+    }
+
+    let (_ed_signer, sig_sk, sig_pk) = Ed25519Signer::generate();
+    let (_pq_signer, sig_pq_sk, sig_pq_pk) = MlDsaSigner::generate();
+    let (kem_sk, kem_pk) = HybridKem::generate_keypair();
+    let keyring = Keyring::new(kem_pk, kem_sk, sig_pq_pk, sig_pq_sk, sig_pk, sig_sk);
+
+    let json = serde_json::json!({
+        "host_id": 1,
+        "keyring": hex::encode(keyring.to_bytes()),
+    });
+    fs::write(path, serde_json::to_string_pretty(&json)?)?;
+    println!("Generated host keys -> {path}");
+    Ok((1u64, keyring))
+}
+
+/*
+ * Load the client database from disk.  Each entry maps a numeric
+ * client ID to its PublicKeyBundle.  next_id starts one past the
+ * highest known ID (or 1000 if the DB is empty).
+ */
+fn load_client_db(
+    path: &str,
+) -> Result<(Arc<Mutex<HashMap<u64, PublicKeyBundle>>>, Arc<Mutex<u64>>), Box<dyn std::error::Error>>
+{
+    let clients: Arc<Mutex<HashMap<u64, PublicKeyBundle>>> =
+        Arc::new(Mutex::new(if let Ok(data) = fs::read_to_string(path) {
+            serde_json::from_str(&data).unwrap_or_default()
+        } else {
+            HashMap::new()
+        }));
+    let next_id = Arc::new(Mutex::new(
+        clients.lock().unwrap().keys().max().unwrap_or(&999) + 1,
+    ));
+    Ok((clients, next_id))
+}
+
+/*
+ * Build a Pong response carrying a description, a Unix timestamp, and
+ * a custom payload according to the negotiated type map.
+ */
+fn build_pong_response(tm: &TypeMap) -> CommunicationValue {
+    let timestamp = std::time::SystemTime::now()
+        .duration_since(std::time::UNIX_EPOCH)
+        .unwrap()
+        .as_secs();
+
+    CommunicationValue::from_comm(CommunicationType::Pong, tm)
+        .add_data(
+            DataTypeId(tm.data_id_enum(DataType::Description).unwrap()),
+            DataValue::Str("Hello from server!".into()),
+        )
+        .add_data(
+            DataTypeId(tm.data_id_enum(DataType::Timestamp).unwrap()),
+            DataValue::UnsignedNumber(timestamp as u128),
+        )
+        .add_data(
+            DataTypeId(tm.data_id_enum(DataType::Data).unwrap()),
+            DataValue::Str("custom payload".into()),
+        )
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let (cert_pem, key_pem) = load_or_generate_tls("server.pem", "server.key")?;
+    let (host_id, host_keyring) = load_or_generate_host_keys("host_keys.json")?;
+
+    // Export the host's public keys so clients can verify it
+    fs::write(
+        "host_enc_kem_pk.bin",
+        host_keyring.kem_public_key.as_bytes(),
+    )?;
+    fs::write("host_sig_pk.bin", host_keyring.sig_cl_public_key.as_bytes())?;
+    fs::write(
+        "host_sig_pq_pk.bin",
+        host_keyring.sig_pq_public_key.as_bytes(),
+    )?;
+
+    let (clients, next_id) = load_client_db("clients.json")?;
+
+    // Clone the Arc so each closure owns its own reference
+    let clients_for_get = clients.clone();
+    let get_existing_user = Box::new(move |id: u64| -> Option<PublicKeyBundle> {
+        clients_for_get.lock().unwrap().get(&id).cloned()
+    });
+
+    let clients_for_register = clients.clone();
+    let next_id_for_register = next_id.clone();
+    let clients_path = "clients.json".to_string();
+    let complete_register = Box::new(move |bundle: PublicKeyBundle| -> u64 {
+        let mut db = clients_for_register.lock().unwrap();
+        let mut nid = next_id_for_register.lock().unwrap();
+        let id = *nid;
+        *nid += 1;
+        db.insert(id, bundle);
+        fs::write(&clients_path, serde_json::to_string_pretty(&*db).unwrap()).ok();
+        println!("Registered new client with ID: {}", id);
+        id
+    });
 
     println!("Starting MTP server on port 8080 ...");
 
     let config = HostConfig {
-        ip: "::".into(),
+        ip: std::net::IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED),
         port: 8080,
         tls_fullchain: cert_pem,
         tls_key: key_pem,
+        require_authentication: true,
+        host_id,
+        host_keyring,
+        get_existing_user,
+        complete_register,
     };
 
     let mut host = MTPHost::new(config).await?;
     println!("Server listening on {}", host.local_addr());
 
     while let Some(conn) = host.accept().await {
-        println!("\n--- New connection accepted (version {}) ---", conn.version);
+        println!(
+            "\n--- New authenticated connection (version {}) ---",
+            conn.version
+        );
+        println!("Client ID: {}", conn.client_id);
 
         let tm: &TypeMap = conn.codec.registry().get(&conn.version).unwrap();
 
         match conn.receiver.receive().await {
             Ok(msg) => {
                 println!("Received: {msg}");
-
+                let response = build_pong_response(tm);
-                let timestamp = std::time::SystemTime::now()
-                    .duration_since(std::time::UNIX_EPOCH)
-                    .unwrap()
-                    .as_secs();
-
-                let response = CommunicationValue::from_comm(
-                    CommunicationType::Pong,
-                    tm,
-                )
-                .add_data(
-                    mtp_codec::DataTypeId(tm.data_id_enum(DataType::Description).unwrap()),
-                    DataValue::Str("Hello from server!".into()),
-                )
-                .add_data(
-                    mtp_codec::DataTypeId(tm.data_id_enum(DataType::Timestamp).unwrap()),
-                    DataValue::UnsignedNumber(timestamp as u128),
-                )
-                .add_data(
-                    mtp_codec::DataTypeId(tm.data_id_enum(DataType::Data).unwrap()),
-                    DataValue::Str("custom payload".into()),
-                );
-
                 println!("Sending: {response}");
                 conn.sender.send(&response).await?;
             }

host/Cargo.toml

@@ -7,4 +7,9 @@ edition = "2024"
 mtp-common = { path = "../common" }
 mtp-codec = { path = "../codec", features = ["registry"] }
 mtp-transport = { path = "../transport", features = ["host"] }
+mtp-crypto = { path = "../crypto", optional = true }
 tokio = { version = "1", features = ["full"] }
+rand = "0.8"
+
+[features]
+crypto = ["dep:mtp-crypto", "mtp-codec/crypto"]

host/src/lib.rs

@@ -1,16 +1,38 @@
+#[cfg(feature = "crypto")]
+use mtp_codec::DataType;
 use mtp_codec::{
     CommunicationValue, DataTypeId, DataValue, Version,
     registry::{Registry, VersionedCodec},
 };
 use mtp_common::CommunicationError;
 use mtp_transport::{Policy, Receiver, Sender};
+use std::net::IpAddr;
 
 // Host configuration.
 pub struct HostConfig {
-    pub ip: String,
+    pub ip: IpAddr,
     pub port: u16,
     pub tls_fullchain: Vec<u8>,
     pub tls_key: Vec<u8>,
+
+    #[cfg(feature = "crypto")]
+    pub require_authentication: bool,
+    #[cfg(feature = "crypto")]
+    pub host_id: u64,
+    #[cfg(feature = "crypto")]
+    pub host_keyring: mtp_crypto::Keyring,
+    #[cfg(feature = "crypto")]
+    pub get_existing_user: Box<dyn Fn(u64) -> Option<mtp_crypto::PublicKeyBundle> + Send>,
+    #[cfg(feature = "crypto")]
+    pub complete_register: Box<dyn Fn(mtp_crypto::PublicKeyBundle) -> u64 + Send>,
+}
+
+#[cfg(feature = "crypto")]
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum AuthState {
+    Pending,
+    Authenticated,
+    Failed,
 }
 
 // A connection that has completed version negotiation.
@@ -19,12 +41,20 @@ pub struct MTPConnection {
     pub codec: VersionedCodec,
     pub sender: Sender,
     pub receiver: Receiver,
+    #[cfg(feature = "crypto")]
+    pub auth_state: AuthState,
+    #[cfg(feature = "crypto")]
+    pub client_id: u64,
+    #[cfg(feature = "crypto")]
+    pub client_public_key: Option<mtp_crypto::PublicKeyBundle>,
 }
 
 // High-level MTP host with built-in version negotiation.
 pub struct MTPHost {
     transport: mtp_transport::Host,
     registry: Registry,
+    #[cfg(feature = "crypto")]
+    config: HostConfig,
 }
 
 impl MTPHost {
@@ -32,9 +62,10 @@ impl MTPHost {
         let registry = Registry::builtin();
 
         let transport = mtp_transport::host(
+            config.ip,
             config.port,
-            config.tls_fullchain,
+            config.tls_fullchain.clone(),
-            config.tls_key,
+            config.tls_key.clone(),
             Policy::default(),
         )
         .await?;
@@ -42,6 +73,8 @@ impl MTPHost {
         Ok(Self {
             transport,
             registry,
+            #[cfg(feature = "crypto")]
+            config,
         })
     }
 
@@ -55,7 +88,12 @@ impl MTPHost {
     pub async fn accept(&mut self) -> Option<MTPConnection> {
         let (sender, receiver) = self.transport.next().await?;
 
-        /* Read the first message (always encoded with reserved types). */
+        #[cfg(feature = "crypto")]
+        if self.config.require_authentication {
+            return self.accept_authenticated(sender, receiver).await;
+        }
+
+        // Read the first message (always encoded with reserved types).
         let first_msg = receiver.receive().await.ok()?;
 
         /*
@@ -76,6 +114,12 @@ impl MTPHost {
             codec,
             sender,
             receiver,
+            #[cfg(feature = "crypto")]
+            auth_state: AuthState::Authenticated,
+            #[cfg(feature = "crypto")]
+            client_id: 0,
+            #[cfg(feature = "crypto")]
+            client_public_key: None,
         })
     }
 
@@ -88,6 +132,226 @@ impl MTPHost {
     }
 }
 
+#[cfg(feature = "crypto")]
+impl MTPHost {
+    async fn accept_authenticated(
+        &mut self,
+        sender: Sender,
+        receiver: Receiver,
+    ) -> Option<MTPConnection> {
+        use mtp_crypto::{
+            Ed25519Signer, PublicKeyBundle, SignatureScheme, 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)) {
+            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) = if msg.get_type() == mtp_codec::CommunicationTypeId(15) {
+            // LOGIN
+            let cid = match msg.get_data(DataTypeId(6)) {
+                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),
+                    );
+                    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)
+        } else if msg.get_type() == mtp_codec::CommunicationTypeId(17) {
+            // REGISTER
+            let bundle = match msg.get_data(DataTypeId(9)) {
+                DataValue::Bytes(b) => PublicKeyBundle::from_bytes(b).ok()?,
+                _ => {
+                    sender.close();
+                    return None;
+                }
+            };
+            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)
+        } else {
+            sender.close();
+            return None;
+        };
+
+        // 2. Send success response (single host message)
+        let new_nonce: u128 = rand::random();
+
+        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());
+
+        let host_signer = Ed25519Signer::new(&self.config.host_keyring.sig_cl_secret_key).ok()?;
+
+        /* ===== Signature ===== */
+        let host_sig = host_signer.sign(&host_sig_payload).ok()?;
+
+        let mut response =
+            CommunicationValue::new(mtp_codec::CommunicationType::IdentificationResponse)
+                .add_typed_default(DataType::Connected, DataValue::BoolTrue)
+                .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()?;
+            response =
+                response.add_typed_default(DataType::PqSignature, DataValue::Bytes(host_pq_sig));
+        }
+        /* ===== End Signature ===== */
+
+        sender.send(&response).await.ok()?;
+
+        // 3. Version negotiation
+        let negotiated = self.registry.negotiate(&[client_version])?;
+        let codec = VersionedCodec::new(self.registry.clone());
+
+        Some(MTPConnection {
+            version: negotiated,
+            codec,
+            sender,
+            receiver,
+            auth_state: AuthState::Authenticated,
+            client_id: assigned_id,
+            client_public_key: Some(client_bundle),
+        })
+    }
+}
+
 /*
  * Extract the protocol version from an initial `CommunicationValue`.
  *
@@ -110,7 +374,10 @@ mod tests {
     #[test]
     fn version_extraction() {
         let tm = mtp_codec::TypeMap::latest();
-        let msg = CommunicationValue::from_comm(mtp_codec::CommunicationType::Identification, &tm)
+        let msg = mtp_codec::CommunicationValue::from_comm(
+            mtp_codec::CommunicationType::Identification,
+            &tm,
+        )
         .add_data(DataTypeId(3), DataValue::Str("2.0".to_string()));
         let version = extract_version(&msg);
         assert_eq!(version, Some(Version(2, 0)));
@@ -119,14 +386,20 @@ mod tests {
     #[test]
     fn version_extraction_returns_none_for_missing() {
         let tm = mtp_codec::TypeMap::latest();
-        let msg = CommunicationValue::from_comm(mtp_codec::CommunicationType::Identification, &tm);
+        let msg = mtp_codec::CommunicationValue::from_comm(
+            mtp_codec::CommunicationType::Identification,
+            &tm,
+        );
         assert!(extract_version(&msg).is_none());
     }
 
     #[test]
     fn version_extraction_bad_format() {
         let tm = mtp_codec::TypeMap::latest();
-        let msg = CommunicationValue::from_comm(mtp_codec::CommunicationType::Identification, &tm)
+        let msg = mtp_codec::CommunicationValue::from_comm(
+            mtp_codec::CommunicationType::Identification,
+            &tm,
+        )
         .add_data(DataTypeId(3), DataValue::UnsignedNumber(42));
         assert!(extract_version(&msg).is_none());
     }

registry/Cargo.toml

@@ -1,8 +0,0 @@
-[package]
-name = "registry"
-version = "0.1.0"
-edition = "2024"
-
-[dependencies]
-mtp-common = { path = "../common" }
-mtp-type-map = { path = "../type-map", features = ["multi-version"] }

registry/src/lib.rs

@@ -1,127 +0,0 @@
-use std::collections::BTreeMap;
-
-pub use mtp_type_map::{
-    CommunicationType, CommunicationTypeId, DataType, DataTypeId, TypeMap, Version,
-};
-
-/*
- * Reserved communication & data-type namespace (0..32).
- * These are fixed across all protocol versions for version negotiation
- * and security. They are never assigned by user configuration.
- */
-pub const INTERNAL_COMM_RESERVED: std::ops::Range<u16> = 0..32;
-pub const INTERNAL_DATA_RESERVED: std::ops::Range<u16> = 0..32;
-
-/*
- * Multi-version type-map registry.
- *
- * Stores one `TypeMap` per protocol version and provides version
- * negotiation for the host. Hosts build this from the compiled-in
- * type maps (via `Registry::builtin()` or programmatically).
- */
-#[derive(Clone, Debug)]
-pub struct Registry {
-    versions: BTreeMap<Version, TypeMap>,
-}
-
-impl Registry {
-    pub fn new() -> Self {
-        Self {
-            versions: BTreeMap::new(),
-        }
-    }
-
-    // Register a type map, keyed by its version.
-    pub fn register(&mut self, typemap: TypeMap) {
-        self.versions.insert(typemap.version.clone(), typemap);
-    }
-
-    // Look up the type map for an exact version.
-    pub fn get(&self, version: &Version) -> Option<&TypeMap> {
-        self.versions.get(version)
-    }
-
-    // Check if a specific version is supported.
-    pub fn supports(&self, version: &Version) -> bool {
-        self.versions.contains_key(version)
-    }
-
-    /*
-     * Negotiate the highest mutually supported version from a list.
-     *
-     * For single-version clients (the typical case) this is equivalent
-     * to checking `supports(&client_versions[0])`.
-     */
-    pub fn negotiate(&self, client_versions: &[Version]) -> Option<Version> {
-        client_versions
-            .iter()
-            .filter(|v| self.versions.contains_key(v))
-            .max()
-            .cloned()
-    }
-
-    // Return the latest (highest) registered version's type map.
-    pub fn latest(&self) -> Option<&TypeMap> {
-        self.versions.last_key_value().map(|(_, v)| v)
-    }
-
-    /*
-     * Build the registry from the compiled-in configuration.
-     *
-     * Uses the `TypeMap::vX_Y()` constructors generated by
-     * `type-map/build.rs` from the YAML config.
-     */
-    pub fn builtin() -> Self {
-        let mut r = Self::new();
-        for tm in mtp_type_map::builtin_type_maps() {
-            r.register(tm);
-        }
-        r
-    }
-}
-
-impl Default for Registry {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
-/* ================================ TESTS ================================ */
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn builtin_contains_versions() {
-        let r = Registry::builtin();
-        // At least the reserved-only fallback should exist
-        assert!(r.supports(&Version(0, 0)));
-    }
-
-    #[test]
-    fn negotiate_finds_highest() {
-        let mut r = Registry::new();
-        r.register(TypeMap::new(Version(0, 0)));
-        r.register(TypeMap::new(Version(1, 0)));
-        r.register(TypeMap::new(Version(2, 0)));
-
-        let client = &[Version(0, 0)];
-        assert_eq!(r.negotiate(client), Some(Version(0, 0)));
-
-        let client = &[Version(1, 0), Version(0, 0)];
-        assert_eq!(r.negotiate(client), Some(Version(1, 0)));
-
-        let client = &[Version(5, 0)];
-        assert_eq!(r.negotiate(client), None);
-    }
-
-    #[test]
-    fn latest_returns_highest() {
-        let mut r = Registry::new();
-        r.register(TypeMap::new(Version(0, 0)));
-        r.register(TypeMap::new(Version(2, 0)));
-        r.register(TypeMap::new(Version(1, 0)));
-
-        assert_eq!(r.latest().unwrap().version, Version(2, 0));
-    }
-}

src/lib.rs

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

transport/Cargo.toml

@@ -11,14 +11,18 @@ wtransport = { version = "0.7.1", default-features = false, features = [
     "quinn",
     "self-signed",
 ] }
-rustls = { version = "0.23.40" }
+rustls = { version = "0.23.41" }
-quinn = { version = "0.11.9", default-features = false, features = [
+quinn = { version = "0.11.11", default-features = false, features = [
     "rustls-aws-lc-rs",
     "rustls",
 ] }
 tokio = { version = "1", features = ["full"] }
 thiserror = "2.0.18"
 rustls-native-certs = "0.8.4"
+log = "0.4"
+
+[dev-dependencies]
+rcgen = "0.14"
 
 [features]
 default = []

transport/src/connection.rs

@@ -53,7 +53,6 @@ impl Default for Policy {
     }
 }
 
-#[allow(unused)]
 enum ReceivedFrame {
     Message(CommunicationValue),
     ClosedByPeer,
@@ -215,7 +214,7 @@ impl Sender {
             .await
             .map_err(|_| CommunicationError::StreamError)?
         {
-            println!("[Sender] close frame finish failed: {e}");
+            log::warn!("[Sender] close frame finish failed: {e}");
         }
 
         Ok(())
@@ -445,7 +444,7 @@ impl Receiver {
                     if e.kind() == ErrorKind::UnexpectedEof {
                         return Ok(ReceivedFrame::Idle);
                     }
-                    println!("[Receiver] read_u32 failed: {e}");
+                    log::warn!("[Receiver] read_u32 failed: {e}");
                     return Err(CommunicationError::StreamError);
                 }
             }
@@ -468,13 +467,13 @@ impl Receiver {
                 _ => return Err(e.into()),
             },
             Err(_) => {
-                println!("[Receiver] read_exact timed out (len={})", len);
+                log::warn!("[Receiver] read_exact timed out (len={})", len);
                 return Err(CommunicationError::StreamError);
             }
         }
 
         let message = CommunicationValue::from_bytes(&buf)
-            .ok_or(CommunicationError::ParseCommunicationValue)?;
+            .map_err(|_| CommunicationError::ParseCommunicationValue)?;
 
         Ok(ReceivedFrame::Message(message))
     }
@@ -517,3 +516,48 @@ impl Receiver {
         self.handle.close_reason()
     }
 }
+
+/* ================================ TESTS ================================ */
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_send_mode_derive() {
+        assert_eq!(SendMode::PersistentStream, SendMode::PersistentStream);
+        assert_ne!(SendMode::PersistentStream, SendMode::SingleStreamPerMessage);
+    }
+
+    #[test]
+    fn test_policy_default_values() {
+        let p = Policy::default();
+        assert_eq!(p.send_mode, SendMode::PersistentStream);
+        assert_eq!(p.max_message_size, 1_000_000_000);
+        assert_eq!(p.close_frame_len, u32::MAX);
+        assert_eq!(p.application_close_code, 0);
+        assert_eq!(p.open_stream_timeout, Duration::from_millis(2_000));
+        assert_eq!(p.write_timeout, Duration::from_millis(2_000));
+        assert_eq!(p.accept_stream_timeout, Duration::from_millis(10_000));
+        assert_eq!(p.read_timeout, Duration::from_millis(30_000));
+        assert_eq!(p.keep_alive_interval, Some(Duration::from_secs(3)));
+        assert_eq!(p.max_idle_timeout, Some(Duration::from_secs(30)));
+        assert_eq!(p.force_close_delay, Duration::from_millis(300));
+        assert_eq!(p.max_transient_recv_errors, 20);
+        assert_eq!(p.transient_recv_backoff, Duration::from_millis(100));
+        assert_eq!(p.receiver_queue_capacity, 1000);
+    }
+
+    #[test]
+    fn test_policy_clone() {
+        let p = Policy::default();
+        let cloned = p.clone();
+        assert_eq!(p.send_mode, cloned.send_mode);
+    }
+
+    #[test]
+    fn test_policy_debug() {
+        let p = Policy::default();
+        let debug_str = format!("{:?}", p);
+        assert!(debug_str.contains("Policy"));
+    }
+}

transport/src/connection_handle.rs

@@ -63,3 +63,70 @@ impl Default for ConnectionHandle {
         Self::new()
     }
 }
+
+/* ================================ TESTS ================================ */
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_new_is_open() {
+        let h = ConnectionHandle::new();
+        assert!(h.is_open());
+        assert!(!h.is_closed());
+    }
+
+    #[test]
+    fn test_close_transitions_state() {
+        let h = ConnectionHandle::new();
+        h.close(Some(CommunicationError::StreamClosed));
+        assert!(!h.is_open());
+        assert!(h.is_closed());
+    }
+
+    #[test]
+    fn test_close_reason_some() {
+        let h = ConnectionHandle::new();
+        h.close(Some(CommunicationError::UseAfterClosed));
+        assert!(h.close_reason().is_some());
+    }
+
+    #[test]
+    fn test_close_reason_none() {
+        let h = ConnectionHandle::new();
+        h.close(None);
+        assert!(h.close_reason().is_none());
+    }
+
+    #[test]
+    fn test_default_is_new() {
+        let h = ConnectionHandle::default();
+        assert!(h.is_open());
+    }
+
+    #[test]
+    fn test_close_with_error() {
+        let h = ConnectionHandle::new();
+        h.close_with_error(CommunicationError::MessageTooLarge);
+        assert!(h.is_closed());
+        assert!(h.close_reason().is_some());
+    }
+
+    #[test]
+    fn test_multiple_close_first_wins() {
+        let h = ConnectionHandle::new();
+        h.close(Some(CommunicationError::StreamClosed));
+        h.close(Some(CommunicationError::UseAfterClosed));
+        // First close reason is preserved
+        assert!(h.close_reason().is_some());
+    }
+
+    #[test]
+    fn test_close_sends_reason() {
+        let h = ConnectionHandle::new();
+        let mut rx = h.subscribe_close();
+        h.close(Some(CommunicationError::ClosedLocally));
+        // After close, the watch channel is updated
+        assert!(rx.borrow_and_update().is_some());
+    }
+}

transport/src/host.rs

@@ -1,4 +1,5 @@
 use crate::{ConnectionHandle, Policy, Receiver, Sender};
+use log;
 use mtp_common::CommunicationError;
 use rustls::pki_types::{PrivateKeyDer, pem::PemObject};
 use std::net::{IpAddr, SocketAddr};
@@ -13,7 +14,13 @@ pub struct Host {
 
 impl Host {
     pub async fn next(&mut self) -> Option<(Sender, Receiver)> {
-        self.incoming.recv().await
+        log::warn!("[transport Host::next] waiting on recv...");
+        let result = self.incoming.recv().await;
+        match &result {
+            Some(_) => log::warn!("[transport Host::next] received connection"),
+            None => log::warn!("[transport Host::next] incoming channel closed - sender dropped"),
+        }
+        result
     }
 
     pub fn local_addr(&self) -> std::net::SocketAddr {
@@ -22,6 +29,7 @@ impl Host {
 }
 
 pub async fn host(
+    ip: IpAddr,
     port: u16,
     cert_pem: Vec<u8>,
     key_pem: Vec<u8>,
@@ -29,7 +37,7 @@ pub async fn host(
 ) -> Result<Host, CommunicationError> {
     let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
 
-    let server_config = configure_server(port, cert_pem, key_pem, &policy).await?;
+    let server_config = configure_server(ip, port, cert_pem, key_pem, &policy).await?;
     let endpoint = Endpoint::server(server_config)
         .map_err(|e| CommunicationError::Other(format!("Endpoint creation failed: {}", e)))?;
 
@@ -42,25 +50,37 @@ pub async fn host(
     let policy = Arc::new(policy);
 
     let task = tokio::spawn(async move {
+        log::warn!("[transport bg task] started");
         loop {
+            log::warn!("[transport bg task] waiting for connection...");
             let incoming_session = endpoint.accept().await;
+            log::warn!("[transport bg task] got incoming session");
 
             let request = match incoming_session.await {
-                Ok(req) => req,
+                Ok(req) => {
-                Err(_) => {
+                    log::warn!("[transport bg task] got request");
+                    req
+                }
+                Err(e) => {
+                    log::warn!("[transport bg task] incoming session error: {e}");
                     continue;
                 }
             };
 
             let connection = match request.accept().await {
-                Ok(conn) => conn,
+                Ok(conn) => {
-                Err(_) => {
+                    log::warn!("[transport bg task] connection accepted");
+                    conn
+                }
+                Err(e) => {
+                    log::warn!("[transport bg task] accept error: {e}");
                     continue;
                 }
             };
 
             let incoming_tx = incoming_tx.clone();
             let policy = policy.clone();
+            eprintln!("[transport bg task] spawning handle_connection");
             tokio::spawn(handle_connection(connection, incoming_tx, policy));
         }
     });
@@ -85,6 +105,7 @@ async fn handle_connection(
 }
 
 async fn configure_server(
+    bind_ip: IpAddr,
     port: u16,
     cert_pem: Vec<u8>,
     key_pem: Vec<u8>,
@@ -104,11 +125,6 @@ async fn configure_server(
 
     tls_config.alpn_protocols = vec![b"h3".to_vec()];
 
-    let bind_ip = std::env::var("mtp_BIND")
-        .ok()
-        .and_then(|s| if s.is_empty() { None } else { Some(s) })
-        .unwrap_or_else(|| "::".to_string())
-        .parse::<IpAddr>()?;
     let bind_addr = SocketAddr::new(bind_ip, port);
 
     let server_config = ServerConfig::builder()

transport/tests/integration.rs

@@ -0,0 +1,174 @@
+use std::net::{IpAddr, Ipv4Addr};
+
+use mtp_transport::{Policy, host, connect};
+
+fn generate_self_signed_cert() -> (Vec<u8>, Vec<u8>) {
+    let key_pair = rcgen::KeyPair::generate().unwrap();
+    let params =
+        rcgen::CertificateParams::new(vec!["localhost".into(), "127.0.0.1".into()]).unwrap();
+    let cert = params.self_signed(&key_pair).unwrap();
+    let cert_pem = cert.pem();
+    let key_pem = key_pair.serialize_pem();
+    (cert_pem.into_bytes(), key_pem.into_bytes())
+}
+
+#[tokio::test]
+async fn test_host_start_and_stop() {
+    let (cert_pem, key_pem) = generate_self_signed_cert();
+    let h = host(
+        IpAddr::V4(Ipv4Addr::LOCALHOST),
+        0,
+        cert_pem,
+        key_pem,
+        Policy::default(),
+    )
+    .await
+    .unwrap();
+    let addr = h.local_addr();
+    // Port should be non-zero (OS-assigned)
+    assert!(addr.port() > 0);
+}
+
+#[tokio::test]
+async fn test_send_receive_roundtrip() {
+    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();
+
+    // Accept on host side
+    let (host_tx, host_rx) = h.next().await.unwrap();
+
+    // Client sends a simple message
+    let msg = mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::Ping)
+        .add_data(
+            mtp_codec::DataTypeId(6),
+            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!(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),
+            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_val, mtp_codec::DataValue::UnsignedNumber(99));
+
+    // Close both sides
+    client_tx.close();
+    host_tx.close();
+}
+
+#[tokio::test]
+async fn test_concurrent_messages() {
+    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();
+
+    // 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),
+                mtp_codec::DataValue::UnsignedNumber(i),
+            );
+        client_tx.send(&msg).await.unwrap();
+    }
+
+    // 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();
+        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),
+                mtp_codec::DataValue::UnsignedNumber(i * 10),
+            );
+        client_tx.send(&msg).await.unwrap();
+    }
+
+    for i in 0..3u128 {
+        let received = host_rx.receive().await.unwrap();
+        let val = received.get_data(mtp_codec::DataTypeId(6)).clone();
+        assert_eq!(val, mtp_codec::DataValue::UnsignedNumber(i * 10));
+    }
+
+    client_tx.close();
+}
+
+#[tokio::test]
+async fn test_close_detection() {
+    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();
+
+    // Send a message then close
+    let msg = mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::Ping);
+    client_tx.send(&msg).await.unwrap();
+    client_tx.close();
+
+    // Host should still receive the message
+    let received = host_rx.receive().await.unwrap();
+    assert_eq!(received.get_type(), mtp_codec::CommunicationTypeId(19)); // Ping
+
+    // Host should get an error or closed signal on next receive
+    let result = host_rx.receive().await;
+    assert!(result.is_err());
+}

type-map/Cargo.toml

@@ -6,7 +6,9 @@ build = "build.rs"
 
 [features]
 default = []
-multi-version = []
+# Enables multi-version type-map constructors, builtin_type_maps(), and
+# the Registry struct for version negotiation. Used by host, not client.
+registry = []
 
 [dependencies]
 

type-map/build.rs

@@ -24,56 +24,160 @@ struct ReservedEntry {
 }
 
 const RESERVED_COMM_TYPES: &[ReservedEntry] = &[
-    ReservedEntry { name: "Error", id: 0 },
+    ReservedEntry {
-    ReservedEntry { name: "ErrorParsing", id: 1 },
+        name: "Error",
-    ReservedEntry { name: "ErrorBadVersion", id: 2 },
+        id: 0,
-    ReservedEntry { name: "Disconnect", id: 3 },
+    },
-    ReservedEntry { name: "Redirect", id: 4 },
+    ReservedEntry {
-    ReservedEntry { name: "Shutdown", id: 5 },
+        name: "ErrorParsing",
-    ReservedEntry { name: "BadRequest", id: 6 },
+        id: 1,
-    ReservedEntry { name: "Unauthorized", id: 7 },
+    },
-    ReservedEntry { name: "Forbidden", id: 8 },
+    ReservedEntry {
-    ReservedEntry { name: "NotFound", id: 9 },
+        name: "ErrorBadVersion",
-    ReservedEntry { name: "TooManyRequests", id: 10 },
+        id: 2,
-    ReservedEntry { name: "InternalServerError", id: 11 },
+    },
-    ReservedEntry { name: "BadGateway", id: 12 },
+    ReservedEntry {
-    ReservedEntry { name: "ServiceUnavailable", id: 13 },
+        name: "Disconnect",
-    ReservedEntry { name: "GatewayTimeout", id: 14 },
+        id: 3,
-    ReservedEntry { name: "Identification", id: 15 },
+    },
-    ReservedEntry { name: "IdentificationResponse", id: 16 },
+    ReservedEntry {
-    ReservedEntry { name: "Register", id: 17 },
+        name: "Redirect",
-    ReservedEntry { name: "RegisterResponse", id: 18 },
+        id: 4,
-    ReservedEntry { name: "Ping", id: 19 },
+    },
-    ReservedEntry { name: "Pong", id: 20 },
+    ReservedEntry {
+        name: "Shutdown",
+        id: 5,
+    },
+    ReservedEntry {
+        name: "BadRequest",
+        id: 6,
+    },
+    ReservedEntry {
+        name: "Unauthorized",
+        id: 7,
+    },
+    ReservedEntry {
+        name: "Forbidden",
+        id: 8,
+    },
+    ReservedEntry {
+        name: "NotFound",
+        id: 9,
+    },
+    ReservedEntry {
+        name: "TooManyRequests",
+        id: 10,
+    },
+    ReservedEntry {
+        name: "InternalServerError",
+        id: 11,
+    },
+    ReservedEntry {
+        name: "BadGateway",
+        id: 12,
+    },
+    ReservedEntry {
+        name: "ServiceUnavailable",
+        id: 13,
+    },
+    ReservedEntry {
+        name: "GatewayTimeout",
+        id: 14,
+    },
+    ReservedEntry {
+        name: "Identification",
+        id: 15,
+    },
+    ReservedEntry {
+        name: "IdentificationResponse",
+        id: 16,
+    },
+    ReservedEntry {
+        name: "Register",
+        id: 17,
+    },
+    ReservedEntry {
+        name: "RegisterResponse",
+        id: 18,
+    },
+    ReservedEntry {
+        name: "Ping",
+        id: 19,
+    },
+    ReservedEntry {
+        name: "Pong",
+        id: 20,
+    },
 ];
 
 const RESERVED_DATA_TYPES: &[ReservedEntry] = &[
-    ReservedEntry { name: "Error", id: 0 },
+    ReservedEntry {
-    ReservedEntry { name: "ErrorParsing", id: 1 },
+        name: "Error",
-    ReservedEntry { name: "ErrorMessage", id: 2 },
+        id: 0,
-    ReservedEntry { name: "Version", id: 3 },
+    },
-    ReservedEntry { name: "Description", id: 4 },
+    ReservedEntry {
-    ReservedEntry { name: "Timestamp", id: 5 },
+        name: "ErrorParsing",
+        id: 1,
+    },
+    ReservedEntry {
+        name: "ErrorMessage",
+        id: 2,
+    },
+    ReservedEntry {
+        name: "Version",
+        id: 3,
+    },
+    ReservedEntry {
+        name: "Description",
+        id: 4,
+    },
+    ReservedEntry {
+        name: "Timestamp",
+        id: 5,
+    },
     ReservedEntry { name: "Id", id: 6 },
-    ReservedEntry { name: "Nonce", id: 7 },
+    ReservedEntry {
-    ReservedEntry { name: "PublicKeys", id: 8 },
+        name: "ClientNonce",
-    ReservedEntry { name: "Signature", id: 9 },
+        id: 7,
+    },
+    ReservedEntry {
+        name: "ServerNonce",
+        id: 8,
+    },
+    ReservedEntry {
+        name: "PublicKeys",
+        id: 9,
+    },
+    ReservedEntry {
+        name: "Signature",
+        id: 10,
+    },
+    ReservedEntry {
+        name: "Connected",
+        id: 11,
+    },
+    ReservedEntry {
+        name: "PqSignature",
+        id: 12,
+    },
 ];
 
 fn main() {
     let out = std::path::PathBuf::from(std::env::var("OUT_DIR").unwrap());
-    let multi_version = std::env::var("CARGO_FEATURE_MULTI_VERSION").is_ok();
+    let multi_version = std::env::var("CARGO_FEATURE_REGISTRY").is_ok();
+
+    println!("cargo:rerun-if-env-changed=MTP_TYPE_MAPS");
 
     let config = match std::env::var("MTP_TYPE_MAPS") {
         Ok(config_path) => {
+            println!("cargo:rerun-if-changed={}", config_path);
+
             let content =
                 std::fs::read_to_string(&config_path).expect("Failed to read type-maps.yaml");
             serde_yaml::from_str(&content).expect("Failed to parse type-maps.yaml")
         }
         Err(_) => {
-            eprintln!(
+            eprint!("warning: MTP_TYPE_MAPS not set; generating types with reserved entries only");
-                "warning: MTP_TYPE_MAPS not set; generating types with reserved entries only"
-            );
             Config {
                 protocol_version: String::new(),
                 type_maps: BTreeMap::new(),
@@ -251,6 +355,33 @@ fn generate_comm_type_enum(out: &mut String, user_names: &BTreeSet<&str>) {
     writeln!(out, "}}").unwrap();
     writeln!(out).unwrap();
 
+    writeln!(out, "impl CommunicationType {{").unwrap();
+    writeln!(out, "    pub fn from_name(s: &str) -> Option<Self> {{").unwrap();
+    writeln!(out, "        match s {{").unwrap();
+
+    for entry in RESERVED_COMM_TYPES {
+        writeln!(
+            out,
+            "            \"{}\" => Some(CommunicationType::{}),",
+            entry.name, entry.name
+        )
+        .unwrap();
+    }
+    for name in user_names {
+        writeln!(
+            out,
+            "            \"{}\" => Some(CommunicationType::{}),",
+            name, name
+        )
+        .unwrap();
+    }
+
+    writeln!(out, "            _ => None,").unwrap();
+    writeln!(out, "        }}").unwrap();
+    writeln!(out, "    }}").unwrap();
+    writeln!(out, "}}").unwrap();
+    writeln!(out).unwrap();
+
     writeln!(out, "impl std::fmt::Display for CommunicationType {{").unwrap();
     writeln!(
         out,
@@ -303,6 +434,28 @@ fn generate_data_type_enum(out: &mut String, user_names: &BTreeSet<&str>) {
     writeln!(out, "}}").unwrap();
     writeln!(out).unwrap();
 
+    writeln!(out, "impl DataType {{").unwrap();
+    writeln!(out, "    pub fn from_name(s: &str) -> Option<Self> {{").unwrap();
+    writeln!(out, "        match s {{").unwrap();
+
+    for entry in RESERVED_DATA_TYPES {
+        writeln!(
+            out,
+            "            \"{}\" => Some(DataType::{}),",
+            entry.name, entry.name
+        )
+        .unwrap();
+    }
+    for name in user_names {
+        writeln!(out, "            \"{}\" => Some(DataType::{}),", name, name).unwrap();
+    }
+
+    writeln!(out, "            _ => None,").unwrap();
+    writeln!(out, "        }}").unwrap();
+    writeln!(out, "    }}").unwrap();
+    writeln!(out, "}}").unwrap();
+    writeln!(out).unwrap();
+
     writeln!(out, "impl std::fmt::Display for DataType {{").unwrap();
     writeln!(
         out,
@@ -429,11 +582,7 @@ fn generate_single_version_lookup(out: &mut String, config: &Config) {
     )
     .unwrap();
     writeln!(out, "        match self.version {{").unwrap();
-    writeln!(
+    writeln!(out, "            PROTOCOL_VERSION => match ct {{").unwrap();
-        out,
-        "            PROTOCOL_VERSION => match ct {{"
-    )
-    .unwrap();
     for entry in RESERVED_COMM_TYPES {
         writeln!(
             out,
@@ -466,11 +615,7 @@ fn generate_single_version_lookup(out: &mut String, config: &Config) {
     )
     .unwrap();
     writeln!(out, "        match self.version {{").unwrap();
-    writeln!(
+    writeln!(out, "            PROTOCOL_VERSION => match dt {{").unwrap();
-        out,
-        "            PROTOCOL_VERSION => match dt {{"
-    )
-    .unwrap();
     for entry in RESERVED_DATA_TYPES {
         writeln!(
             out,
@@ -529,7 +674,11 @@ fn generate_reverse_lookups(
         }
     }
 
-    writeln!(out, "pub fn communication_type_name(id: u16) -> Option<&'static str> {{").unwrap();
+    writeln!(
+        out,
+        "pub fn communication_type_name(id: u16) -> Option<&'static str> {{"
+    )
+    .unwrap();
     writeln!(out, "    match id {{").unwrap();
     for (id, name) in &id_to_comm {
         writeln!(out, "        {} => Some(\"{}\"),", id, name).unwrap();
@@ -557,7 +706,11 @@ fn generate_reverse_lookups(
         }
     }
 
-    writeln!(out, "pub fn data_type_name(id: u16) -> Option<&'static str> {{").unwrap();
+    writeln!(
+        out,
+        "pub fn data_type_name(id: u16) -> Option<&'static str> {{"
+    )
+    .unwrap();
     writeln!(out, "    match id {{").unwrap();
     for (id, name) in &id_to_data {
         writeln!(out, "        {} => Some(\"{}\"),", id, name).unwrap();
@@ -570,7 +723,11 @@ fn generate_reverse_lookups(
 
 fn generate_enum_conversion_methods(out: &mut String) {
     writeln!(out, "impl CommunicationType {{").unwrap();
-    writeln!(out, "    pub fn to_id(self, tm: &TypeMap) -> CommunicationTypeId {{").unwrap();
+    writeln!(
+        out,
+        "    pub fn to_id(self, tm: &TypeMap) -> CommunicationTypeId {{"
+    )
+    .unwrap();
     writeln!(
         out,
         "        CommunicationTypeId(tm.comm_id_enum(self).expect(\"CommunicationType not found in this TypeMap version\"))"

type-map/src/lib.rs

@@ -1,9 +1,40 @@
+pub const INTERNAL_COMM_RESERVED: std::ops::Range<u16> = 0..32;
+pub const INTERNAL_DATA_RESERVED: std::ops::Range<u16> = 0..32;
+
 #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub struct CommunicationTypeId(pub u16);
 
+impl CommunicationTypeId {
+    pub fn try_new(id: u16) -> Result<Self, ()> {
+        if INTERNAL_COMM_RESERVED.contains(&id) {
+            Err(())
+        } else {
+            Ok(Self(id))
+        }
+    }
+
+    pub fn is_reserved(&self) -> bool {
+        INTERNAL_COMM_RESERVED.contains(&self.0)
+    }
+}
+
 #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub struct DataTypeId(pub u16);
 
+impl DataTypeId {
+    pub fn try_new(id: u16) -> Result<Self, ()> {
+        if INTERNAL_DATA_RESERVED.contains(&id) {
+            Err(())
+        } else {
+            Ok(Self(id))
+        }
+    }
+
+    pub fn is_reserved(&self) -> bool {
+        INTERNAL_DATA_RESERVED.contains(&self.0)
+    }
+}
+
 #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub struct Version(pub u16, pub u16);
 
@@ -43,3 +74,110 @@ impl TypeMap {
 }
 
 include!(concat!(env!("OUT_DIR"), "/types.rs"));
+
+/* ============================= REGISTRY ============================= */
+#[cfg(feature = "registry")]
+pub use registry::*;
+
+#[cfg(feature = "registry")]
+mod registry {
+    use crate::{TypeMap, Version, builtin_type_maps};
+    use std::collections::BTreeMap;
+
+    /*
+     * Multi-version type-map registry.
+     *
+     * Stores one `TypeMap` per protocol version and provides version
+     * negotiation for the host. Hosts build this from the compiled-in
+     * type maps (via `Registry::builtin()` or programmatically).
+     */
+    #[derive(Clone, Debug)]
+    pub struct Registry {
+        versions: BTreeMap<Version, TypeMap>,
+    }
+
+    impl Registry {
+        pub fn new() -> Self {
+            Self {
+                versions: BTreeMap::new(),
+            }
+        }
+
+        pub fn register(&mut self, typemap: TypeMap) {
+            self.versions.insert(typemap.version.clone(), typemap);
+        }
+
+        pub fn get(&self, version: &Version) -> Option<&TypeMap> {
+            self.versions.get(version)
+        }
+
+        pub fn supports(&self, version: &Version) -> bool {
+            self.versions.contains_key(version)
+        }
+
+        pub fn negotiate(&self, client_versions: &[Version]) -> Option<Version> {
+            client_versions
+                .iter()
+                .filter(|v| self.versions.contains_key(v))
+                .max()
+                .cloned()
+        }
+
+        pub fn latest(&self) -> Option<&TypeMap> {
+            self.versions.last_key_value().map(|(_, v)| v)
+        }
+
+        pub fn builtin() -> Self {
+            let mut r = Self::new();
+            for tm in builtin_type_maps() {
+                r.register(tm);
+            }
+            r
+        }
+    }
+
+    impl Default for Registry {
+        fn default() -> Self {
+            Self::new()
+        }
+    }
+
+    /* ================================ TESTS ================================ */
+    #[cfg(test)]
+    mod tests {
+        use super::*;
+
+        #[test]
+        fn builtin_contains_versions() {
+            let r = Registry::builtin();
+            assert!(r.supports(&Version(0, 0)));
+        }
+
+        #[test]
+        fn negotiate_finds_highest() {
+            let mut r = Registry::new();
+            r.register(TypeMap::new(Version(0, 0)));
+            r.register(TypeMap::new(Version(1, 0)));
+            r.register(TypeMap::new(Version(2, 0)));
+
+            let client = &[Version(0, 0)];
+            assert_eq!(r.negotiate(client), Some(Version(0, 0)));
+
+            let client = &[Version(1, 0), Version(0, 0)];
+            assert_eq!(r.negotiate(client), Some(Version(1, 0)));
+
+            let client = &[Version(5, 0)];
+            assert_eq!(r.negotiate(client), None);
+        }
+
+        #[test]
+        fn latest_returns_highest() {
+            let mut r = Registry::new();
+            r.register(TypeMap::new(Version(0, 0)));
+            r.register(TypeMap::new(Version(2, 0)));
+            r.register(TypeMap::new(Version(1, 0)));
+
+            assert_eq!(r.latest().unwrap().version, Version(2, 0));
+        }
+    }
+}