initial implementation using proper keys

Author: dartdart26

Cargo.lock

@@ -32,6 +32,8 @@ futures = "0.3.32"
 hybrid-array = { version = "0.4.7", features = ["bytemuck"] }
 libc = "0.2.181"
 ml-kem = "0.2.2"
+module-lattice = "0.1.0"
+sha3 = "0.10.8"
 ndarray = "0.17.2"
 num-traits = "0.2.19"
 rand = "0.10.0"

Cargo.toml

@@ -29,8 +29,11 @@ cryprot-net.workspace = true
 cryprot-pprf.workspace = true
 curve25519-dalek = { workspace = true, features = ["rand_core", "serde"] }
 futures.workspace = true
+hybrid-array.workspace = true
 ml-kem.workspace = true
+module-lattice.workspace = true
 rand.workspace = true
+sha3.workspace = true
 serde_bytes.workspace = true
 serde = { workspace = true, features = ["derive"] }
 subtle.workspace = true

cryprot-ot/Cargo.toml

@@ -178,7 +178,8 @@ ek_b = r_b + H(r_{1-b})
      = (ek - H(r_{1-b})) + H(r_{1-b})
      = ek
 ```
-So `ek_b = ek`, the real public key. In step 10, the receiver calls `ML-KEM.Decaps(dk, ct_b)` and recovers the same shared secret `ss_b` that the sender computed via `ML-KEM.Encaps(ek_b)` in step 7.
+So `ek_b = ek`, the real public key. In step 10, the receiver calls `ML-KEM.Decaps(dk, ct_b)` and
+recovers the same shared secret `ss_b` that the sender computed via `ML-KEM.Encaps(ek_b)` in step 7.
 
 **Security:**
 

cryprot-ot/docs/mlkem-ot-protocol.md

@@ -1,30 +1,165 @@
 //! Post-quantum base OT using ML-KEM.
+//!
+//! Implements the MR19 protocol (Masny-Rindal, ePrint 2019/706, Figure 8)
+//! instantiated with ML-KEM as per Section D.3.
+//! See `docs/mlkem-ot-protocol.md` for the full protocol description.
 
-use std::io;
+use std::{io, mem::size_of};
 
 use cryprot_core::{Block, buf::Buf, rand_compat::RngCompat, random_oracle::RandomOracle};
 use cryprot_net::{Connection, ConnectionError};
 use futures::{SinkExt, StreamExt};
+use hybrid_array::typenum::Unsigned;
 // ML-KEM variant: change to MlKem512/MlKem512Params or MlKem768/MlKem768Params
 // for different security levels.
 use ml_kem::{
     Ciphertext as MlKemCiphertext, EncodedSizeUser, KemCore, MlKem1024 as MlKem,
-    MlKem1024Params as MlKemParams, SharedKey,
-    array::typenum::Unsigned,
+    MlKem1024Params as MlKemParams, ParameterSet, SharedKey,
     kem::{Decapsulate, DecapsulationKey, Encapsulate, EncapsulationKey as MlKemEncapsulationKey},
 };
+use module_lattice::{Encode, Field, NttPolynomial};
 use rand::{RngExt, rngs::StdRng};
 use serde::{Deserialize, Serialize};
+use sha3::{
+    Shake128,
+    digest::{ExtendableOutput, Update, XofReader},
+};
 use subtle::{Choice, ConditionallySelectable};
 use tracing::Level;
 
 use crate::{Connected, RotReceiver, RotSender, SemiHonest, phase};
 
+// Define the ML-KEM base field (q = 3329).
+module_lattice::define_field!(MlKemField, u16, u32, u64, 3329);
+
+// Module dimension derived from the chosen ML-KEM parameter set.
+type K = <MlKemParams as ParameterSet>::K;
+
+type NttVector = module_lattice::NttVector<MlKemField, K>;
+
+type U12 = hybrid_array::typenum::U12;
+
 const ENCAPSULATION_KEY_LEN: usize =
     <MlKemEncapsulationKey<MlKemParams> as EncodedSizeUser>::EncodedSize::USIZE;
 const CIPHERTEXT_LEN: usize = <MlKem as KemCore>::CiphertextSize::USIZE;
 const HASH_DOMAIN_SEPARATOR: &[u8] = b"MlKemOt";
 
+// Number of coefficients per polynomial (FIPS 203, Section 2: n = 256).
+const NUM_COEFFICIENTS: usize = 256;
+
+/// rho is a 32-byte seed used to derive the public matrix A_hat (FIPS 203).
+type Rho = [u8; 32];
+
+// Serialized t_hat is the encapsulation key minus the rho suffix.
+const T_HAT_BYTES_LEN: usize = ENCAPSULATION_KEY_LEN - size_of::<Rho>();
+
+// ---------------------------------------------------------------------------
+// Protocol helper functions (see docs/mlkem-ot-protocol.md)
+// ---------------------------------------------------------------------------
+
+/// Parse serialized encapsulation key bytes into (NttVector, rho).
+/// The input is a fixed-size array so the slicing is infallible.
+fn parse_ek(bytes: &[u8; ENCAPSULATION_KEY_LEN]) -> (NttVector, Rho) {
+    let enc = bytes[..T_HAT_BYTES_LEN]
+        .try_into()
+        .expect("t_hat length mismatch");
+    let t_hat = <NttVector as Encode<U12>>::decode(enc);
+    let rho = bytes[T_HAT_BYTES_LEN..]
+        .try_into()
+        .expect("rho length mismatch");
+    (t_hat, rho)
+}
+
+/// Serialize NttVector + rho back into encapsulation key bytes.
+fn serialize_ek(t_hat: &NttVector, rho: &Rho) -> [u8; ENCAPSULATION_KEY_LEN] {
+    let encoded = <NttVector as Encode<U12>>::encode(t_hat);
+    let mut out = [0u8; ENCAPSULATION_KEY_LEN];
+    out[..T_HAT_BYTES_LEN].copy_from_slice(encoded.as_slice());
+    out[T_HAT_BYTES_LEN..].copy_from_slice(rho);
+    out
+}
+
+/// XOF(rho, j, i) from FIPS 203, Algorithm 2 SHAKE128example.
+/// In Algorithm 13 (K-PKE.KeyGen), this is called as XOF(rho, j, i) where
+/// j is the column index (byte 32) and i is the row index (byte 33),
+/// using 0-based indexing.
+fn xof(seed: &Rho, j: u8, i: u8) -> impl XofReader {
+    let mut h = Shake128::default();
+    h.update(seed);
+    h.update(&[i, j]);
+    h.finalize_xof()
+}
+
+/// FIPS 203 Algorithm 7: SampleNTT.
+/// Rejection sampling from a byte stream to produce a pseudorandom NTT
+/// polynomial.
+///
+/// Adapted from the ml-kem crate's `sample_ntt`.
+fn sample_ntt_poly(xof: &mut impl XofReader) -> NttPolynomial<MlKemField> {
+    const Q: u16 = MlKemField::Q;
+    // Read 32 triples (3 bytes each) at a time from the XOF.
+    const BUF_LEN: usize = 96;
+    let mut poly = NttPolynomial::<MlKemField>::default();
+    let mut buf = [0u8; BUF_LEN];
+    let mut pos = BUF_LEN; // start at end to trigger first read
+    let mut i = 0;
+
+    while i < NUM_COEFFICIENTS {
+        if pos >= BUF_LEN {
+            xof.read(&mut buf);
+            pos = 0;
+        }
+
+        let d1 = u16::from(buf[pos]) | ((u16::from(buf[pos + 1]) & 0x0F) << 8);
+        let d2 = (u16::from(buf[pos + 1]) >> 4) | (u16::from(buf[pos + 2]) << 4);
+        pos += 3;
+
+        if d1 < Q {
+            poly.0[i] = module_lattice::Elem::new(d1);
+            i += 1;
+        }
+        if i < NUM_COEFFICIENTS && d2 < Q {
+            poly.0[i] = module_lattice::Elem::new(d2);
+            i += 1;
+        }
+    }
+
+    poly
+}
+
+/// SampleNTTVector: call SampleNTT k times with FIPS 203 domain separation.
+/// Produces a pseudorandom NttVector<k> from a 32-byte seed.
+/// Each polynomial j uses XOF(seed || j || 0).
+fn sample_ntt_vector(seed: &Rho) -> NttVector {
+    NttVector::new(
+        (0..K::USIZE)
+            .map(|j| {
+                let mut reader = xof(seed, j as u8, 0);
+                sample_ntt_poly(&mut reader)
+            })
+            .collect(),
+    )
+}
+
+/// H(ek): hash-to-key. Maps an NttVector to another NttVector via SHA3-256.
+/// Corresponds to libOTe's `pkHash`.
+fn hash_to_key(t_hat: &NttVector) -> NttVector {
+    use sha3::Digest;
+    let encoded = <NttVector as Encode<U12>>::encode(t_hat);
+    let seed: Rho = sha3::Sha3_256::digest(encoded.as_slice()).into();
+    sample_ntt_vector(&seed)
+}
+
+/// RandomEK: generate a random NttVector from a random seed.
+fn random_ek(rng: &mut StdRng) -> NttVector {
+    let seed: Rho = rng.random();
+    sample_ntt_vector(&seed)
+}
+
+// ---------------------------------------------------------------------------
+// Wire types and protocol implementation
+// ---------------------------------------------------------------------------
+
 #[derive(thiserror::Error, Debug)]
 pub enum Error {
     #[error("quic connection error")]
@@ -67,8 +202,7 @@ impl ConditionallySelectable for CtBytes {
     }
 }
 
-// Message from receiver to sender: two encapsulation keys per OT.
-// For choice bit c, ek{c} is a real key, ek{1-c} is random bytes.
+// Message from receiver to sender: two values (r_0, r_1) per OT.
 #[derive(Serialize, Deserialize)]
 struct EncapsulationKeysMessage {
     eks0: Vec<EncapKeyBytes>,
@@ -133,16 +267,23 @@ impl RotSender for MlKemOt {
 
         let mut cts0 = Vec::with_capacity(count);
         let mut cts1 = Vec::with_capacity(count);
-        for (i, (ek0, ek1)) in receiver_msg
+        for (i, (r0_bytes, r1_bytes)) in receiver_msg
             .eks0
             .iter()
             .zip(receiver_msg.eks1.iter())
             .enumerate()
         {
-            let (ct0, key0) = encapsulate(ek0, &mut self.rng);
+            // Reconstruct encapsulation keys: ek_j = r_j + H(r_{1-j})
+            let (r0, rho) = parse_ek(&r0_bytes.0);
+            let (r1, _) = parse_ek(&r1_bytes.0);
+
+            let ek0_bytes = serialize_ek(&(&r0 + &hash_to_key(&r1)), &rho);
+            let ek1_bytes = serialize_ek(&(&r1 + &hash_to_key(&r0)), &rho);
+
+            let (ct0, key0) = encapsulate(&EncapKeyBytes(ek0_bytes), &mut self.rng);
             let key0 = hash(&key0, i);
 
-            let (ct1, key1) = encapsulate(ek1, &mut self.rng);
+            let (ct1, key1) = encapsulate(&EncapKeyBytes(ek1_bytes), &mut self.rng);
             let key1 = hash(&key1, i);
 
             cts0.push(ct0);
@@ -180,18 +321,30 @@ impl RotReceiver for MlKemOt {
         let mut eks1 = Vec::with_capacity(count);
 
         for choice in choices.iter() {
-            // Generate real keypair.
+            // Step 1: Generate real keypair.
             let (dk, ek) = MlKem::generate(&mut RngCompat(&mut self.rng));
-            let real_ek = EncapKeyBytes(
-                ek.as_bytes()
-                    .as_slice()
-                    .try_into()
-                    .expect("incorrect encapsulation key size"),
-            );
-            let fake_ek = EncapKeyBytes(self.rng.random());
+            let ek_bytes: [u8; ENCAPSULATION_KEY_LEN] = ek
+                .as_bytes()
+                .as_slice()
+                .try_into()
+                .expect("incorrect encapsulation key size");
+            let (real_t_hat, rho) = parse_ek(&ek_bytes);
+
+            // Step 2: Sample random key for position 1-b.
+            let rand_t_hat = random_ek(&mut self.rng);
+
+            // Step 3: Compute correlated key for position b: r_b = ek - H(r_{1-b}).
+            let correlated_t_hat = &real_t_hat - &hash_to_key(&rand_t_hat);
+
+            // Serialize both keys with the same rho.
+            let correlated_bytes = EncapKeyBytes(serialize_ek(&correlated_t_hat, &rho));
+            let random_bytes = EncapKeyBytes(serialize_ek(&rand_t_hat, &rho));
 
-            let ek0 = EncapKeyBytes::conditional_select(&real_ek, &fake_ek, *choice);
-            let ek1 = EncapKeyBytes::conditional_select(&fake_ek, &real_ek, *choice);
+            // Step 4: Select (r_0, r_1) based on choice bit (constant-time).
+            // If b=0: r_0 = correlated (real side), r_1 = random.
+            // If b=1: r_0 = random, r_1 = correlated (real side).
+            let ek0 = EncapKeyBytes::conditional_select(&correlated_bytes, &random_bytes, *choice);
+            let ek1 = EncapKeyBytes::conditional_select(&random_bytes, &correlated_bytes, *choice);
 
             decap_keys.push(dk);
             eks0.push(ek0);
@@ -217,15 +370,18 @@ impl RotReceiver for MlKemOt {
             });
         }
 
-        // Decapsulate the chosen ciphertext for each OT.
+        // Step 10-11: Decapsulate the chosen ciphertext and derive OT key.
         for (i, ((dk, choice), (ct0, ct1))) in decap_keys
             .iter()
             .zip(choices.iter())
             .zip(sender_msg.cts0.iter().zip(sender_msg.cts1.iter()))
             .enumerate()
         {
-            let chosen_ct: MlKemCiphertext<MlKem> =
-                CtBytes::conditional_select(ct0, ct1, *choice).0.into();
+            let ct_bytes = CtBytes::conditional_select(ct0, ct1, *choice).0;
+            let chosen_ct: MlKemCiphertext<MlKem> = ct_bytes
+                .as_slice()
+                .try_into()
+                .expect("incorrect ciphertext size");
             let shared_key = dk
                 .decapsulate(&chosen_ct)
                 .map_err(|_| Error::Decapsulation)?;