update benchmarking setup

Update manifests
Add KKT
2025-10-29 14:55:12 +01:00 · 2025-10-29 08:08:21 +01:00 · 2025-10-29 08:08:03 +01:00 · 2025-10-07 19:11:47 +02:00 · 2025-08-19 13:56:00 +02:00 · 2025-08-08 17:23:56 +02:00
24 changed files with 3179 additions and 81 deletions
@@ -1298,6 +1298,16 @@ version = "0.7.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f46ad14479a25103f283c0f10005961cf086d8dc42205bb44c46ac563475dca6"
 [[package]]
 name = "classic-mceliece-rust"
 version = "3.2.0"
 source = "git+https://github.com/georgio/classic-mceliece-rust#d2367683c06330b694047132d7bd38924ac97da8"
 dependencies = [
 "rand 0.9.2",
 "sha3",
 "zeroize",
 ]
 [[package]]
 name = "colorchoice"
 version = "1.0.3"
@@ -1457,6 +1467,16 @@ version = "0.8.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b"
 [[package]]
 name = "core-models"
 version = "0.0.3"
 source = "git+https://github.com/cryspen/libcrux#7cfce8fcf0666dc016f4a8e8ed2c6260285052ac"
 dependencies = [
 "hax-lib",
 "pastey",
 "rand 0.9.2",
 ]
 [[package]]
 name = "cosmos-sdk-proto"
 version = "0.26.1"
@@ -1878,6 +1898,7 @@ dependencies = [
 "curve25519-dalek-derive",
 "digest 0.10.7",
 "fiat-crypto",
 "rand_core 0.6.4",
 "rustc_version 0.4.1",
 "serde",
 "subtle 2.6.1",
@@ -3141,6 +3162,43 @@ dependencies = [
 "hashbrown 0.15.2",
 ]
 [[package]]
 name = "hax-lib"
 version = "0.3.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "74d9ba66d1739c68e0219b2b2238b5c4145f491ebf181b9c6ab561a19352ae86"
 dependencies = [
 "hax-lib-macros",
 "num-bigint",
 "num-traits",
 ]
 [[package]]
 name = "hax-lib-macros"
 version = "0.3.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "24ba777a231a58d1bce1d68313fa6b6afcc7966adef23d60f45b8a2b9b688bf1"
 dependencies = [
 "hax-lib-macros-types",
 "proc-macro-error2",
 "proc-macro2",
 "quote",
 "syn 2.0.98",
 ]
 [[package]]
 name = "hax-lib-macros-types"
 version = "0.3.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "867e19177d7425140b417cd27c2e05320e727ee682e98368f88b7194e80ad515"
 dependencies = [
 "proc-macro2",
 "quote",
 "serde",
 "serde_json",
 "uuid",
 ]
 [[package]]
 name = "hdrhistogram"
 version = "7.5.4"
@@ -3250,7 +3308,7 @@ dependencies = [
 "idna",
 "ipnet",
 "once_cell",
- "rand 0.9.0",
+ "rand 0.9.2",
 "ring",
 "rustls 0.23.25",
 "thiserror 2.0.12",
@@ -3275,7 +3333,7 @@ dependencies = [
 "moka",
 "once_cell",
 "parking_lot",
- "rand 0.9.0",
+ "rand 0.9.2",
 "resolv-conf",
 "rustls 0.23.25",
 "smallvec",
@@ -4072,6 +4130,15 @@ dependencies = [
 "signature",
 ]
 [[package]]
 name = "keccak"
 version = "0.1.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ecc2af9a1119c51f12a14607e783cb977bde58bc069ff0c3da1095e635d70654"
 dependencies = [
 "cpufeatures",
 ]
 [[package]]
 name = "keystream"
 version = "1.0.0"
@@ -4150,6 +4217,213 @@ version = "0.2.171"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c19937216e9d3aa9956d9bb8dfc0b0c8beb6058fc4f7a4dc4d850edf86a237d6"
 [[package]]
 name = "libcrux-chacha20poly1305"
 version = "0.0.3"
 source = "git+https://github.com/cryspen/libcrux#7cfce8fcf0666dc016f4a8e8ed2c6260285052ac"
 dependencies = [
 "libcrux-hacl-rs",
 "libcrux-macros",
 "libcrux-poly1305",
 "libcrux-secrets",
 "libcrux-traits",
 ]
 [[package]]
 name = "libcrux-curve25519"
 version = "0.0.3"
 source = "git+https://github.com/cryspen/libcrux#7cfce8fcf0666dc016f4a8e8ed2c6260285052ac"
 dependencies = [
 "libcrux-hacl-rs",
 "libcrux-macros",
 "libcrux-secrets",
 "libcrux-traits",
 ]
 [[package]]
 name = "libcrux-ecdh"
 version = "0.0.3"
 source = "git+https://github.com/cryspen/libcrux#7cfce8fcf0666dc016f4a8e8ed2c6260285052ac"
 dependencies = [
 "libcrux-curve25519",
 "libcrux-p256",
 "rand 0.9.2",
 "tls_codec",
 ]
 [[package]]
 name = "libcrux-ed25519"
 version = "0.0.3"
 source = "git+https://github.com/cryspen/libcrux#7cfce8fcf0666dc016f4a8e8ed2c6260285052ac"
 dependencies = [
 "libcrux-hacl-rs",
 "libcrux-macros",
 "libcrux-sha2",
 "rand_core 0.9.1",
 "tls_codec",
 ]
 [[package]]
 name = "libcrux-hacl-rs"
 version = "0.0.3"
 source = "git+https://github.com/cryspen/libcrux#7cfce8fcf0666dc016f4a8e8ed2c6260285052ac"
 dependencies = [
 "libcrux-macros",
 ]
 [[package]]
 name = "libcrux-hkdf"
 version = "0.0.3"
 source = "git+https://github.com/cryspen/libcrux#7cfce8fcf0666dc016f4a8e8ed2c6260285052ac"
 dependencies = [
 "libcrux-hacl-rs",
 "libcrux-hmac",
 "libcrux-secrets",
 ]
 [[package]]
 name = "libcrux-hmac"
 version = "0.0.3"
 source = "git+https://github.com/cryspen/libcrux#7cfce8fcf0666dc016f4a8e8ed2c6260285052ac"
 dependencies = [
 "libcrux-hacl-rs",
 "libcrux-macros",
 "libcrux-sha2",
 ]
 [[package]]
 name = "libcrux-intrinsics"
 version = "0.0.3"
 source = "git+https://github.com/cryspen/libcrux#7cfce8fcf0666dc016f4a8e8ed2c6260285052ac"
 dependencies = [
 "core-models",
 "hax-lib",
 ]
 [[package]]
 name = "libcrux-kem"
 version = "0.0.3"
 source = "git+https://github.com/cryspen/libcrux#7cfce8fcf0666dc016f4a8e8ed2c6260285052ac"
 dependencies = [
 "libcrux-curve25519",
 "libcrux-ecdh",
 "libcrux-ml-kem",
 "libcrux-p256",
 "libcrux-sha3",
 "libcrux-traits",
 "rand 0.9.2",
 "tls_codec",
 ]
 [[package]]
 name = "libcrux-macros"
 version = "0.0.3"
 source = "git+https://github.com/cryspen/libcrux#7cfce8fcf0666dc016f4a8e8ed2c6260285052ac"
 dependencies = [
 "quote",
 "syn 2.0.98",
 ]
 [[package]]
 name = "libcrux-ml-kem"
 version = "0.0.3"
 source = "git+https://github.com/cryspen/libcrux#7cfce8fcf0666dc016f4a8e8ed2c6260285052ac"
 dependencies = [
 "hax-lib",
 "libcrux-intrinsics",
 "libcrux-platform",
 "libcrux-secrets",
 "libcrux-sha3",
 "libcrux-traits",
 "rand 0.9.2",
 "tls_codec",
 ]
 [[package]]
 name = "libcrux-p256"
 version = "0.0.3"
 source = "git+https://github.com/cryspen/libcrux#7cfce8fcf0666dc016f4a8e8ed2c6260285052ac"
 dependencies = [
 "libcrux-hacl-rs",
 "libcrux-macros",
 "libcrux-secrets",
 "libcrux-sha2",
 "libcrux-traits",
 ]
 [[package]]
 name = "libcrux-platform"
 version = "0.0.2"
 source = "git+https://github.com/cryspen/libcrux#7cfce8fcf0666dc016f4a8e8ed2c6260285052ac"
 dependencies = [
 "libc",
 ]
 [[package]]
 name = "libcrux-poly1305"
 version = "0.0.3"
 source = "git+https://github.com/cryspen/libcrux#7cfce8fcf0666dc016f4a8e8ed2c6260285052ac"
 dependencies = [
 "libcrux-hacl-rs",
 "libcrux-macros",
 ]
 [[package]]
 name = "libcrux-psq"
 version = "0.0.3"
 source = "git+https://github.com/cryspen/libcrux#7cfce8fcf0666dc016f4a8e8ed2c6260285052ac"
 dependencies = [
 "libcrux-chacha20poly1305",
 "libcrux-ecdh",
 "libcrux-ed25519",
 "libcrux-hkdf",
 "libcrux-hmac",
 "libcrux-kem",
 "libcrux-ml-kem",
 "libcrux-sha2",
 "libcrux-traits",
 "rand 0.9.2",
 "tls_codec",
 ]
 [[package]]
 name = "libcrux-secrets"
 version = "0.0.3"
 source = "git+https://github.com/cryspen/libcrux#7cfce8fcf0666dc016f4a8e8ed2c6260285052ac"
 dependencies = [
 "hax-lib",
 ]
 [[package]]
 name = "libcrux-sha2"
 version = "0.0.3"
 source = "git+https://github.com/cryspen/libcrux#7cfce8fcf0666dc016f4a8e8ed2c6260285052ac"
 dependencies = [
 "libcrux-hacl-rs",
 "libcrux-macros",
 "libcrux-traits",
 ]
 [[package]]
 name = "libcrux-sha3"
 version = "0.0.3"
 source = "git+https://github.com/cryspen/libcrux#7cfce8fcf0666dc016f4a8e8ed2c6260285052ac"
 dependencies = [
 "hax-lib",
 "libcrux-intrinsics",
 "libcrux-platform",
 "libcrux-traits",
 ]
 [[package]]
 name = "libcrux-traits"
 version = "0.0.3"
 source = "git+https://github.com/cryspen/libcrux#7cfce8fcf0666dc016f4a8e8ed2c6260285052ac"
 dependencies = [
 "libcrux-secrets",
 "rand 0.9.2",
 ]
 [[package]]
 name = "libm"
 version = "0.2.11"
@@ -6223,6 +6497,35 @@ dependencies = [
 "url",
 ]
 [[package]]
 name = "nym-kkt"
 version = "0.1.0"
 dependencies = [
 "aead",
 "arc-swap",
 "blake3",
 "bytes",
 "classic-mceliece-rust",
 "criterion",
 "curve25519-dalek",
 "futures",
 "libcrux-ecdh",
 "libcrux-kem",
 "libcrux-ml-kem",
 "libcrux-psq",
 "libcrux-sha3",
 "libcrux-traits",
 "nym-crypto 0.4.0",
 "pin-project",
 "rand 0.9.2",
 "strum 0.26.3",
 "thiserror 2.0.12",
 "tokio",
 "tokio-util",
 "tracing",
 "zeroize",
 ]
 [[package]]
 name = "nym-ledger"
 version = "0.1.0"
@@ -6781,9 +7084,14 @@ dependencies = [
 "arc-swap",
 "bytes",
 "futures",
 "libcrux-ed25519",
 "libcrux-kem",
 "libcrux-psq",
 "nym-crypto 0.4.0",
 "nym-noise-keys",
 "nympsq",
 "pin-project",
 "rand 0.9.2",
 "rand_chacha 0.3.1",
 "sha2 0.10.9",
 "snow",
@@ -7968,6 +8276,21 @@ dependencies = [
 "x25519-dalek",
 ]
 [[package]]
 name = "nympsq"
 version = "0.1.0"
 dependencies = [
 "criterion",
 "libcrux-ed25519",
 "libcrux-kem",
 "libcrux-psq",
 "libcrux-traits",
 "nym-crypto 0.4.0",
 "rand 0.9.2",
 "thiserror 2.0.12",
 "tls_codec",
 ]
 [[package]]
 name = "nymvisor"
 version = "0.1.22"
@@ -8286,6 +8609,12 @@ version = "1.0.15"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a"
 [[package]]
 name = "pastey"
 version = "0.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b3a8cb46bdc156b1c90460339ae6bfd45ba0394e5effbaa640badb4987fdc261"
 [[package]]
 name = "peg"
 version = "0.8.4"
@@ -8606,7 +8935,7 @@ dependencies = [
 "hmac",
 "md-5",
 "memchr",
- "rand 0.9.0",
+ "rand 0.9.2",
 "sha2 0.10.9",
 "stringprep",
 ]
@@ -8857,7 +9186,7 @@ checksum = "b820744eb4dc9b57a3398183639c511b5a26d2ed702cedd3febaa1393caa22cc"
 dependencies = [
 "bytes",
 "getrandom 0.3.1",
- "rand 0.9.0",
+ "rand 0.9.2",
 "ring",
 "rustc-hash",
 "rustls 0.23.25",
@@ -8911,13 +9240,12 @@ dependencies = [
 [[package]]
 name = "rand"
-version = "0.9.0"
+version = "0.9.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3779b94aeb87e8bd4e834cee3650289ee9e0d5677f976ecdb6d219e5f4f6cd94"
+checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1"
 dependencies = [
 "rand_chacha 0.9.0",
 "rand_core 0.9.1",
 "zerocopy 0.8.20",
 ]
 [[package]]
@@ -9949,6 +10277,16 @@ dependencies = [
 "digest 0.10.7",
 ]
 [[package]]
 name = "sha3"
 version = "0.10.8"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "75872d278a8f37ef87fa0ddbda7802605cb18344497949862c0d4dcb291eba60"
 dependencies = [
 "digest 0.10.7",
 "keccak",
 ]
 [[package]]
 name = "sharded-slab"
 version = "0.1.7"
@@ -10923,6 +11261,27 @@ version = "0.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"
 [[package]]
 name = "tls_codec"
 version = "0.4.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0de2e01245e2bb89d6f05801c564fa27624dbd7b1846859876c7dad82e90bf6b"
 dependencies = [
 "tls_codec_derive",
 "zeroize",
 ]
 [[package]]
 name = "tls_codec_derive"
 version = "0.4.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2d2e76690929402faae40aebdda620a2c0e25dd6d3b9afe48867dfd95991f4bd"
 dependencies = [
 "proc-macro2",
 "quote",
 "syn 2.0.98",
 ]
 [[package]]
 name = "tokio"
 version = "1.45.1"
@@ -10982,7 +11341,7 @@ dependencies = [
 "pin-project-lite",
 "postgres-protocol",
 "postgres-types",
- "rand 0.9.0",
+ "rand 0.9.2",
 "socket2",
 "tokio",
 "tokio-util",
@@ -66,8 +66,10 @@ members = [
    "common/nym-id",
    "common/nym-metrics",
    "common/nym_offline_compact_ecash",
    "common/nymkkt",
    "common/nymnoise",
    "common/nymnoise/keys",
    "common/nympsq",
    "common/nymsphinx",
    "common/nymsphinx/acknowledgements",
    "common/nymsphinx/addressing",
@@ -133,7 +133,7 @@ impl ManagedConnection {
                    debug!("Managed to establish connection to {}", self.address);
                    let noise_stream =
-                        match upgrade_noise_initiator(stream, &self.noise_config).await {
+                        match upgrade_noise_initiator(stream, &self.noise_config, None).await {
                            Ok(noise_stream) => noise_stream,
                            Err(err) => {
                                error!("Failed to perform Noise handshake with {address} - {err}");
@@ -0,0 +1,47 @@
 [package]
 name = "nym-kkt"
 version = "0.1.0"
 authors = ["Georgio Nicolas <georgio@nymtech.net>"]
 edition = "2021"
 license.workspace = true
 [dependencies]
 arc-swap = { workspace = true }
 bytes = { workspace = true }
 futures = { workspace = true }
 tracing = { workspace = true }
 pin-project = { workspace = true }
 blake3 = { workspace = true }
 aead = { workspace = true }
 strum = { workspace = true, features = ["derive"] }
 thiserror = { workspace = true }
 tokio = { workspace = true, features = ["full"] }
 tokio-util = { workspace = true, features = ["codec"] }
 # internal
 nym-crypto = { path = "../crypto", features = ["asymmetric", "serde"]}
 libcrux-traits = { git = "https://github.com/cryspen/libcrux" }
 libcrux-kem = { git = "https://github.com/cryspen/libcrux" }
 libcrux-psq = { git = "https://github.com/cryspen/libcrux", features = ["test-utils"] }
 libcrux-sha3 = { git = "https://github.com/cryspen/libcrux" }
 libcrux-ml-kem = { git = "https://github.com/cryspen/libcrux" }
 libcrux-ecdh = { git = "https://github.com/cryspen/libcrux", features = ["codec"]}
 rand = "0.9.2"
 curve25519-dalek = {version = "4.1.3", features = ["rand_core", "serde"] }
 zeroize = { workspace = true, features = ["zeroize_derive"] }
 classic-mceliece-rust = { git = "https://github.com/georgio/classic-mceliece-rust", features = ["mceliece460896f","zeroize"]}
 [dev-dependencies]
 criterion = {workspace = true}
 [[bench]]
 name = "benches"
 harness = false
 [lints]
 workspace = true
@@ -0,0 +1,518 @@
 // Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 use criterion::{criterion_group, criterion_main, Criterion};
 use nym_crypto::asymmetric::ed25519;
 use nym_kkt::{
    ciphersuite::{Ciphersuite, EncapsulationKey, HashFunction, SignatureScheme, KEM},
    context::KKTMode,
    frame::KKTFrame,
    key_utils::{generate_keypair_libcrux, generate_keypair_mceliece, hash_encapsulation_key},
    session::{
        anonymous_initiator_process, initiator_ingest_response, initiator_process,
        responder_ingest_message, responder_process,
    },
 };
 use rand::prelude::*;
 pub fn gen_ed25519_keypair(c: &mut Criterion) {
    c.bench_function("Generate Ed25519 Keypair", |b| {
        b.iter(|| {
            let mut s: [u8; 32] = [0u8; 32];
            rand::rng().fill_bytes(&mut s);
            ed25519::KeyPair::from_secret(s, 0)
        });
    });
 }
 pub fn gen_mlkem768_keypair(c: &mut Criterion) {
    c.bench_function("Generate MlKem768 Keypair", |b| {
        b.iter(|| {
            libcrux_kem::key_gen(libcrux_kem::Algorithm::MlKem768, &mut rand::rng()).unwrap()
        });
    });
 }
 pub fn kkt_benchmark(c: &mut Criterion) {
    let mut rng = rand::rng();
    // generate ed25519 keys
    let mut secret_initiator: [u8; 32] = [0u8; 32];
    rng.fill_bytes(&mut secret_initiator);
    let initiator_ed25519_keypair = ed25519::KeyPair::from_secret(secret_initiator, 0);
    let mut secret_responder: [u8; 32] = [0u8; 32];
    rng.fill_bytes(&mut secret_responder);
    let responder_ed25519_keypair = ed25519::KeyPair::from_secret(secret_responder, 1);
    for kem in [KEM::MlKem768, KEM::XWing, KEM::X25519, KEM::McEliece] {
        for hash_function in [
            HashFunction::Blake3,
            HashFunction::SHA256,
            HashFunction::SHAKE128,
            HashFunction::SHAKE256,
        ] {
            let ciphersuite = Ciphersuite::resolve_ciphersuite(
                kem,
                hash_function,
                SignatureScheme::Ed25519,
                None,
            )
            .unwrap();
            // generate kem public keys
            let (responder_kem_public_key, initiator_kem_public_key) = match kem {
                KEM::MlKem768 => (
                    EncapsulationKey::MlKem768(generate_keypair_libcrux(&mut rng, kem).unwrap().1),
                    EncapsulationKey::MlKem768(generate_keypair_libcrux(&mut rng, kem).unwrap().1),
                ),
                KEM::XWing => (
                    EncapsulationKey::XWing(generate_keypair_libcrux(&mut rng, kem).unwrap().1),
                    EncapsulationKey::XWing(generate_keypair_libcrux(&mut rng, kem).unwrap().1),
                ),
                KEM::X25519 => (
                    EncapsulationKey::X25519(generate_keypair_libcrux(&mut rng, kem).unwrap().1),
                    EncapsulationKey::X25519(generate_keypair_libcrux(&mut rng, kem).unwrap().1),
                ),
                KEM::McEliece => (
                    EncapsulationKey::McEliece(generate_keypair_mceliece(&mut rng).1),
                    EncapsulationKey::McEliece(generate_keypair_mceliece(&mut rng).1),
                ),
            };
            let i_kem_key_bytes = initiator_kem_public_key.encode();
            let r_kem_key_bytes = responder_kem_public_key.encode();
            let i_dir_hash = hash_encapsulation_key(
                &ciphersuite.hash_function(),
                ciphersuite.hash_len(),
                &i_kem_key_bytes,
            );
            let r_dir_hash = hash_encapsulation_key(
                &ciphersuite.hash_function(),
                ciphersuite.hash_len(),
                &r_kem_key_bytes,
            );
            // Anonymous Initiator, OneWay
            {
                c.bench_function(
                    &format!(
                        "{}, {} | Anonymous Initiator: Generate Request",
                        kem, hash_function
                    ),
                    |b| {
                        b.iter(|| anonymous_initiator_process(&mut rng, ciphersuite).unwrap());
                    },
                );
                let (mut i_context, i_frame) =
                    anonymous_initiator_process(&mut rng, ciphersuite).unwrap();
                c.bench_function(
                    &format!(
                        "{}, {} | Anonymous Initiator: Encode Frame - Request",
                        kem, hash_function
                    ),
                    |b| b.iter(|| i_frame.to_bytes()),
                );
                let i_frame_bytes = i_frame.to_bytes();
                c.bench_function(
                    &format!(
                        "{}, {} | Anonymous Initiator: Decode Frame - Request",
                        kem, hash_function
                    ),
                    |b| b.iter(|| KKTFrame::from_bytes(&i_frame_bytes).unwrap()),
                );
                let (i_frame_r, r_context) = KKTFrame::from_bytes(&i_frame_bytes).unwrap();
                c.bench_function(
                    &format!(
                        "{}, {} | Anonymous Initiator: Responder Ingest Frame",
                        kem, hash_function
                    ),
                    |b| {
                        b.iter(|| {
                            responder_ingest_message(&r_context, None, None, &i_frame_r).unwrap()
                        });
                    },
                );
                let (mut r_context, _) =
                    responder_ingest_message(&r_context, None, None, &i_frame_r).unwrap();
                c.bench_function(
                    &format!(
                        "{}, {} | Anonymous Initiator: Responder Generate Response",
                        kem, hash_function
                    ),
                    |b| {
                        b.iter(|| {
                            responder_process(
                                &mut r_context,
                                i_frame_r.session_id_ref(),
                                responder_ed25519_keypair.private_key(),
                                &responder_kem_public_key,
                            )
                            .unwrap()
                        });
                    },
                );
                let r_frame = responder_process(
                    &mut r_context,
                    i_frame_r.session_id_ref(),
                    responder_ed25519_keypair.private_key(),
                    &responder_kem_public_key,
                )
                .unwrap();
                c.bench_function(
                    &format!(
                        "{}, {} | Anonymous Initiator: Responder Encode Frame",
                        kem, hash_function
                    ),
                    |b| b.iter(|| r_frame.to_bytes()),
                );
                let r_bytes = r_frame.to_bytes();
                c.bench_function(
                    &format!(
                        "{}, {} | Anonymous Initiator: Initiator Ingest Response",
                        kem, hash_function
                    ),
                    |b| {
                        b.iter(|| {
                            initiator_ingest_response(
                                &mut i_context,
                                responder_ed25519_keypair.public_key(),
                                &r_dir_hash,
                                &r_bytes,
                            )
                            .unwrap()
                        });
                    },
                );
                let obtained_key = initiator_ingest_response(
                    &mut i_context,
                    responder_ed25519_keypair.public_key(),
                    &r_dir_hash,
                    &r_bytes,
                )
                .unwrap();
                assert_eq!(obtained_key.encode(), r_kem_key_bytes)
            }
            // Initiator, OneWay
            {
                let (mut i_context, i_frame) = initiator_process(
                    &mut rng,
                    KKTMode::OneWay,
                    ciphersuite,
                    initiator_ed25519_keypair.private_key(),
                    None,
                )
                .unwrap();
                c.bench_function(
                    &format!(
                        "{}, {} | Initiator OneWay: Generate Request",
                        kem, hash_function
                    ),
                    |b| {
                        b.iter(|| {
                            initiator_process(
                                &mut rng,
                                KKTMode::OneWay,
                                ciphersuite,
                                initiator_ed25519_keypair.private_key(),
                                None,
                            )
                            .unwrap()
                        });
                    },
                );
                c.bench_function(
                    &format!(
                        "{}, {} | Initiator OneWay: Encode Frame - Request",
                        kem, hash_function
                    ),
                    |b| b.iter(|| i_frame.to_bytes()),
                );
                let i_frame_bytes = i_frame.to_bytes();
                c.bench_function(
                    &format!(
                        "{}, {} | Initiator OneWay: Decode Frame - Request",
                        kem, hash_function
                    ),
                    |b| b.iter(|| KKTFrame::from_bytes(&i_frame_bytes).unwrap()),
                );
                let (i_frame_r, r_context) = KKTFrame::from_bytes(&i_frame_bytes).unwrap();
                c.bench_function(
                    &format!(
                        "{}, {} | Initiator OneWay: Responder Ingest Frame",
                        kem, hash_function
                    ),
                    |b| {
                        b.iter(|| {
                            responder_ingest_message(
                                &r_context,
                                Some(initiator_ed25519_keypair.public_key()),
                                None,
                                &i_frame_r,
                            )
                            .unwrap()
                        });
                    },
                );
                let (mut r_context, r_obtained_key) = responder_ingest_message(
                    &r_context,
                    Some(initiator_ed25519_keypair.public_key()),
                    None,
                    &i_frame_r,
                )
                .unwrap();
                assert!(r_obtained_key.is_none());
                c.bench_function(
                    &format!(
                        "{}, {} | Initiator OneWay: Responder Generate Response",
                        kem, hash_function
                    ),
                    |b| {
                        b.iter(|| {
                            responder_process(
                                &mut r_context,
                                i_frame_r.session_id_ref(),
                                responder_ed25519_keypair.private_key(),
                                &responder_kem_public_key,
                            )
                            .unwrap()
                        });
                    },
                );
                let r_frame = responder_process(
                    &mut r_context,
                    i_frame_r.session_id_ref(),
                    responder_ed25519_keypair.private_key(),
                    &responder_kem_public_key,
                )
                .unwrap();
                c.bench_function(
                    &format!(
                        "{}, {} | Initiator OneWay: Responder Encode Frame",
                        kem, hash_function
                    ),
                    |b| {
                        b.iter(|| r_frame.to_bytes());
                    },
                );
                let r_bytes = r_frame.to_bytes();
                c.bench_function(
                    &format!(
                        "{}, {} | Initiator OneWay: Initiator Ingest Response",
                        kem, hash_function
                    ),
                    |b| {
                        b.iter(|| {
                            initiator_ingest_response(
                                &mut i_context,
                                responder_ed25519_keypair.public_key(),
                                &r_dir_hash,
                                &r_bytes,
                            )
                            .unwrap()
                        });
                    },
                );
                let i_obtained_key = initiator_ingest_response(
                    &mut i_context,
                    responder_ed25519_keypair.public_key(),
                    &r_dir_hash,
                    &r_bytes,
                )
                .unwrap();
                assert_eq!(i_obtained_key.encode(), r_kem_key_bytes)
            }
            // Initiator, Mutual
            {
                c.bench_function(
                    &format!(
                        "{}, {} | Initiator Mutual: Generate Request",
                        kem, hash_function
                    ),
                    |b| {
                        b.iter(|| {
                            initiator_process(
                                &mut rng,
                                KKTMode::Mutual,
                                ciphersuite,
                                initiator_ed25519_keypair.private_key(),
                                Some(&initiator_kem_public_key),
                            )
                            .unwrap()
                        });
                    },
                );
                let (mut i_context, i_frame) = initiator_process(
                    &mut rng,
                    KKTMode::Mutual,
                    ciphersuite,
                    initiator_ed25519_keypair.private_key(),
                    Some(&initiator_kem_public_key),
                )
                .unwrap();
                c.bench_function(
                    &format!(
                        "{}, {} | Initiator Mutual: Encode Frame - Request",
                        kem, hash_function
                    ),
                    |b| {
                        b.iter(|| i_frame.to_bytes());
                    },
                );
                let i_frame_bytes = i_frame.to_bytes();
                c.bench_function(
                    &format!(
                        "{}, {} | Initiator Mutual: Decode Frame - Request",
                        kem, hash_function
                    ),
                    |b| {
                        b.iter(|| KKTFrame::from_bytes(&i_frame_bytes).unwrap());
                    },
                );
                let (i_frame_r, r_context) = KKTFrame::from_bytes(&i_frame_bytes).unwrap();
                c.bench_function(
                    &format!(
                        "{}, {} | Initiator Mutual: Responder Ingest Frame",
                        kem, hash_function
                    ),
                    |b| {
                        b.iter(|| {
                            responder_ingest_message(
                                &r_context,
                                Some(initiator_ed25519_keypair.public_key()),
                                Some(&i_dir_hash),
                                &i_frame_r,
                            )
                            .unwrap()
                        });
                    },
                );
                let (mut r_context, r_obtained_key) = responder_ingest_message(
                    &r_context,
                    Some(initiator_ed25519_keypair.public_key()),
                    Some(&i_dir_hash),
                    &i_frame_r,
                )
                .unwrap();
                assert_eq!(r_obtained_key.unwrap().encode(), i_kem_key_bytes);
                c.bench_function(
                    &format!(
                        "{}, {} | Initiator Mutual: Responder Generate Response",
                        kem, hash_function
                    ),
                    |b| {
                        b.iter(|| {
                            responder_process(
                                &mut r_context,
                                i_frame_r.session_id_ref(),
                                responder_ed25519_keypair.private_key(),
                                &responder_kem_public_key,
                            )
                            .unwrap()
                        });
                    },
                );
                let r_frame = responder_process(
                    &mut r_context,
                    i_frame_r.session_id_ref(),
                    responder_ed25519_keypair.private_key(),
                    &responder_kem_public_key,
                )
                .unwrap();
                c.bench_function(
                    &format!(
                        "{}, {} | Initiator Mutual: Responder Encode Frame",
                        kem, hash_function
                    ),
                    |b| {
                        b.iter(|| {
                            r_frame.to_bytes();
                        });
                    },
                );
                let r_bytes = r_frame.to_bytes();
                c.bench_function(
                    &format!(
                        "{}, {} | Initiator Mutual: Initiator Ingest Response",
                        kem, hash_function
                    ),
                    |b| {
                        b.iter(|| {
                            initiator_ingest_response(
                                &mut i_context,
                                responder_ed25519_keypair.public_key(),
                                &r_dir_hash,
                                &r_bytes,
                            )
                            .unwrap()
                        });
                    },
                );
                let obtained_key = initiator_ingest_response(
                    &mut i_context,
                    responder_ed25519_keypair.public_key(),
                    &r_dir_hash,
                    &r_bytes,
                )
                .unwrap();
                assert_eq!(obtained_key.encode(), r_kem_key_bytes)
            }
        }
    }
 }
 criterion_group!(
    benches,
    gen_ed25519_keypair,
    gen_mlkem768_keypair,
    kkt_benchmark
 );
 criterion_main!(benches);
@@ -0,0 +1,311 @@
 // Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 use std::fmt::Display;
 use libcrux_kem::{Algorithm, MlKem768PublicKey};
 use nym_crypto::asymmetric::ed25519;
 use crate::error::KKTError;
 pub const HASH_LEN_256: u8 = 32;
 pub const CIPHERSUITE_ENCODING_LEN: usize = 4;
 pub const CURVE25519_KEY_LEN: usize = 32;
 #[derive(Clone, Copy, Debug)]
 pub enum HashFunction {
    Blake3,
    SHAKE128,
    SHAKE256,
    SHA256,
 }
 impl Display for HashFunction {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(match self {
            HashFunction::Blake3 => "Blake3",
            HashFunction::SHAKE128 => "SHAKE128",
            HashFunction::SHAKE256 => "SHAKE256",
            HashFunction::SHA256 => "SHA256",
        })
    }
 }
 pub enum EncapsulationKey<'a> {
    MlKem768(libcrux_kem::PublicKey),
    XWing(libcrux_kem::PublicKey),
    X25519(libcrux_kem::PublicKey),
    McEliece(classic_mceliece_rust::PublicKey<'a>),
 }
 pub enum DecapsulationKey<'a> {
    MlKem768(libcrux_kem::PrivateKey),
    XWing(libcrux_kem::PrivateKey),
    X25519(libcrux_kem::PrivateKey),
    McEliece(classic_mceliece_rust::SecretKey<'a>),
 }
 impl<'a> EncapsulationKey<'a> {
    pub(crate) fn decode(kem: KEM, bytes: &[u8]) -> Result<Self, KKTError> {
        match kem {
            KEM::McEliece => {
                if bytes.len() != classic_mceliece_rust::CRYPTO_PUBLICKEYBYTES {
                    Err(KKTError::KEMError {
                        info: "Received McEliece Encapsulation Key with Invalid Length",
                    })
                } else {
                    let mut public_key_bytes =
                        Box::new([0u8; classic_mceliece_rust::CRYPTO_PUBLICKEYBYTES]);
                    // Size must be correct due to KKTFrame::from_bytes(message_bytes)?
                    public_key_bytes.clone_from_slice(bytes);
                    Ok(EncapsulationKey::McEliece(
                        classic_mceliece_rust::PublicKey::from(public_key_bytes),
                    ))
                }
            }
            KEM::X25519 => Ok(EncapsulationKey::X25519(libcrux_kem::PublicKey::decode(
                map_kem_to_libcrux_kem(kem),
                bytes,
            )?)),
            KEM::MlKem768 => Ok(EncapsulationKey::MlKem768(libcrux_kem::PublicKey::decode(
                map_kem_to_libcrux_kem(kem),
                bytes,
            )?)),
            KEM::XWing => Ok(EncapsulationKey::XWing(libcrux_kem::PublicKey::decode(
                map_kem_to_libcrux_kem(kem),
                bytes,
            )?)),
        }
    }
    pub fn encode(&self) -> Vec<u8> {
        match self {
            EncapsulationKey::XWing(public_key)
            | EncapsulationKey::MlKem768(public_key)
            | EncapsulationKey::X25519(public_key) => public_key.encode(),
            EncapsulationKey::McEliece(public_key) => Vec::from(public_key.as_array()),
        }
    }
 }
 #[derive(Clone, Copy, Debug)]
 pub enum SignatureScheme {
    Ed25519,
 }
 impl Display for SignatureScheme {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(match self {
            SignatureScheme::Ed25519 => "Ed25519",
        })
    }
 }
 #[derive(Clone, Copy, Debug)]
 pub enum KEM {
    MlKem768,
    XWing,
    X25519,
    McEliece,
 }
 impl Display for KEM {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(match self {
            KEM::MlKem768 => "MlKem768",
            KEM::XWing => "XWing",
            KEM::X25519 => "x25519",
            KEM::McEliece => "McEliece",
        })
    }
 }
 #[derive(Clone, Copy, Debug)]
 pub struct Ciphersuite {
    hash_function: HashFunction,
    signature_scheme: SignatureScheme,
    kem: KEM,
    hash_length: u8,
    encapsulation_key_length: usize,
    signing_key_length: usize,
    verification_key_length: usize,
    signature_length: usize,
 }
 impl Ciphersuite {
    pub fn kem_key_len(&self) -> usize {
        self.encapsulation_key_length
    }
    pub fn signature_len(&self) -> usize {
        self.signature_length
    }
    pub fn signing_key_len(&self) -> usize {
        self.signing_key_length
    }
    pub fn verification_key_len(&self) -> usize {
        self.verification_key_length
    }
    pub fn hash_function(&self) -> HashFunction {
        self.hash_function
    }
    pub fn kem(&self) -> KEM {
        self.kem
    }
    pub fn signature_scheme(&self) -> SignatureScheme {
        self.signature_scheme
    }
    pub fn hash_len(&self) -> usize {
        self.hash_length as usize
    }
    pub fn default() -> Self {
        Self::resolve_ciphersuite(
            KEM::XWing,
            HashFunction::Blake3,
            SignatureScheme::Ed25519,
            None,
        )
        .unwrap()
    }
    pub fn resolve_ciphersuite(
        kem: KEM,
        hash_function: HashFunction,
        signature_scheme: SignatureScheme,
        // This should be None 99.9999% of the time
        custom_hash_length: Option<u8>,
    ) -> Result<Self, KKTError> {
        let hash_len = match custom_hash_length {
            Some(l) => {
                if l < 16 {
                    return Err(KKTError::InsecureHashLen);
                } else {
                    l
                }
            }
            None => HASH_LEN_256,
        };
        Ok(Self {
            hash_function: hash_function,
            signature_scheme: signature_scheme,
            kem: kem,
            hash_length: hash_len,
            encapsulation_key_length: match kem {
                // 1184 bytes
                KEM::MlKem768 => MlKem768PublicKey::len(),
                // 1216 bytes = 1184 + 32
                KEM::XWing => MlKem768PublicKey::len() + CURVE25519_KEY_LEN,
                // 32 bytes
                KEM::X25519 => CURVE25519_KEY_LEN,
                // 524160 bytes
                KEM::McEliece => classic_mceliece_rust::CRYPTO_PUBLICKEYBYTES,
            },
            signing_key_length: match signature_scheme {
                // 32 bytes
                SignatureScheme::Ed25519 => ed25519::SECRET_KEY_LENGTH,
            },
            verification_key_length: match signature_scheme {
                // 32 bytes
                SignatureScheme::Ed25519 => ed25519::PUBLIC_KEY_LENGTH,
            },
            signature_length: match signature_scheme {
                // 64 bytes
                SignatureScheme::Ed25519 => ed25519::SIGNATURE_LENGTH,
            },
        })
    }
    pub fn encode(&self) -> [u8; 4] {
        // [kem, hash, hashlen, sig]
        [
            match self.kem {
                KEM::XWing => 0,
                KEM::MlKem768 => 1,
                KEM::McEliece => 2,
                KEM::X25519 => 255,
            },
            match self.hash_function {
                HashFunction::Blake3 => 0,
                HashFunction::SHAKE256 => 1,
                HashFunction::SHAKE128 => 2,
                HashFunction::SHA256 => 3,
            },
            match self.hash_length {
                HASH_LEN_256 => 0,
                _ => self.hash_length as u8,
            },
            match self.signature_scheme {
                SignatureScheme::Ed25519 => 0,
            },
        ]
    }
    pub fn decode(encoding: &[u8]) -> Result<Self, KKTError> {
        if encoding.len() == 4 {
            let kem = match encoding[0] {
                0 => KEM::XWing,
                1 => KEM::MlKem768,
                2 => KEM::McEliece,
                255 => KEM::X25519,
                _ => {
                    return Err(KKTError::CiphersuiteDecodingError {
                        info: format!("Undefined KEM: {}", encoding[0]),
                    })
                }
            };
            let hash_function = match encoding[1] {
                0 => HashFunction::Blake3,
                1 => HashFunction::SHAKE256,
                2 => HashFunction::SHAKE128,
                3 => HashFunction::SHA256,
                _ => {
                    return Err(KKTError::CiphersuiteDecodingError {
                        info: format!("Undefined Hash Function: {}", encoding[1]),
                    })
                }
            };
            let custom_hash_length = match encoding[2] {
                0 => None,
                _ => Some(encoding[2]),
            };
            let signature_scheme = match encoding[3] {
                0 => SignatureScheme::Ed25519,
                _ => {
                    return Err(KKTError::CiphersuiteDecodingError {
                        info: format!("Undefined Signature Scheme: {}", encoding[3]),
                    })
                }
            };
            Self::resolve_ciphersuite(kem, hash_function, signature_scheme, custom_hash_length)
        } else {
            return Err(KKTError::CiphersuiteDecodingError {
                info: format!(
                    "Incorrect Encoding Length: actual: {} != expected: {}",
                    encoding.len(),
                    CIPHERSUITE_ENCODING_LEN
                ),
            });
        }
    }
 }
 impl Display for Ciphersuite {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(
            &format!(
                "{}_{}({})_{}",
                self.kem, self.hash_function, self.hash_length, self.signature_scheme
            )
            .to_ascii_lowercase(),
        )
    }
 }
 pub const fn map_kem_to_libcrux_kem(kem: KEM) -> Algorithm {
    match kem {
        KEM::MlKem768 => Algorithm::MlKem768,
        KEM::XWing => Algorithm::XWingKemDraft06,
        KEM::X25519 => Algorithm::X25519,
        _ => unreachable!(),
    }
 }
@@ -0,0 +1,258 @@
 // Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 use std::fmt::Display;
 use crate::{ciphersuite::Ciphersuite, error::KKTError, frame::KKT_SESSION_ID_LEN, KKT_VERSION};
 pub const KKT_CONTEXT_LEN: usize = 7;
 #[derive(Clone, Copy, PartialEq, Debug)]
 pub enum KKTStatus {
    Ok,
    InvalidRequestFormat,
    InvalidResponseFormat,
    InvalidSignature,
    UnsupportedCiphersuite,
    UnsupportedKKTVersion,
    InvalidKey,
    Timeout,
 }
 impl Display for KKTStatus {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(match self {
            KKTStatus::Ok => "Ok",
            KKTStatus::InvalidRequestFormat => "Invalid Request Format",
            KKTStatus::InvalidResponseFormat => "Invalid Response Format",
            KKTStatus::InvalidSignature => "Invalid Signature",
            KKTStatus::UnsupportedCiphersuite => "Unsupported Ciphersuite",
            KKTStatus::UnsupportedKKTVersion => "Unsupported KKT Version",
            KKTStatus::InvalidKey => "Invalid Key",
            KKTStatus::Timeout => "Timeout",
        })
    }
 }
 #[derive(Clone, Copy, PartialEq, Debug)]
 pub enum KKTRole {
    Initiator,
    AnonymousInitiator,
    Responder,
 }
 #[derive(Clone, Copy, PartialEq, Debug)]
 pub enum KKTMode {
    OneWay,
    Mutual,
 }
 #[derive(Copy, Clone, Debug)]
 pub struct KKTContext {
    version: u8,
    message_sequence: u8,
    status: KKTStatus,
    mode: KKTMode,
    role: KKTRole,
    ciphersuite: Ciphersuite,
 }
 impl KKTContext {
    pub fn new(role: KKTRole, mode: KKTMode, ciphersuite: Ciphersuite) -> Result<Self, KKTError> {
        if role == KKTRole::AnonymousInitiator && mode != KKTMode::OneWay {
            return Err(KKTError::IncompatibilityError {
                info: "Anonymous Initiator can only use OneWay mode",
            });
        }
        Ok(Self {
            version: KKT_VERSION,
            message_sequence: 0,
            status: KKTStatus::Ok,
            mode: mode,
            role: role,
            ciphersuite: ciphersuite,
        })
    }
    pub fn derive_responder_header(&self) -> Result<Self, KKTError> {
        let mut responder_header = self.clone();
        responder_header.increment_message_sequence_count()?;
        responder_header.role = KKTRole::Responder;
        Ok(responder_header)
    }
    pub fn increment_message_sequence_count(&mut self) -> Result<(), KKTError> {
        if self.message_sequence + 1 < (1 << 4) {
            self.message_sequence += 1;
            Ok(())
        } else {
            Err(KKTError::MessageCountLimitReached)
        }
    }
    pub fn update_status(&mut self, status: KKTStatus) {
        self.status = status;
    }
    pub fn version(&self) -> u8 {
        self.version
    }
    pub fn status(&self) -> KKTStatus {
        self.status
    }
    pub fn ciphersuite(&self) -> Ciphersuite {
        self.ciphersuite
    }
    pub fn role(&self) -> KKTRole {
        self.role
    }
    pub fn mode(&self) -> KKTMode {
        self.mode
    }
    pub fn body_len(&self) -> usize {
        if self.status != KKTStatus::Ok
            || (self.mode == KKTMode::OneWay
                && (self.role == KKTRole::Initiator || self.role == KKTRole::AnonymousInitiator))
        {
            0
        } else {
            self.ciphersuite.kem_key_len()
        }
    }
    pub fn signature_len(&self) -> usize {
        match self.role {
            KKTRole::Initiator | KKTRole::Responder => self.ciphersuite.signature_len(),
            KKTRole::AnonymousInitiator => 0,
        }
    }
    pub fn header_len(&self) -> usize {
        KKT_CONTEXT_LEN
    }
    pub fn session_id_len(&self) -> usize {
        // match self.role {
        //     KKTRole::Initiator | KKTRole::Responder => SESSION_ID_LENGTH,
        // It doesn't make sense to send a session_id if we send messages in the clear
        //     KKTRole::AnonymousInitiator => 0,
        // }
        KKT_SESSION_ID_LEN
    }
    pub fn full_message_len(&self) -> usize {
        self.body_len() + self.signature_len() + self.header_len() + self.session_id_len()
    }
    pub fn encode(&self) -> Result<Vec<u8>, KKTError> {
        let mut header_bytes: Vec<u8> = Vec::with_capacity(KKT_CONTEXT_LEN);
        if self.message_sequence >= 1 << 4 {
            return Err(KKTError::MessageCountLimitReached);
        }
        header_bytes.push((KKT_VERSION << 4) + self.message_sequence);
        header_bytes.push(
            match self.status {
                KKTStatus::Ok => 0,
                KKTStatus::InvalidRequestFormat => 0b001_000_00,
                KKTStatus::InvalidResponseFormat => 0b010_000_00,
                KKTStatus::InvalidSignature => 0b011_000_00,
                KKTStatus::UnsupportedCiphersuite => 0b100_000_00,
                KKTStatus::UnsupportedKKTVersion => 0b101_000_00,
                KKTStatus::InvalidKey => 0b110_000_00,
                KKTStatus::Timeout => 0b111_000_00,
            } + match self.mode {
                KKTMode::OneWay => 0,
                KKTMode::Mutual => 0b000_001_00,
            } + match self.role {
                KKTRole::Initiator => 0,
                KKTRole::Responder => 1,
                KKTRole::AnonymousInitiator => 2,
            },
        );
        header_bytes.extend_from_slice(&self.ciphersuite.encode());
        header_bytes.push(0);
        Ok(header_bytes)
    }
    pub fn try_decode(header_bytes: &[u8]) -> Result<Self, KKTError> {
        if header_bytes.len() == KKT_CONTEXT_LEN {
            let kkt_version = header_bytes[0] & 0b1111_0000;
            let message_sequence_counter = header_bytes[0] & 0b0000_1111;
            // We only check if stuff is valid here, not necessarily if it's compatible
            if (kkt_version >> 4) > KKT_VERSION {
                return Err(KKTError::FrameDecodingError {
                    info: format!("Header - Invalid KKT Version: {}", kkt_version >> 4),
                });
            }
            let status = match header_bytes[1] & 0b111_000_00 {
                0 => KKTStatus::Ok,
                0b001_000_00 => KKTStatus::InvalidRequestFormat,
                0b010_000_00 => KKTStatus::InvalidResponseFormat,
                0b011_000_00 => KKTStatus::InvalidSignature,
                0b100_000_00 => KKTStatus::UnsupportedCiphersuite,
                0b101_000_00 => KKTStatus::UnsupportedKKTVersion,
                0b110_000_00 => KKTStatus::InvalidKey,
                0b111_000_00 => KKTStatus::Timeout,
                _ => {
                    return Err(KKTError::FrameDecodingError {
                        info: format!(
                            "Header - Invalid KKT Status: {}",
                            header_bytes[1] & 0b111_000_00
                        ),
                    })
                }
            };
            let role = match header_bytes[1] & 0b000_000_11 {
                0 => KKTRole::Initiator,
                1 => KKTRole::Responder,
                2 => KKTRole::AnonymousInitiator,
                _ => {
                    return Err(KKTError::FrameDecodingError {
                        info: format!(
                            "Header - Invalid KKT Role: {}",
                            header_bytes[1] & 0b000_000_11
                        ),
                    })
                }
            };
            let mode = match (header_bytes[1] & 0b000_111_00) >> 2 {
                0 => KKTMode::OneWay,
                1 => KKTMode::Mutual,
                _ => {
                    return Err(KKTError::FrameDecodingError {
                        info: format!(
                            "Header - Invalid KKT Mode: {}",
                            (header_bytes[1] & 0b000_111_00) >> 2
                        ),
                    })
                }
            };
            Ok(KKTContext {
                version: kkt_version,
                status: status,
                mode: mode,
                role: role,
                ciphersuite: Ciphersuite::decode(&header_bytes[2..6])?,
                message_sequence: message_sequence_counter,
            })
        } else {
            Err(KKTError::FrameDecodingError {
                info: format!(
                    "Header - Invalid Header Length: actual: {} != expected: {}",
                    header_bytes.len(),
                    KKT_CONTEXT_LEN
                ),
            })
        }
    }
 }
@@ -0,0 +1,95 @@
 use core::hash;
 use blake3::{Hash, Hasher};
 use curve25519_dalek::digest::DynDigest;
 use libcrux_psq::traits::Ciphertext;
 use nym_crypto::symmetric::aead::{AeadKey, Nonce};
 use nym_crypto::{
    aes::Aes256,
    asymmetric::x25519::{self, PrivateKey, PublicKey},
    generic_array::GenericArray,
    Aes256GcmSiv,
 };
 // use rand::{CryptoRng, RngCore};
 use zeroize::Zeroize;
 use nym_crypto::aes::cipher::crypto_common::rand_core::{CryptoRng, RngCore};
 use crate::error::KKTError;
 fn generate_round_trip_symmetric_key<R>(
    rng: &mut R,
    remote_public_key: &PublicKey,
 ) -> ([u8; 64], [u8; 32])
 where
    R: CryptoRng + RngCore,
 {
    let mut s = x25519::PrivateKey::new(rng);
    let gs = s.public_key();
    let mut gbs = s.diffie_hellman(remote_public_key);
    s.zeroize();
    let mut message: [u8; 64] = [0u8; 64];
    message[0..32].clone_from_slice(gs.as_bytes());
    let mut hasher = Hasher::new();
    hasher.update(&gbs);
    gbs.zeroize();
    let key: [u8; 32] = hasher.finalize().as_bytes().to_owned();
    hasher.update(remote_public_key.as_bytes());
    hasher.update(gs.as_bytes());
    hasher.finalize_into_reset(&mut message[32..64]);
    (message, key)
 }
 fn extract_shared_secret(b: &PrivateKey, message: &[u8; 64]) -> Result<[u8; 32], KKTError> {
    let gs = PublicKey::from_bytes(&message[0..32])?;
    let mut gsb = b.diffie_hellman(&gs);
    let mut hasher = Hasher::new();
    hasher.update(&gsb);
    gsb.zeroize();
    let key: [u8; 32] = hasher.finalize().as_bytes().to_owned();
    hasher.update(b.public_key().as_bytes());
    hasher.update(gs.as_bytes());
    // This runs in constant time
    if hasher.finalize() == message[32..64] {
        Ok(key)
    } else {
        Err(KKTError::X25519Error {
            info: format!("Symmetric Key Hash Validation Error"),
        })
    }
 }
 fn encrypt(mut key: [u8; 32], message: &[u8]) -> Result<Vec<u8>, KKTError> {
    // The empty nonce is fine since we use the key once.
    let nonce = Nonce::<Aes256GcmSiv>::from_slice(&[]);
    let ciphertext =
        nym_crypto::symmetric::aead::encrypt::<Aes256GcmSiv>(&key.into(), nonce, message)?;
    key.zeroize();
    Ok(ciphertext)
 }
 fn decrypt(key: [u8; 32], ciphertext: Vec<u8>) -> Vec<u8> {
    // The empty nonce is fine since we use the key once.
    let nonce = Nonce::<Aes256>::from_slice(&[]);
    let ciphertext =
        nym_crypto::symmetric::aead::encrypt::<Aes256GcmSiv>(&key.into(), nonce, message)?;
    key.zeroize();
    Ok(ciphertext)
 }
@@ -0,0 +1,90 @@
 // Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 use thiserror::Error;
 use crate::context::KKTStatus;
 #[derive(Error, Debug)]
 pub enum KKTError {
    #[error("Signature constructor error")]
    SigConstructorError,
    #[error("Signature verification error")]
    SigVerifError,
    #[error("Ciphersuite Decoding Error: {}", info)]
    CiphersuiteDecodingError { info: String },
    #[error("Insecure Encapsulation Key Hash Length")]
    InsecureHashLen,
    #[error("KKT Frame Decoding Error: {}", info)]
    FrameDecodingError { info: String },
    #[error("KKT Frame Encoding Error: {}", info)]
    FrameEncodingError { info: String },
    #[error("KKT Incompatibility Error: {}", info)]
    IncompatibilityError { info: &'static str },
    #[error("KKT Responder Flagged Error: {}", status)]
    ResponderFlaggedError { status: KKTStatus },
    #[error("KKT Message Count Limit Reached")]
    MessageCountLimitReached,
    #[error("PSQ KEM Error: {}", info)]
    KEMError { info: &'static str },
    #[error("Local Function Input Error: {}", info)]
    FunctionInputError { info: &'static str },
    #[error("{}", info)]
    X25519Error { info: &'static str },
    // #[error("Protocol did not complete")]
    // ProtocolError,
    // #[error("encountered an IO error: {0}")]
    // IoError(#[from] io::Error),s
    // #[error("Handshake timeout")]
    // HandshakeTimeout(#[from] tokio::time::error::Elapsed),
 }
 impl From<libcrux_kem::Error> for KKTError {
    fn from(err: libcrux_kem::Error) -> Self {
        match err {
            libcrux_kem::Error::EcDhError(_) => KKTError::KEMError { info: "ECDH Error" },
            libcrux_kem::Error::KeyGen => KKTError::KEMError {
                info: "Key Generation Error",
            },
            libcrux_kem::Error::Encapsulate => KKTError::KEMError {
                info: "Encapsulation Error",
            },
            libcrux_kem::Error::Decapsulate => KKTError::KEMError {
                info: "Decapsulation Error",
            },
            libcrux_kem::Error::UnsupportedAlgorithm => KKTError::KEMError {
                info: "libcrux Unsupported Algorithm",
            },
            libcrux_kem::Error::InvalidPrivateKey => KKTError::KEMError {
                info: "Invalid Private Key",
            },
            libcrux_kem::Error::InvalidPublicKey => KKTError::KEMError {
                info: "Invalid Public Key",
            },
            libcrux_kem::Error::InvalidCiphertext => KKTError::KEMError {
                info: "Invalid Ciphertext",
            },
        }
    }
 }
 impl From<libcrux_ecdh::Error> for KKTError {
    fn from(err: libcrux_ecdh::Error) -> Self {
        match err {
            libcrux_ecdh::Error::InvalidPoint => KKTError::KEMError {
                info: "Invalid Remote Public Key",
            },
            _ => todo!(),
        }
    }
 }
@@ -0,0 +1,146 @@
 // Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 // | 0 | 1 | 2, 3, 4, 5 | 6 | 7
 // [0] => KKT version (4 bits) + Message Sequence Count (4 bits)
 // [1] => Status (3 bits) + Mode (3 bits) + Role (2 bits)
 // [2..=5] => Ciphersuite
 // [6] => Reserved
 use nym_crypto::asymmetric::x25519::PublicKey;
 use crate::{
    context::{KKTContext, KKT_CONTEXT_LEN},
    error::KKTError,
 };
 pub const KKT_SESSION_ID_LEN: usize = 16;
 pub struct KKTFrame {
    context: Vec<u8>,
    session_id: Vec<u8>,
    body: Vec<u8>,
    signature: Vec<u8>,
 }
 // if oneway and message coming from initiator => body is empty, signature contains signature of context + session id (64 bytes).
 // if message coming from anonymous initiator => body is empty, there is no signature.
 // if mutual and message coming from initiator => body has the initiator's kem public key and the signature is over the context + body + session_id.
 // if coming from responder => body has the responder's kem public key and the signature is over the context + body + session_id.
 impl KKTFrame {
    pub fn new(context: &[u8], body: &[u8], session_id: &[u8], signature: &[u8]) -> Self {
        Self {
            context: Vec::from(context),
            body: Vec::from(body),
            session_id: Vec::from(session_id),
            signature: Vec::from(signature),
        }
    }
    pub fn context_ref<'a>(&'a self) -> &'a [u8] {
        &self.context
    }
    pub fn signature_ref<'a>(&'a self) -> &'a [u8] {
        &self.signature
    }
    pub fn body_ref<'a>(&'a self) -> &'a [u8] {
        &self.body
    }
    pub fn session_id_ref<'a>(&'a self) -> &'a [u8] {
        &self.session_id
    }
    pub fn signature_mut<'a>(&'a mut self) -> &'a mut [u8] {
        &mut self.signature
    }
    pub fn body_mut<'a>(&'a mut self) -> &'a mut [u8] {
        &mut self.body
    }
    pub fn session_id_mut<'a>(&'a mut self) -> &'a mut [u8] {
        &mut self.session_id
    }
    pub fn frame_length(&self) -> usize {
        self.context.len() + self.session_id.len() + self.body.len() + self.signature.len()
    }
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut bytes = Vec::with_capacity(self.frame_length());
        bytes.extend_from_slice(&self.context);
        bytes.extend_from_slice(&self.body);
        bytes.extend_from_slice(&self.session_id);
        bytes.extend_from_slice(&self.signature);
        bytes
    }
    pub fn from_bytes(bytes: &[u8]) -> Result<(Self, KKTContext), KKTError> {
        if bytes.len() < KKT_CONTEXT_LEN {
            return Err(KKTError::FrameDecodingError {
                info: format!(
                    "Frame is shorter than expected context length: actual {} != expected {}",
                    bytes.len(),
                    KKT_CONTEXT_LEN
                ),
            });
        } else {
            let context_bytes = Vec::from(&bytes[0..KKT_CONTEXT_LEN]);
            let context = KKTContext::try_decode(&context_bytes)?;
            let (mut session_id, mut body, mut signature): (Vec<u8>, Vec<u8>, Vec<u8>) =
                (vec![], vec![], vec![]);
            if bytes.len() == context.full_message_len() {
                if context.body_len() > 0 {
                    body.extend_from_slice(
                        &bytes[KKT_CONTEXT_LEN..KKT_CONTEXT_LEN + context.body_len()],
                    );
                }
                if context.session_id_len() > 0 {
                    session_id.extend_from_slice(
                        &bytes[KKT_CONTEXT_LEN + context.body_len()
                            ..KKT_CONTEXT_LEN + context.body_len() + context.session_id_len()],
                    );
                }
                if context.signature_len() > 0 {
                    signature.extend_from_slice(
                        &bytes[KKT_CONTEXT_LEN + context.body_len() + context.session_id_len()
                            ..KKT_CONTEXT_LEN
                                + context.body_len()
                                + context.session_id_len()
                                + context.signature_len()],
                    );
                }
                Ok((
                    KKTFrame::new(&context_bytes, &body, &session_id, &signature),
                    context,
                ))
            } else {
                return Err(KKTError::FrameDecodingError {
                    info: format!(
                        "Frame is shorter than expected: actual {} != expected {}",
                        bytes.len(),
                        context.full_message_len()
                    ),
                });
            }
        }
    }
 }
 pub struct EncryptedKKTRequest {
    ephermeral_key: PublicKey,
    body: Vec<u8>,
 }
 impl EncryptedKKTRequest {
    pub fn to_bytes(mut self) -> Vec<u8> {
        self.body.extend(self.ephermeral_key.to_bytes());
        self.body
    }
 }
 pub struct EncryptedKKTResponse {
    body: Vec<u8>,
 }
@@ -0,0 +1,107 @@
 use crate::{
    ciphersuite::{HashFunction, KEM},
    error::KKTError,
 };
 use classic_mceliece_rust::keypair_boxed;
 use libcrux_kem::{key_gen, Algorithm};
 use libcrux_sha3;
 use rand::{CryptoRng, RngCore};
 // (decapsulation_key, encapsulation_key)
 pub fn generate_keypair_libcrux<R>(
    rng: &mut R,
    kem: KEM,
 ) -> Result<(libcrux_kem::PrivateKey, libcrux_kem::PublicKey), KKTError>
 where
    R: RngCore + CryptoRng,
 {
    match kem {
        KEM::MlKem768 => Ok(key_gen(Algorithm::MlKem768, rng)?),
        KEM::XWing => Ok(key_gen(Algorithm::XWingKemDraft06, rng)?),
        KEM::X25519 => Ok(key_gen(Algorithm::X25519, rng)?),
        _ => Err(KKTError::KEMError {
            info: "Key Generation Error: Unsupported Libcrux Algorithm",
        }),
    }
 }
 // (decapsulation_key, encapsulation_key)
 pub fn generate_keypair_mceliece<'a, R>(
    rng: &mut R,
 ) -> (
    classic_mceliece_rust::SecretKey<'a>,
    classic_mceliece_rust::PublicKey<'a>,
 )
 where
    // this is annoying because mceliece lib uses rand 0.8.5...
    R: RngCore + CryptoRng,
 {
    let (encapsulation_key, decapsulation_key) = keypair_boxed(rng);
    (decapsulation_key, encapsulation_key)
 }
 pub fn hash_key_bytes(
    hash_function: &HashFunction,
    hash_length: usize,
    key_bytes: &[u8],
 ) -> Vec<u8> {
    let mut hashed_key: Vec<u8> = vec![0u8; hash_length];
    match hash_function {
        HashFunction::Blake3 => {
            let mut hasher = blake3::Hasher::new();
            hasher.update(key_bytes);
            hasher.finalize_xof().fill(&mut hashed_key);
            hasher.reset();
        }
        HashFunction::SHAKE256 => {
            libcrux_sha3::shake256_ema(&mut hashed_key, &key_bytes);
        }
        HashFunction::SHAKE128 => {
            libcrux_sha3::shake128_ema(&mut hashed_key, &key_bytes);
        }
        HashFunction::SHA256 => {
            libcrux_sha3::sha256_ema(&mut hashed_key, &key_bytes);
        }
    }
    hashed_key
 }
 /// This does NOT run in constant time.
 // It's fine for KKT since we are comparing hashes.
 fn compare_hashes(a: &[u8], b: &[u8]) -> bool {
    a == b
 }
 pub fn validate_encapsulation_key(
    hash_function: &HashFunction,
    hash_length: usize,
    encapsulation_key: &[u8],
    expected_hash_bytes: &[u8],
 ) -> bool {
    compare_hashes(
        &hash_encapsulation_key(hash_function, hash_length, encapsulation_key),
        expected_hash_bytes,
    )
 }
 pub fn validate_key_bytes(
    hash_function: &HashFunction,
    hash_length: usize,
    key_bytes: &[u8],
    expected_hash_bytes: &[u8],
 ) -> bool {
    compare_hashes(
        &hash_key_bytes(hash_function, hash_length, key_bytes),
        expected_hash_bytes,
    )
 }
 pub fn hash_encapsulation_key(
    hash_function: &HashFunction,
    hash_length: usize,
    encapsulation_key: &[u8],
 ) -> Vec<u8> {
    hash_key_bytes(hash_function, hash_length, &encapsulation_key)
 }
@@ -0,0 +1,231 @@
 // Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 pub mod ciphersuite;
 pub mod context;
 // pub mod encryption;
 pub mod error;
 pub mod frame;
 pub mod key_utils;
 pub mod session;
 // pub mod psq;
 // This must be less than 4 bits
 pub const KKT_VERSION: u8 = 1;
 const _: () = assert!(KKT_VERSION < 1 << 4);
 #[cfg(test)]
 mod test {
    use nym_crypto::asymmetric::ed25519;
    use rand::prelude::*;
    use crate::{
        ciphersuite::{Ciphersuite, EncapsulationKey, HashFunction, KEM},
        frame::KKTFrame,
        key_utils::{generate_keypair_libcrux, generate_keypair_mceliece, hash_encapsulation_key},
        session::{
            anonymous_initiator_process, initiator_ingest_response, initiator_process,
            responder_ingest_message, responder_process,
        },
    };
    #[test]
    fn test_kkt_psq_e2e_clear() {
        let mut rng = rand::rng();
        // generate ed25519 keys
        let mut secret_initiator: [u8; 32] = [0u8; 32];
        rng.fill_bytes(&mut secret_initiator);
        let initiator_ed25519_keypair = ed25519::KeyPair::from_secret(secret_initiator, 0);
        let mut secret_responder: [u8; 32] = [0u8; 32];
        rng.fill_bytes(&mut secret_responder);
        let responder_ed25519_keypair = ed25519::KeyPair::from_secret(secret_responder, 1);
        for kem in [KEM::MlKem768, KEM::XWing, KEM::X25519, KEM::McEliece] {
            for hash_function in [
                HashFunction::Blake3,
                HashFunction::SHA256,
                HashFunction::SHAKE128,
                HashFunction::SHAKE256,
            ] {
                let ciphersuite = Ciphersuite::resolve_ciphersuite(
                    kem,
                    hash_function,
                    crate::ciphersuite::SignatureScheme::Ed25519,
                    None,
                )
                .unwrap();
                // generate kem public keys
                let (responder_kem_public_key, initiator_kem_public_key) = match kem {
                    KEM::MlKem768 => (
                        EncapsulationKey::MlKem768(
                            generate_keypair_libcrux(&mut rng, kem).unwrap().1,
                        ),
                        EncapsulationKey::MlKem768(
                            generate_keypair_libcrux(&mut rng, kem).unwrap().1,
                        ),
                    ),
                    KEM::XWing => (
                        EncapsulationKey::XWing(generate_keypair_libcrux(&mut rng, kem).unwrap().1),
                        EncapsulationKey::XWing(generate_keypair_libcrux(&mut rng, kem).unwrap().1),
                    ),
                    KEM::X25519 => (
                        EncapsulationKey::X25519(
                            generate_keypair_libcrux(&mut rng, kem).unwrap().1,
                        ),
                        EncapsulationKey::X25519(
                            generate_keypair_libcrux(&mut rng, kem).unwrap().1,
                        ),
                    ),
                    KEM::McEliece => (
                        EncapsulationKey::McEliece(generate_keypair_mceliece(&mut rng).1),
                        EncapsulationKey::McEliece(generate_keypair_mceliece(&mut rng).1),
                    ),
                };
                let i_kem_key_bytes = initiator_kem_public_key.encode();
                let r_kem_key_bytes = responder_kem_public_key.encode();
                let i_dir_hash = hash_encapsulation_key(
                    &ciphersuite.hash_function(),
                    ciphersuite.hash_len(),
                    &i_kem_key_bytes,
                );
                let r_dir_hash = hash_encapsulation_key(
                    &ciphersuite.hash_function(),
                    ciphersuite.hash_len(),
                    &r_kem_key_bytes,
                );
                // Anonymous Initiator, OneWay
                {
                    let (mut i_context, i_frame) =
                        anonymous_initiator_process(&mut rng, ciphersuite).unwrap();
                    let i_frame_bytes = i_frame.to_bytes();
                    let (i_frame_r, r_context) = KKTFrame::from_bytes(&i_frame_bytes).unwrap();
                    let (mut r_context, _) =
                        responder_ingest_message(&r_context, None, None, &i_frame_r).unwrap();
                    let r_frame = responder_process(
                        &mut r_context,
                        i_frame_r.session_id_ref(),
                        responder_ed25519_keypair.private_key(),
                        &responder_kem_public_key,
                    )
                    .unwrap();
                    let r_bytes = r_frame.to_bytes();
                    let obtained_key = initiator_ingest_response(
                        &mut i_context,
                        responder_ed25519_keypair.public_key(),
                        &r_dir_hash,
                        &r_bytes,
                    )
                    .unwrap();
                    assert_eq!(obtained_key.encode(), r_kem_key_bytes)
                }
                // Initiator, OneWay
                {
                    let (mut i_context, i_frame) = initiator_process(
                        &mut rng,
                        crate::context::KKTMode::OneWay,
                        ciphersuite,
                        initiator_ed25519_keypair.private_key(),
                        None,
                    )
                    .unwrap();
                    let i_frame_bytes = i_frame.to_bytes();
                    let (i_frame_r, r_context) = KKTFrame::from_bytes(&i_frame_bytes).unwrap();
                    let (mut r_context, r_obtained_key) = responder_ingest_message(
                        &r_context,
                        Some(initiator_ed25519_keypair.public_key()),
                        None,
                        &i_frame_r,
                    )
                    .unwrap();
                    assert!(r_obtained_key.is_none());
                    let r_frame = responder_process(
                        &mut r_context,
                        i_frame_r.session_id_ref(),
                        responder_ed25519_keypair.private_key(),
                        &responder_kem_public_key,
                    )
                    .unwrap();
                    let r_bytes = r_frame.to_bytes();
                    let i_obtained_key = initiator_ingest_response(
                        &mut i_context,
                        responder_ed25519_keypair.public_key(),
                        &r_dir_hash,
                        &r_bytes,
                    )
                    .unwrap();
                    assert_eq!(i_obtained_key.encode(), r_kem_key_bytes)
                }
                // Initiator, Mutual
                {
                    let (mut i_context, i_frame) = initiator_process(
                        &mut rng,
                        crate::context::KKTMode::Mutual,
                        ciphersuite,
                        initiator_ed25519_keypair.private_key(),
                        Some(&initiator_kem_public_key),
                    )
                    .unwrap();
                    let i_frame_bytes = i_frame.to_bytes();
                    let (i_frame_r, r_context) = KKTFrame::from_bytes(&i_frame_bytes).unwrap();
                    let (mut r_context, r_obtained_key) = responder_ingest_message(
                        &r_context,
                        Some(initiator_ed25519_keypair.public_key()),
                        Some(&i_dir_hash),
                        &i_frame_r,
                    )
                    .unwrap();
                    assert_eq!(r_obtained_key.unwrap().encode(), i_kem_key_bytes);
                    let r_frame = responder_process(
                        &mut r_context,
                        i_frame_r.session_id_ref(),
                        responder_ed25519_keypair.private_key(),
                        &responder_kem_public_key,
                    )
                    .unwrap();
                    let r_bytes = r_frame.to_bytes();
                    let obtained_key = initiator_ingest_response(
                        &mut i_context,
                        responder_ed25519_keypair.public_key(),
                        &r_dir_hash,
                        &r_bytes,
                    )
                    .unwrap();
                    assert_eq!(obtained_key.encode(), r_kem_key_bytes)
                }
            }
        }
    }
 }
@@ -0,0 +1,228 @@
 use nym_crypto::asymmetric::ed25519::{self, Signature};
 use rand::{CryptoRng, RngCore};
 use crate::{
    ciphersuite::{Ciphersuite, EncapsulationKey},
    context::{KKTContext, KKTMode, KKTRole, KKTStatus},
    error::KKTError,
    frame::{KKTFrame, KKT_SESSION_ID_LEN},
    key_utils::validate_encapsulation_key,
 };
 pub fn initiator_process<'a, R>(
    rng: &mut R,
    mode: KKTMode,
    ciphersuite: Ciphersuite,
    signing_key: &ed25519::PrivateKey,
    own_encapsulation_key: Option<&EncapsulationKey<'a>>,
 ) -> Result<(KKTContext, KKTFrame), KKTError>
 where
    R: CryptoRng + RngCore,
 {
    let context = KKTContext::new(KKTRole::Initiator, mode, ciphersuite)?;
    let context_bytes = context.encode()?;
    let mut session_id = [0; KKT_SESSION_ID_LEN];
    // Generate Session ID
    rng.fill_bytes(&mut session_id);
    let body: &[u8] = match mode {
        KKTMode::OneWay => &[],
        KKTMode::Mutual => match own_encapsulation_key {
            Some(encaps_key) => &encaps_key.encode(),
            // Missing key
            None => {
                return Err(KKTError::FunctionInputError {
                    info: "KEM Key Not Provided",
                })
            }
        },
    };
    let mut bytes_to_sign =
        Vec::with_capacity(context.full_message_len() - context.signature_len());
    bytes_to_sign.extend_from_slice(&context_bytes);
    bytes_to_sign.extend_from_slice(&body);
    bytes_to_sign.extend_from_slice(&session_id);
    let signature = signing_key.sign(bytes_to_sign).to_bytes();
    Ok((
        context,
        KKTFrame::new(&context_bytes, &body, &session_id, &signature),
    ))
 }
 pub fn anonymous_initiator_process<R>(
    rng: &mut R,
    ciphersuite: Ciphersuite,
 ) -> Result<(KKTContext, KKTFrame), KKTError>
 where
    R: CryptoRng + RngCore,
 {
    let context = KKTContext::new(KKTRole::AnonymousInitiator, KKTMode::OneWay, ciphersuite)?;
    let context_bytes = context.encode()?;
    let mut session_id = [0u8; KKT_SESSION_ID_LEN];
    rng.fill_bytes(&mut session_id);
    Ok((
        context,
        KKTFrame::new(&context_bytes, &[], &session_id, &[]),
    ))
 }
 pub fn initiator_ingest_response<'a>(
    own_context: &mut KKTContext,
    remote_verification_key: &ed25519::PublicKey,
    expected_hash: &[u8],
    message_bytes: &[u8],
 ) -> Result<EncapsulationKey<'a>, KKTError> {
    // sizes have to be correct
    let (frame, remote_context) = KKTFrame::from_bytes(message_bytes)?;
    check_compatibility(&own_context, &remote_context)?;
    match remote_context.status() {
        KKTStatus::Ok => {
            let mut bytes_to_verify: Vec<u8> = Vec::with_capacity(
                remote_context.full_message_len() - remote_context.signature_len(),
            );
            bytes_to_verify.extend_from_slice(&remote_context.encode()?);
            bytes_to_verify.extend_from_slice(frame.body_ref());
            bytes_to_verify.extend_from_slice(frame.session_id_ref());
            match Signature::from_bytes(frame.signature_ref()) {
                Ok(sig) => match remote_verification_key.verify(bytes_to_verify, &sig) {
                    Ok(()) => {
                        let received_encapsulation_key = EncapsulationKey::decode(
                            own_context.ciphersuite().kem(),
                            frame.body_ref(),
                        )?;
                        match validate_encapsulation_key(
                    &own_context.ciphersuite().hash_function(),
                    own_context.ciphersuite().hash_len(),
                    &frame.body_ref(),
                    expected_hash,
                ) {
                    true => Ok(received_encapsulation_key),
                    // The key does not match the hash obtained from the directory
                    false => Err(KKTError::KEMError { info: "Hash of received encapsulation key does not match the value stored on the directory." }),
                }
                    }
                    Err(_) => Err(KKTError::SigVerifError),
                },
                Err(_) => Err(KKTError::SigConstructorError),
            }
        }
        _ => Err(KKTError::ResponderFlaggedError {
            status: remote_context.status(),
        }),
    }
 }
 // todo: figure out how to handle errors using status codes
 pub fn responder_ingest_message<'a>(
    remote_context: &KKTContext,
    remote_verification_key: Option<&ed25519::PublicKey>,
    expected_hash: Option<&[u8]>,
    remote_frame: &KKTFrame,
 ) -> Result<(KKTContext, Option<EncapsulationKey<'a>>), KKTError> {
    let own_context = remote_context.derive_responder_header()?;
    match remote_context.role() {
        KKTRole::AnonymousInitiator => Ok((own_context, None)),
        KKTRole::Initiator => {
            match remote_verification_key {
                Some(remote_verif_key) => {
                    let mut bytes_to_verify: Vec<u8> = Vec::with_capacity(
                        own_context.full_message_len() - own_context.signature_len(),
                    );
                    bytes_to_verify.extend_from_slice(remote_frame.context_ref());
                    bytes_to_verify.extend_from_slice(remote_frame.body_ref());
                    bytes_to_verify.extend_from_slice(remote_frame.session_id_ref());
                    match Signature::from_bytes(remote_frame.signature_ref()) {
                        Ok(sig) => match remote_verif_key.verify(bytes_to_verify, &sig) {
                            Ok(()) => {
                                // using own_context here because maybe for whatever reason we want to ignore the remote kem key
                                match own_context.mode() {
                                    KKTMode::OneWay => return Ok((own_context, None)),
                                    KKTMode::Mutual => {
                                        match expected_hash {
    Some(expected_hash) =>{
      let received_encapsulation_key =
                    EncapsulationKey::decode(own_context.ciphersuite().kem(), remote_frame.body_ref())?;
                    if
                validate_encapsulation_key(
                    &own_context.ciphersuite().hash_function(),
                    own_context.ciphersuite().hash_len(),
                    &remote_frame.body_ref(),
                    expected_hash,
                ){
                    Ok((own_context, Some(received_encapsulation_key)))
                }
                    // The key does not match the hash obtained from the directory
                  else {
                        Err(KKTError::KEMError { info: "Hash of received encapsulation key does not match the value stored on the directory." })
                    }
            }
            None => Err(KKTError::FunctionInputError { info: "Expected hash of the remote encapsulation key is not provided." }),
        }
                                    }
                                }
                            }
                            Err(_) => Err(KKTError::SigVerifError),
                        },
                        Err(_) => Err(KKTError::SigConstructorError),
                    }
                }
                None => Err(KKTError::FunctionInputError {
                    info: "Remote Signature Verification Key Not Provided",
                }),
            }
        }
        KKTRole::Responder => Err(KKTError::IncompatibilityError {
            info: "Responder received a request from another responder.",
        }),
    }
 }
 pub fn responder_process<'a>(
    own_context: &mut KKTContext,
    session_id: &[u8],
    signing_key: &ed25519::PrivateKey,
    encapsulation_key: &EncapsulationKey<'a>,
 ) -> Result<KKTFrame, KKTError> {
    let body = encapsulation_key.encode();
    let context_bytes = own_context.encode()?;
    let mut bytes_to_sign =
        Vec::with_capacity(own_context.full_message_len() - own_context.signature_len());
    bytes_to_sign.extend_from_slice(&own_context.encode()?);
    bytes_to_sign.extend_from_slice(&body);
    bytes_to_sign.extend_from_slice(&session_id);
    let signature = signing_key.sign(bytes_to_sign).to_bytes();
    Ok(KKTFrame::new(
        &context_bytes,
        &body,
        &session_id,
        &signature,
    ))
 }
 fn check_compatibility(
    _own_context: &KKTContext,
    _remote_context: &KKTContext,
 ) -> Result<(), KKTError> {
    // todo: check ciphersuite/context compatibility
    Ok(())
 }
@@ -22,6 +22,12 @@ tokio-util = { workspace = true, features = ["codec"] }
 nym-crypto = { path = "../crypto" }
 nym-noise-keys = { path = "keys" }
 nympsq = {path = "../nympsq"}
 libcrux-ed25519 = { git = "https://github.com/cryspen/libcrux" }
 libcrux-psq = { git = "https://github.com/cryspen/libcrux", features = ["test-utils"] }
 libcrux-kem = { git = "https://github.com/cryspen/libcrux" }
 rand = "0.9.2"
 [dev-dependencies]
 anyhow = { workspace = true }
 tokio = { workspace = true, features = ["full"] }
@@ -9,6 +9,7 @@ use serde::{Deserialize, Serialize};
 #[serde(from = "u8", into = "u8")]
 pub enum NoiseVersion {
    V1,
    V2,
    Unknown(u8), //Implies a newer version we don't know
 }
@@ -16,6 +17,7 @@ impl From<u8> for NoiseVersion {
    fn from(value: u8) -> Self {
        match value {
            1 => NoiseVersion::V1,
            2 => NoiseVersion::V2,
            other => NoiseVersion::Unknown(other),
        }
    }
@@ -25,6 +27,7 @@ impl From<NoiseVersion> for u8 {
    fn from(version: NoiseVersion) -> Self {
        match version {
            NoiseVersion::V1 => 1,
            NoiseVersion::V2 => 2,
            NoiseVersion::Unknown(other) => other,
        }
    }
@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 use nym_noise_keys::NoiseVersion;
 use nympsq::error::PSQError;
 use snow::Error;
 use std::io;
 use thiserror::Error;
@@ -70,6 +71,12 @@ pub enum NoiseError {
    #[error("Handshake timeout")]
    HandshakeTimeout(#[from] tokio::time::error::Elapsed),
    #[error("Missing Field")]
    MissingField { info: &'static str },
    #[error("PSQ Error")]
    PSQError(PSQError),
 }
 impl NoiseError {
@@ -81,6 +88,18 @@ impl NoiseError {
    }
 }
 impl From<PSQError> for NoiseError {
    fn from(err: PSQError) -> Self {
        NoiseError::PSQError(err)
    }
 }
 impl From<libcrux_kem::Error> for NoiseError {
    fn from(err: libcrux_kem::Error) -> Self {
        NoiseError::PSQError(err.into())
    }
 }
 impl From<Error> for NoiseError {
    fn from(err: Error) -> Self {
        match err {
@@ -1,9 +1,11 @@
 // Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 use nym_crypto::asymmetric::ed25519;
 use nym_noise_keys::NoiseVersion;
-use snow::error::Prerequisite;
+use nympsq::psq::CONTEXT_LEN;
-use snow::Error;
+
 use snow::{error::Prerequisite, Error};
 use tokio::net::TcpStream;
 use tracing::{error, warn};
@@ -12,48 +14,138 @@ pub mod connection;
 pub mod error;
 pub mod stream;
-use crate::config::NoiseConfig;
+use crate::{
-use crate::connection::Connection;
+    config::NoiseConfig, connection::Connection, error::NoiseError, stream::NoiseStreamBuilder,
-use crate::error::NoiseError;
+};
 use crate::stream::NoiseStreamBuilder;
 const NOISE_PSK_PREFIX: &[u8] = b"NYMTECH_NOISE_dQw4w9WgXcQ";
-pub const LATEST_NOISE_VERSION: NoiseVersion = NoiseVersion::V1;
+pub(crate) const NOISE_PSQ_DEFAULT_CONTEXT: &[u8; CONTEXT_LEN] = b"Exsl88AD2ccS99kk";
 pub(crate) const NOISE_PSQ_DEFAULT_DURATION_SECS: u64 = 1000;
 pub const LATEST_NOISE_VERSION: NoiseVersion = NoiseVersion::V2;
 // TODO: this should be behind some trait because presumably, depending on the version,
 // other arguments would be needed
 mod psk_gen {
-    use crate::error::NoiseError;
+    use std::time::Duration;
-    use crate::stream::Psk;
+
-    use crate::NOISE_PSK_PREFIX;
+    use crate::{error::NoiseError, stream::Psk, NOISE_PSK_PREFIX};
-    use nym_crypto::asymmetric::x25519;
+    use libcrux_ed25519::VerificationKey;
-    use nym_noise_keys::NoiseVersion;
+    use libcrux_kem::{PrivateKey, PublicKey};
    use libcrux_psq::impls::{MlKem768, XWingKemDraft06, X25519};
    use nym_crypto::asymmetric::ed25519;
    use nympsq::{
        error::PSQError,
        psq::{PSQInitiator, PSQResponder, CONTEXT_LEN, PSK_HANDLE_LEN},
    };
    use sha2::{Digest, Sha256};
-    pub(crate) fn generate_psk(
+    pub fn generate_psk_v1(responder_pub_key: &[u8]) -> [u8; 32] {
        responder_pub_key: x25519::PublicKey,
        version: NoiseVersion,
    ) -> Result<Psk, NoiseError> {
        match version {
            NoiseVersion::V1 => Ok(generate_psk_v1(responder_pub_key)),
            NoiseVersion::Unknown(noise_version) => {
                Err(NoiseError::PskGenerationFailure { noise_version })
            }
        }
    }
    fn generate_psk_v1(responder_pub_key: x25519::PublicKey) -> [u8; 32] {
        let mut hasher = Sha256::new();
        hasher.update(NOISE_PSK_PREFIX);
-        hasher.update(responder_pub_key.to_bytes());
+        hasher.update(responder_pub_key);
        hasher.finalize().into()
    }
    pub(crate) fn psq_initiate_x25519(
        initiator: &mut PSQInitiator<X25519>,
        responder_pub_key: impl AsRef<[u8]>,
        context: &[u8; CONTEXT_LEN],
        psq_ttl: Duration,
    ) -> Result<Vec<u8>, NoiseError> {
        let pub_key =
            PublicKey::decode(libcrux_kem::Algorithm::X25519, responder_pub_key.as_ref())?;
        let message =
            initiator.compute_initiator_message(&mut rand::rng(), &pub_key, context, psq_ttl)?;
        Ok(message)
    }
    pub(crate) fn psq_respond_x25519(
        responder_private_key: impl AsRef<[u8]>,
        responder_public_key: impl AsRef<[u8]>,
        initiator_verification_key: &ed25519::PublicKey,
        initiator_message: &mut [u8],
        context: &[u8; CONTEXT_LEN],
        psq_ttl: Duration,
        psk_handle: &[u8; PSK_HANDLE_LEN],
    ) -> Result<(Psk, Vec<u8>), PSQError> {
        let kem_private_key = PrivateKey::decode(
            libcrux_kem::Algorithm::X25519,
            responder_private_key.as_ref(),
        )?;
        let kem_public_key = PublicKey::decode(
            libcrux_kem::Algorithm::X25519,
            responder_public_key.as_ref(),
        )?;
        let responder: PSQResponder<X25519> = PSQResponder::init(&kem_private_key, &kem_public_key);
        let res = responder.compute_responder_message(
            &VerificationKey::from_bytes(initiator_verification_key.to_bytes()),
            initiator_message,
            context,
            psq_ttl,
            psk_handle,
        )?;
        Ok(res)
    }
    pub(crate) fn psq_initiate_xwing(
        initiator: &mut PSQInitiator<XWingKemDraft06>,
        responder_pub_key: impl AsRef<[u8]>,
        context: &[u8; CONTEXT_LEN],
        psq_ttl: Duration,
    ) -> Result<Vec<u8>, NoiseError> {
        let pub_key = PublicKey::decode(
            libcrux_kem::Algorithm::XWingKemDraft06,
            responder_pub_key.as_ref(),
        )?;
        let message =
            initiator.compute_initiator_message(&mut rand::rng(), &pub_key, context, psq_ttl)?;
        Ok(message)
    }
    pub(crate) fn psq_respond_xwing(
        responder_private_key: impl AsRef<[u8]>,
        responder_public_key: impl AsRef<[u8]>,
        initiator_verification_key: &ed25519::PublicKey,
        initiator_message: &mut [u8],
        context: &[u8; CONTEXT_LEN],
        psq_ttl: Duration,
        psk_handle: &[u8; PSK_HANDLE_LEN],
    ) -> Result<(Psk, Vec<u8>), PSQError> {
        let kem_private_key = PrivateKey::decode(
            libcrux_kem::Algorithm::XWingKemDraft06,
            responder_private_key.as_ref(),
        )?;
        let kem_public_key = PublicKey::decode(
            libcrux_kem::Algorithm::XWingKemDraft06,
            responder_public_key.as_ref(),
        )?;
        let responder: PSQResponder<XWingKemDraft06> =
            PSQResponder::init(&kem_private_key, &kem_public_key);
        let res = responder.compute_responder_message(
            &VerificationKey::from_bytes(initiator_verification_key.to_bytes()),
            initiator_message,
            context,
            psq_ttl,
            psk_handle,
        )?;
        Ok(res)
    }
 }
 pub async fn upgrade_noise_initiator(
    conn: TcpStream,
    config: &NoiseConfig,
    initiator_identity_keypair: Option<ed25519::KeyPair>,
 ) -> Result<Connection<TcpStream>, NoiseError> {
    if config.unsafe_disabled {
        warn!("Noise is disabled in the config. Not attempting any handshake");
@@ -73,6 +165,7 @@ pub async fn upgrade_noise_initiator(
    let handshake_version = match key.supported_version {
        NoiseVersion::V1 => NoiseVersion::V1,
        NoiseVersion::V2 => NoiseVersion::V2,
        // We're talking to a more recent node, but we can't adapt. Let's try to do our best and if it fails, it fails.
        // If that node sees we're older, it will try to adapt too.
@@ -82,14 +175,29 @@ pub async fn upgrade_noise_initiator(
        }
    };
    let (signing_key, verification_key): (Option<[u8; 32]>, Option<[u8; 32]>) =
        match initiator_identity_keypair {
            Some(keypair) => (
                Some(keypair.private_key().to_bytes()),
                Some(keypair.public_key().to_bytes()),
            ),
            None => (None, None),
        };
    NoiseStreamBuilder::new(conn)
-        .perform_initiator_handshake(config, handshake_version, key.x25519_pubkey)
+        .perform_initiator_handshake(
            config,
            handshake_version,
            key.x25519_pubkey,
            signing_key,
            verification_key,
        )
        .await
        .map(|stream| Connection::Noise(Box::new(stream)))
 }
 pub async fn upgrade_noise_responder(
    conn: TcpStream,
    config: &NoiseConfig,
    initiator_verification_key: Option<ed25519::PublicKey>,
 ) -> Result<Connection<TcpStream>, NoiseError> {
    if config.unsafe_disabled {
        warn!("Noise is disabled in the config. Not attempting any handshake");
@@ -112,7 +220,7 @@ pub async fn upgrade_noise_responder(
    };
    NoiseStreamBuilder::new(conn)
-        .perform_responder_handshake(config)
+        .perform_responder_handshake(config, initiator_verification_key)
        .await
        .map(|stream| Connection::Noise(Box::new(stream)))
 }
@@ -1,20 +1,21 @@
 // Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
-use crate::config::{NoiseConfig, NoisePattern};
+use crate::{
-use crate::error::NoiseError;
+    config::{NoiseConfig, NoisePattern},
-use crate::psk_gen::generate_psk;
+    error::NoiseError,
-use crate::stream::codec::NymNoiseCodec;
+    psk_gen::{generate_psk_v1, psq_initiate_x25519, psq_respond_x25519},
-use crate::stream::framing::NymNoiseFrame;
+    stream::{codec::NymNoiseCodec, framing::NymNoiseFrame},
    NOISE_PSQ_DEFAULT_CONTEXT, NOISE_PSQ_DEFAULT_DURATION_SECS,
 };
 use bytes::{Bytes, BytesMut};
 use futures::{Sink, SinkExt, Stream, StreamExt};
-use nym_crypto::asymmetric::x25519;
+use nym_crypto::asymmetric::{ed25519, x25519};
 use nym_noise_keys::NoiseVersion;
 use nympsq::psq::{PSQInitiator, PSK_HANDLE_LEN};
 use snow::{Builder, HandshakeState, TransportState};
-use std::io;
+use std::{cmp::min, io, pin::Pin, task::ready, task::Poll, time::Duration};
 use std::pin::Pin;
 use std::task::Poll;
 use std::{cmp::min, task::ready};
 use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
 use tokio_util::codec::Framed;
@@ -22,6 +23,7 @@ mod codec;
 mod framing;
 const TAGLEN: usize = 16;
 // This needs to be increased when using psq
 const HANDSHAKE_MAX_LEN: usize = 1024; // using this constant to limit the handshake's buffer size
 pub(crate) type Psk = [u8; 32];
@@ -45,19 +47,65 @@ impl<C> NoiseStreamBuilder<C> {
        pattern: NoisePattern,
        local_private_key: impl AsRef<[u8]>,
        remote_pub_key: impl AsRef<[u8]>,
-        psk: Psk,
+        local_signing_key: Option<impl AsRef<[u8]>>,
        local_verification_key: Option<impl AsRef<[u8]>>,
        version: NoiseVersion,
    ) -> Result<NoiseStream<C>, NoiseError>
    where
        C: AsyncRead + AsyncWrite + Unpin,
    {
-        let handshake = Builder::new(pattern.as_noise_params())
+        let mut handshake = Builder::new(pattern.as_noise_params())
            .local_private_key(local_private_key.as_ref())
            .remote_public_key(remote_pub_key.as_ref())
            .psk(pattern.psk_position(), &psk)
            .build_initiator()?;
-        self.perform_handshake(handshake, version, pattern).await
+        let payload: Vec<u8> = match version {
            NoiseVersion::V1 => {
                handshake.set_psk(
                    pattern.psk_position() as usize,
                    &generate_psk_v1(remote_pub_key.as_ref()),
                )?;
                vec![]
            }
            NoiseVersion::V2 => match (local_signing_key, local_verification_key) {
                (Some(signing_key), Some(verification_key)) => {
                    let mut psq_initiator = PSQInitiator::init(signing_key, verification_key);
                    let payload = psq_initiate_x25519(
                        &mut psq_initiator,
                        &remote_pub_key,
                        NOISE_PSQ_DEFAULT_CONTEXT,
                        Duration::from_secs(NOISE_PSQ_DEFAULT_DURATION_SECS),
                    )?;
                    handshake.set_psk(
                        pattern.psk_position() as usize,
                        //
                        &psq_initiator.get_psk().unwrap(),
                    )?;
                    payload
                }
                (None, None) => {
                    return Err(NoiseError::MissingField {
                        info: "Missing Local Signing Key and Local Verification Key",
                    })
                }
                (None, _) => {
                    return Err(NoiseError::MissingField {
                        info: "Local Signing Key",
                    })
                }
                (_, None) => {
                    return Err(NoiseError::MissingField {
                        info: "Local Verification Key",
                    })
                }
            },
            NoiseVersion::Unknown(version) => {
                return Err(NoiseError::UnknownVersion { encoded: version })
            }
        };
        self.perform_handshake(handshake, payload, version, pattern)
            .await
    }
    pub(crate) async fn perform_initiator_handshake(
@@ -65,12 +113,12 @@ impl<C> NoiseStreamBuilder<C> {
        config: &NoiseConfig,
        version: NoiseVersion,
        remote_pub_key: x25519::PublicKey,
        local_signing_key: Option<impl AsRef<[u8]>>,
        local_verification_key: Option<impl AsRef<[u8]>>,
    ) -> Result<NoiseStream<C>, NoiseError>
    where
        C: AsyncRead + AsyncWrite + Unpin,
    {
        let psk = generate_psk(remote_pub_key, version)?;
        let timeout = config.timeout;
        tokio::time::timeout(
            timeout,
@@ -78,7 +126,8 @@ impl<C> NoiseStreamBuilder<C> {
                config.pattern,
                config.local_key.private_key(),
                remote_pub_key,
-                psk,
+                local_signing_key,
                local_verification_key,
                version,
            ),
        )
@@ -89,7 +138,8 @@ impl<C> NoiseStreamBuilder<C> {
        mut self,
        noise_pattern: NoisePattern,
        local_private_key: impl AsRef<[u8]>,
-        local_pub_key: x25519::PublicKey,
+        local_public_key: impl AsRef<[u8]>,
        initiator_verification_key: Option<ed25519::PublicKey>,
    ) -> Result<NoiseStream<C>, NoiseError>
    where
        C: AsyncRead + AsyncWrite + Unpin,
@@ -117,26 +167,58 @@ impl<C> NoiseStreamBuilder<C> {
            });
        }
        // 2. generate psk and handshake state
        let psk = generate_psk(local_pub_key, initial_frame.header.version)?;
        let mut handshake = Builder::new(pattern.as_noise_params())
            .local_private_key(local_private_key.as_ref())
            .psk(pattern.psk_position(), &psk)
            .build_responder()?;
        // update handshake state with initial frame
        let mut buf = BytesMut::zeroed(HANDSHAKE_MAX_LEN);
-        handshake.read_message(&initial_frame.data, &mut buf)?;
+        let len = handshake.read_message(&initial_frame.data, &mut buf)?;
        buf.truncate(len);
        let payload: Vec<u8> = match initial_frame.version() {
            NoiseVersion::V1 => {
                handshake.set_psk(
                    pattern.psk_position() as usize,
                    &generate_psk_v1(local_public_key.as_ref()),
                )?;
                vec![]
            }
            NoiseVersion::V2 => match initiator_verification_key {
                Some(verif_key) => {
                    let (psk, message) = psq_respond_x25519(
                        &local_private_key,
                        &local_public_key,
                        &verif_key,
                        &mut buf,
                        NOISE_PSQ_DEFAULT_CONTEXT,
                        Duration::from_secs(NOISE_PSQ_DEFAULT_DURATION_SECS),
                        &[1; PSK_HANDLE_LEN],
                    )?;
                    handshake.set_psk(pattern.psk_position() as usize, psk.as_slice())?;
                    message
                }
                None => {
                    return Err(NoiseError::MissingField {
                        info: "Initiator Verification Key",
                    })
                }
            },
            NoiseVersion::Unknown(version) => {
                return Err(NoiseError::UnknownVersion { encoded: version })
            }
        };
        // 3. run handshake to completion
-        self.perform_handshake(handshake, initial_frame.version(), pattern)
+        self.perform_handshake(handshake, payload, initial_frame.version(), pattern)
            .await
    }
    pub(crate) async fn perform_responder_handshake(
        self,
        config: &NoiseConfig,
        initiator_verification_key: Option<ed25519::PublicKey>,
    ) -> Result<NoiseStream<C>, NoiseError>
    where
        C: AsyncRead + AsyncWrite + Unpin,
@@ -147,7 +229,8 @@ impl<C> NoiseStreamBuilder<C> {
            self.perform_responder_handshake_inner(
                config.pattern,
                config.local_key.private_key(),
-                *config.local_key.public_key(),
+                config.local_key.public_key(),
                initiator_verification_key,
            ),
        )
        .await?
@@ -156,6 +239,7 @@ impl<C> NoiseStreamBuilder<C> {
    async fn send_handshake_msg(
        &mut self,
        handshake: &mut HandshakeState,
        payload: impl AsRef<[u8]>,
        version: NoiseVersion,
        pattern: NoisePattern,
    ) -> Result<(), NoiseError>
@@ -163,7 +247,7 @@ impl<C> NoiseStreamBuilder<C> {
        C: AsyncRead + AsyncWrite + Unpin,
    {
        let mut buf = BytesMut::zeroed(HANDSHAKE_MAX_LEN); // we're in the handshake, we can afford a smaller buffer
-        let len = handshake.write_message(&[], &mut buf)?;
+        let len = handshake.write_message(payload.as_ref(), &mut buf)?;
        buf.truncate(len);
        let frame = NymNoiseFrame::new_handshake_frame(buf.freeze(), version, pattern)?;
@@ -176,7 +260,7 @@ impl<C> NoiseStreamBuilder<C> {
        handshake: &mut HandshakeState,
        version: NoiseVersion,
        pattern: NoisePattern,
-    ) -> Result<(), NoiseError>
+    ) -> Result<Bytes, NoiseError>
    where
        C: AsyncRead + AsyncWrite + Unpin,
    {
@@ -199,9 +283,12 @@ impl<C> NoiseStreamBuilder<C> {
                    });
                }
-                let mut buf = BytesMut::zeroed(HANDSHAKE_MAX_LEN); // we're in the handshake, we can afford a smaller buffer
+                // The buffer below is the payload...
                // we're in the handshake, we can afford a smaller buffer
                let mut buf = BytesMut::zeroed(HANDSHAKE_MAX_LEN);
                handshake.read_message(&frame.data, &mut buf)?;
-                Ok(())
+                Ok(buf.into())
            }
            Some(Err(err)) => Err(err),
            None => Err(NoiseError::HandshakeError),
@@ -211,6 +298,7 @@ impl<C> NoiseStreamBuilder<C> {
    async fn perform_handshake(
        mut self,
        mut handshake_state: HandshakeState,
        payload: impl AsRef<[u8]>,
        version: NoiseVersion,
        pattern: NoisePattern,
    ) -> Result<NoiseStream<C>, NoiseError>
@@ -219,7 +307,7 @@ impl<C> NoiseStreamBuilder<C> {
    {
        while !handshake_state.is_handshake_finished() {
            if handshake_state.is_my_turn() {
-                self.send_handshake_msg(&mut handshake_state, version, pattern)
+                self.send_handshake_msg(&mut handshake_state, &payload, version, pattern)
                    .await?;
            } else {
                self.recv_handshake_msg(&mut handshake_state, version, pattern)
@@ -519,24 +607,26 @@ mod tests {
    }
    #[tokio::test]
-    async fn noise_handshake() -> anyhow::Result<()> {
+    async fn noise_handshake_v1() -> anyhow::Result<()> {
        let dummy_seed = [42u8; 32];
        let mut rng = rand_chacha::ChaCha20Rng::from_seed(dummy_seed);
        let initiator_keys = Arc::new(x25519::KeyPair::new(&mut rng));
        let responder_keys = Arc::new(x25519::KeyPair::new(&mut rng));
        let no_key: Option<[u8; 0]> = None;
        let (initiator_stream, responder_stream) = mock_streams();
        let psk = generate_psk(*responder_keys.public_key(), NoiseVersion::V1)?;
        let pattern = NoisePattern::default();
        let stream_initiator = NoiseStreamBuilder::new(initiator_stream)
            .perform_initiator_handshake_inner(
                pattern,
                initiator_keys.private_key().to_bytes(),
-                responder_keys.public_key().to_bytes(),
+                *responder_keys.public_key(),
-                psk,
+                no_key,
                no_key,
                NoiseVersion::V1,
            );
@@ -545,6 +635,75 @@ mod tests {
                pattern,
                responder_keys.private_key().to_bytes(),
                *responder_keys.public_key(),
                None,
            );
        let initiator_fut =
            tokio::spawn(
                async move { timeout(Duration::from_millis(200), stream_initiator).await },
            );
        let responder_fut =
            tokio::spawn(
                async move { timeout(Duration::from_millis(200), stream_responder).await },
            );
        let (initiator, responder) = join!(initiator_fut, responder_fut);
        let mut initiator = initiator???;
        let mut responder = responder???;
        let msg = b"hello there";
        // if noise was successful we should be able to write a proper message across
        timeout(Duration::from_millis(200), initiator.write_all(msg)).await??;
        initiator.inner_stream.flush().await?;
        let inner_buf = initiator.inner_stream.get_ref().unchecked_tx_data();
        let mut buf = [0u8; 11];
        timeout(Duration::from_millis(200), responder.read(&mut buf)).await??;
        assert_eq!(&buf[..], msg);
        // the inner content is different from the actual msg since it was encrypted
        assert_ne!(inner_buf, buf);
        assert_ne!(inner_buf.len(), msg.len());
        Ok(())
    }
    #[tokio::test]
    async fn noise_handshake_v2() -> anyhow::Result<()> {
        let dummy_seed = [42u8; 32];
        let mut rng = rand_chacha::ChaCha20Rng::from_seed(dummy_seed);
        let initiator_keys = Arc::new(x25519::KeyPair::new(&mut rng));
        let initator_identity_keys = Arc::new(ed25519::KeyPair::new(&mut rng));
        let responder_keys = Arc::new(x25519::KeyPair::new(&mut rng));
        let init_verif_key = initator_identity_keys.public_key().clone();
        let (initiator_stream, responder_stream) = mock_streams();
        let pattern = NoisePattern::default();
        let stream_initiator = NoiseStreamBuilder::new(initiator_stream)
            .perform_initiator_handshake_inner(
                pattern,
                initiator_keys.private_key().to_bytes(),
                responder_keys.public_key().to_bytes(),
                Some(initator_identity_keys.private_key().to_bytes()),
                Some(initator_identity_keys.public_key().to_bytes()),
                NoiseVersion::V2,
            );
        let stream_responder = NoiseStreamBuilder::new(responder_stream)
            .perform_responder_handshake_inner(
                pattern,
                responder_keys.private_key().to_bytes(),
                *responder_keys.public_key(),
                Some(init_verif_key),
            );
        let initiator_fut =
@@ -0,0 +1,30 @@
 [package]
 name = "nympsq"
 version = "0.1.0"
 authors = ["Georgio Nicolas <georgio@nymtech.net>"]
 edition = "2021"
 license.workspace = true
 [dependencies]
 thiserror = { workspace = true }
 # internal
 nym-crypto = { path = "../crypto", features = ["asymmetric", "serde"]}
 libcrux-traits = { git = "https://github.com/cryspen/libcrux" }
 libcrux-kem = { git = "https://github.com/cryspen/libcrux" }
 libcrux-psq = { git = "https://github.com/cryspen/libcrux", features = ["test-utils"] }
 libcrux-ed25519 = { git = "https://github.com/cryspen/libcrux" }
 rand = "0.9.2"
 tls_codec = { version = "0.4.2", features = ["derive"] }
 [dev-dependencies]
 criterion = {workspace = true}
 [[bench]]
 name = "benches"
 harness = false
 [lints]
 workspace = true
@@ -0,0 +1,47 @@
 // Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 use criterion::{criterion_group, criterion_main, Criterion};
 use nym_crypto::asymmetric::ed25519;
 use rand::prelude::*;
 pub fn gen_ed25519_keypair(c: &mut Criterion) {
    c.bench_function("Generate Ed25519 Keypair", |b| {
        b.iter(|| {
            let mut s: [u8; 32] = [0u8; 32];
            rand::rng().fill_bytes(&mut s);
            ed25519::KeyPair::from_secret(s, 0)
        });
    });
 }
 pub fn gen_mlkem768_keypair(c: &mut Criterion) {
    c.bench_function("Generate MlKem768 Keypair", |b| {
        b.iter(|| {
            libcrux_kem::key_gen(libcrux_kem::Algorithm::MlKem768, &mut rand::rng()).unwrap()
        });
    });
 }
 pub fn gen_x25519_keypair(c: &mut Criterion) {
    c.bench_function("Generate DHKem Keypair", |b| {
        b.iter(|| libcrux_kem::key_gen(libcrux_kem::Algorithm::X25519, &mut rand::rng()).unwrap());
    });
 }
 pub fn gen_xwing_keypair(c: &mut Criterion) {
    c.bench_function("Generate XWingKem Keypair", |b| {
        b.iter(|| {
            libcrux_kem::key_gen(libcrux_kem::Algorithm::XWingKemDraft06, &mut rand::rng()).unwrap()
        });
    });
 }
 criterion_group!(
    benches,
    gen_ed25519_keypair,
    gen_mlkem768_keypair,
    gen_x25519_keypair,
    gen_xwing_keypair
 );
 criterion_main!(benches);
@@ -0,0 +1,58 @@
 // Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 use std::io;
 use thiserror::Error;
 #[derive(Error, Debug)]
 pub enum PSQError {
    #[error("encountered a decryption error")]
    DecryptionError,
    #[error("encountered a KEM error")]
    KEMError,
    #[error("encountered a PSQ error")]
    PSQError,
    #[error("encountered a Serialization/Deserialization error")]
    SerializationError,
    #[error("encountered an IO error: {0}")]
    IoError(#[from] io::Error),
    #[error("Incorrect state")]
    IncorrectStateError,
    #[error("Handshake did not complete")]
    HandshakeError,
    #[error("Unknown PSQ version")]
    UnknownVersion,
 }
 impl From<libcrux_kem::Error> for PSQError {
    fn from(err: libcrux_kem::Error) -> Self {
        match err {
            // Error::Decrypt => PSQError::DecryptionError,
            _err => PSQError::KEMError,
        }
    }
 }
 impl From<libcrux_psq::Error> for PSQError {
    fn from(err: libcrux_psq::Error) -> Self {
        match err {
            // Error::Decrypt => PSQError::DecryptionError,
            _err => PSQError::PSQError,
        }
    }
 }
 impl From<tls_codec::Error> for PSQError {
    fn from(err: tls_codec::Error) -> Self {
        match err {
            _err => PSQError::SerializationError,
        }
    }
 }
@@ -0,0 +1,5 @@
 // Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 pub mod error;
 pub mod psq;
@@ -0,0 +1,270 @@
 // Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 use std::{io::Cursor, marker::PhantomData, time::Duration};
 use libcrux_ed25519::VerificationKey;
 use libcrux_psq::{
    cred::Ed25519,
    psk_registration::{Initiator, InitiatorMsg, Responder, ResponderMsg},
    traits::PSQ,
 };
 use tls_codec::{Deserialize, Serialize};
 use libcrux_traits::kem::KEM;
 use rand::{CryptoRng, RngCore};
 use crate::error::PSQError;
 pub const PSK_LENGTH: usize = 32;
 pub const CONTEXT_LEN: usize = 16;
 pub const PSK_HANDLE_LEN: usize = 16;
 pub struct PSQInitiator<T: PSQ> {
    signing_key: [u8; 32],
    verification_key: VerificationKey,
    state: Option<Initiator>,
    _t: PhantomData<T>,
 }
 impl<T: PSQ> PSQInitiator<T> {
    pub fn init(signing_key: impl AsRef<[u8]>, verification_key: impl AsRef<[u8]>) -> Self {
        let mut sig_key: [u8; 32] = [0u8; 32];
        sig_key.clone_from_slice(signing_key.as_ref());
        let mut verif_key: [u8; 32] = [0u8; 32];
        verif_key.clone_from_slice(verification_key.as_ref());
        Self {
            signing_key: sig_key,
            verification_key: VerificationKey::from_bytes(verif_key),
            state: None,
            _t: PhantomData,
        }
    }
    pub fn compute_initiator_message<R>(
        &mut self,
        rng: &mut R,
        responder_kem_public_key: &<T::InnerKEM as KEM>::EncapsulationKey,
        context: &[u8; CONTEXT_LEN],
        psk_ttl: Duration,
    ) -> Result<Vec<u8>, PSQError>
    where
        <T as PSQ>::InnerKEM: KEM,
        R: RngCore + CryptoRng,
    {
        let (state, message) = Initiator::send_initial_message::<Ed25519, T>(
            context,
            psk_ttl,
            responder_kem_public_key,
            &self.signing_key,
            &self.verification_key,
            rng,
        )?;
        self.state = Some(state);
        Ok(message.tls_serialize_detached()?)
    }
    // The initator generates the PSK themselves while computing the first message.
    // It's possible to get_psk() this after sending the first message.
    // The key will only be valid if we're sure that the responder has received it.
    //
    // If we use noise right after PSQ, we can just extract the PSK and expect the responder
    // to just send the first Noise message instead of responding to with a PSQ message.
    // In other words, instead of using the responder's PSQ message for key confirmation,
    // we could use a Noise handshake for key confirmation.
    pub fn get_psk(&self) -> Option<[u8; 32]> {
        match &self.state {
            Some(initiator_state) => Some(*initiator_state.unregistered_psk()),
            None => None,
        }
    }
    pub fn finalize(&self, responder_message: &[u8]) -> Result<[u8; PSK_LENGTH], PSQError> {
        match &self.state {
            Some(state) => match ResponderMsg::tls_deserialize_exact(responder_message) {
                Ok(deserialized_responder_message) => {
                    match state.complete_handshake(&deserialized_responder_message) {
                        Ok(psk) => Ok(psk.psk),
                        Err(err) => Err(err.into()),
                    }
                }
                Err(err) => Err(err.into()),
            },
            None => Err(PSQError::IncorrectStateError),
        }
    }
 }
 pub struct PSQResponder<'a, T: PSQ>
 where
    <T as PSQ>::InnerKEM: KEM,
 {
    kem_private_key: &'a <T::InnerKEM as KEM>::DecapsulationKey,
    kem_public_key: &'a <T::InnerKEM as KEM>::EncapsulationKey,
    _t: PhantomData<T>,
 }
 impl<'a, T: PSQ> PSQResponder<'a, T>
 where
    <T as PSQ>::InnerKEM: KEM,
 {
    pub fn init(
        kem_private_key: &'a <T::InnerKEM as KEM>::DecapsulationKey,
        kem_public_key: &'a <T::InnerKEM as KEM>::EncapsulationKey,
    ) -> Self {
        Self {
            kem_private_key,
            kem_public_key,
            _t: PhantomData,
        }
    }
    pub fn compute_responder_message(
        &self,
        initiator_verification_key: &VerificationKey,
        initiator_message: &mut [u8],
        context: &[u8; CONTEXT_LEN],
        psk_ttl: Duration,
        psk_handle: &[u8; PSK_HANDLE_LEN],
    ) -> Result<([u8; PSK_LENGTH], Vec<u8>), PSQError> {
        let mut cursor = Cursor::new(initiator_message);
        let deserialized_initiator_message = InitiatorMsg::tls_deserialize(&mut cursor)?;
        // the problem is here the deserialized message is the same on the initiator's side
        let (registered_psk, responder_msg) = Responder::send::<Ed25519, T>(
            psk_handle,
            psk_ttl,
            context,
            self.kem_public_key,
            self.kem_private_key,
            initiator_verification_key,
            &deserialized_initiator_message,
        )?;
        let responder_bytes = responder_msg.tls_serialize_detached()?;
        Ok((registered_psk.psk, responder_bytes))
    }
 }
 #[cfg(test)]
 mod test {
    use std::time::Duration;
    use libcrux_ed25519::VerificationKey;
    use libcrux_psq::impls::{MlKem768, XWingKemDraft06, X25519};
    use libcrux_psq::traits::PSQ;
    use libcrux_traits::kem::KEM;
    use nym_crypto::asymmetric::ed25519;
    use rand::prelude::*;
    use crate::psq::{CONTEXT_LEN, PSK_HANDLE_LEN};
    use super::{PSQInitiator, PSQResponder};
    fn test_helper<T: PSQ, R>(
        rng: &mut R,
        responder_kem_private_key: <T::InnerKEM as KEM>::DecapsulationKey,
        responder_kem_public_key: <T::InnerKEM as KEM>::EncapsulationKey,
    ) where
        R: CryptoRng + rand::RngCore,
    {
        // set PSQ TTL
        let psk_ttl: Duration = Duration::from_secs(3600);
        // generate random context string
        let mut context: [u8; CONTEXT_LEN] = [0u8; CONTEXT_LEN];
        rng.fill_bytes(&mut context);
        // generate random psk handle
        let mut psk_handle: [u8; PSK_HANDLE_LEN] = [0u8; PSK_HANDLE_LEN];
        rng.fill_bytes(&mut psk_handle);
        // generate ed25519 keys
        let mut secret_initiator: [u8; 32] = [0u8; 32];
        rng.fill_bytes(&mut secret_initiator);
        let initiator_ed25519_keypair = ed25519::KeyPair::from_secret(secret_initiator, 0);
        let mut initiator: PSQInitiator<T> = PSQInitiator::init(
            initiator_ed25519_keypair.private_key().to_bytes(),
            initiator_ed25519_keypair.public_key().to_bytes(),
        );
        let responder: PSQResponder<T> =
            PSQResponder::init(&responder_kem_private_key, &responder_kem_public_key);
        let mut initiator_msg = initiator
            .compute_initiator_message(rng, &responder_kem_public_key, &context, psk_ttl)
            .unwrap();
        let pre_response_psk = initiator.get_psk().unwrap();
        let initiator_public_key =
            VerificationKey::from_bytes(initiator_ed25519_keypair.public_key().to_bytes());
        let (responder_psk, responder_msg) = responder
            .compute_responder_message(
                &initiator_public_key,
                &mut initiator_msg,
                &context,
                psk_ttl,
                &psk_handle,
            )
            .unwrap();
        let initiator_psk = initiator.finalize(&responder_msg).unwrap();
        assert_eq!(initiator_psk, pre_response_psk);
        assert_eq!(initiator_psk, responder_psk);
    }
    #[test]
    fn test_psq_e2e_mlkem() {
        let mut rng = rand::rng();
        // generate mlkem keypair
        let (responder_kem_private_key, responder_kem_public_key) =
            libcrux_kem::key_gen(libcrux_kem::Algorithm::MlKem768, &mut rng).unwrap();
        test_helper::<MlKem768, _>(
            &mut rng,
            responder_kem_private_key,
            responder_kem_public_key,
        );
    }
    #[test]
    fn test_psq_e2e_xwing() {
        let mut rng = rand::rng();
        // generate xwing keypair
        let (responder_kem_private_key, responder_kem_public_key) =
            libcrux_kem::key_gen(libcrux_kem::Algorithm::XWingKemDraft06, &mut rand::rng())
                .unwrap();
        test_helper::<XWingKemDraft06, _>(
            &mut rng,
            responder_kem_private_key,
            responder_kem_public_key,
        );
    }
    #[test]
    fn test_psq_e2e_dhkem() {
        let mut rng = rand::rng();
        // generate dhkem keypair
        let (responder_kem_private_key, responder_kem_public_key) =
            libcrux_kem::key_gen(libcrux_kem::Algorithm::X25519, &mut rand::rng()).unwrap();
        test_helper::<X25519, _>(
            &mut rng,
            responder_kem_private_key,
            responder_kem_public_key,
        );
    }
 }
@@ -464,16 +464,17 @@ impl ConnectionHandler {
        )
    )]
    pub(crate) async fn handle_connection(&mut self, socket: TcpStream) {
-        let noise_stream = match upgrade_noise_responder(socket, &self.shared.noise_config).await {
+        let noise_stream =
-            Ok(noise_stream) => noise_stream,
+            match upgrade_noise_responder(socket, &self.shared.noise_config, None).await {
-            Err(err) => {
+                Ok(noise_stream) => noise_stream,
-                error!(
+                Err(err) => {
-                    "Failed to perform Noise handshake with {:?} - {err}",
+                    error!(
-                    self.remote_address
+                        "Failed to perform Noise handshake with {:?} - {err}",
-                );
+                        self.remote_address
-                return;
+                    );
-            }
+                    return;
-        };
+                }
            };
        debug!(
            "Noise responder handshake completed for {:?}",
            self.remote_address
Author	SHA1	Message	Date
Georgio Nicolas	5e2135812e	update benchmarking setup	2025-10-29 14:55:12 +01:00
Georgio Nicolas	48304afd3a	Update manifests	2025-10-29 08:08:21 +01:00
Georgio Nicolas	a66d4002a1	Add KKT	2025-10-29 08:08:03 +01:00
Georgio Nicolas	40e07e2a0c	add xwing support	2025-10-07 19:11:47 +02:00
Georgio Nicolas	d5194722cf	Switch back to cryspen repo	2025-08-19 13:56:00 +02:00
Georgio Nicolas	7d8a31339e	temporary support for modified v1 api	2025-08-08 17:23:56 +02:00
Georgio Nicolas	79bf97ceaa	Tests pass	2025-08-08 12:50:11 +02:00
Georgio Nicolas	84fcf81852	update psq implementation	2025-08-08 11:07:35 +02:00
Georgio Nicolas	481ad79376	Integrated but not tested	2025-08-08 11:07:07 +02:00
Georgio Nicolas	1a9995535f	temp changes	2025-07-07 16:44:45 +02:00