Partial fix for cryspen/libcrux#1275

Co-authored-by: Rolfe Schmidt <rolfe@signal.org>

Cargo.lock

@@ -348,6 +348,12 @@ version = "0.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
 
+[[package]]
+name = "hex"
+version = "0.4.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70"
+
 [[package]]
 name = "hkdf"
 version = "0.12.4"
@@ -469,9 +475,9 @@ dependencies = [
 
 [[package]]
 name = "libcrux-ml-kem"
-version = "0.0.4"
+version = "0.0.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4bb6a88086bf11bd2ec90926c749c4a427f2e59841437dbdede8cde8a96334ab"
+checksum = "22a36f21056e552438dbe6ce413b6682001795e53bd1f0e2c941d7e231238e5a"
 dependencies = [
  "hax-lib",
  "libcrux-intrinsics",
@@ -483,9 +489,9 @@ dependencies = [
 
 [[package]]
 name = "libcrux-platform"
-version = "0.0.2"
+version = "0.0.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "db82d058aa76ea315a3b2092f69dfbd67ddb0e462038a206e1dcd73f058c0778"
+checksum = "1d9e21d7ed31a92ac539bd69a8c970b183ee883872d2d19ce27036e24cb8ecc4"
 dependencies = [
  "libc",
 ]
@@ -512,9 +518,9 @@ dependencies = [
 
 [[package]]
 name = "libcrux-sha3"
-version = "0.0.4"
+version = "0.0.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2400bec764d1c75b8a496d5747cffe32f1fb864a12577f0aca2f55a92021c962"
+checksum = "18e869fdeb9af62c55c3fcd60ce407552eb282d727550ce986abac94a3474479"
 dependencies = [
  "hax-lib",
  "libcrux-intrinsics",
@@ -1002,7 +1008,7 @@ checksum = "d372029cb5195f9ab4e4b9aef550787dce78b124fcaee8d82519925defcd6f0d"
 
 [[package]]
 name = "spqr"
-version = "1.3.0"
+version = "1.4.0"
 dependencies = [
  "cpufeatures",
  "curve25519-dalek",
@@ -1010,6 +1016,7 @@ dependencies = [
  "env_logger",
  "galois_field_2pm",
  "hax-lib",
+ "hex",
  "hkdf",
  "hmac",
  "libcrux-hmac",

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "spqr"
-version = "1.3.0"
+version = "1.4.0"
 edition = "2021"
 license = "AGPL-3.0-only"
 rust-version = "1.83.0"
@@ -13,7 +13,7 @@ displaydoc = "0.2"
 hax-lib = "0.3.5"
 hkdf = "0.12"
 libcrux-hmac = "0.0.4"
-libcrux-ml-kem = { version = "0.0.4", default-features = false, features = ["incremental", "mlkem768"] }
+libcrux-ml-kem = { version = "0.0.5", default-features = false, features = ["incremental", "mlkem768"] }
 log = "0.4.21"
 num_enum = "0.7.3"
 prost = "0.14.1"
@@ -29,6 +29,7 @@ spqr = { path = ".", features = ["test-utils"] }
 
 env_logger = "0.11"
 galois_field_2pm = "0.1.0"
+hex = "0.4"
 hmac = "0.12.1"
 matches = "0.1.10"
 rand_08 = { package = "rand", version = "0.8" }

Cross.toml

@@ -5,3 +5,9 @@ pre-build = [
   "unzip protoc-29.3-linux-x86_64.zip -d protoc",
   "mv protoc/bin/protoc /usr/bin",
   ]
+
+[build.env]
+passthrough = [
+  "RUST_LOG",
+  ]
+volumes = ["/tmp:/tmp"]

src/incremental_mlkem768.rs

@@ -69,6 +69,8 @@ pub fn encaps1<R: Rng + CryptoRng>(
 #[hax_lib::requires(es.len() == 2080 && ek.len() == 1152)]
 #[hax_lib::ensures(|result| result.len() == 128)]
 pub fn encaps2(ek: &EncapsulationKey, es: &EncapsulationState) -> Ciphertext2 {
+    let maybe_fix = potentially_fix_state_incorrectly_encoded_by_libcrux_issue_1275(es);
+    let es = maybe_fix.as_ref().unwrap_or(es);
     let ct2 = incremental::encapsulate2(
         es.as_slice().try_into().expect("size should be correct"),
         ek.as_slice().try_into().expect("size should be correct"),
@@ -76,6 +78,78 @@ pub fn encaps2(ek: &EncapsulationKey, es: &EncapsulationState) -> Ciphertext2 {
     ct2.value.to_vec()
 }
 
+/// Due to https://github.com/cryspen/libcrux/issues/1275, state may
+/// contain incorrect endian-ness.  We need to fix this locally before
+/// using it.  Luckily, this is doable by checking that the values in
+/// error2 are in the range [-2, 2].
+#[hax_lib::requires(es.len() == 2080)]
+#[hax_lib::ensures(|result| if let Some(es) = result {
+    es.len() == 2080
+} else {
+    true
+})]
+#[hax_lib::opaque]
+fn potentially_fix_state_incorrectly_encoded_by_libcrux_issue_1275(
+    es: &EncapsulationState,
+) -> Option<EncapsulationState> {
+    assert_eq!(es.len(), 2080);
+    // Look at each value within the error2 portion of EncapsState.
+    //
+    // The last 32 bytes are a raw random vector, thus they don't have endianness
+    // and we shouldn't flip them.  All other bytes are encoded i16s.
+    // This is the libcrux-specific encoding of the following struct, where all
+    // PolynomialRingElements are encoded as described.
+    //
+    //    pub struct EncapsState<const K: usize, Vector: Operations> {
+    //        pub(super) r_as_ntt: [PolynomialRingElement<Vector>; K],
+    //        pub(super) error2: PolynomialRingElement<Vector>,
+    //        pub(super) randomness: [u8; 32],
+    //    }
+    //
+    // To determine whether it's valid or not, we look at the encoding of `error2`.
+    // All error2 values should be in the range [-2, 2].
+    // This is 𝜂2 from https://nvlpubs.nist.gov/nistpubs/fips/nist.fips.203.pdf
+    // page 39 table 2, generated as 𝑒2 on page 30 algorithm 14 line 17.
+    const NEG1_I16: i16 = 0xFFFFu16 as i16;
+    const NEG2_I16_GOOD: i16 = 0xFFFEu16 as i16;
+    const NEG2_I16_BAD: i16 = 0xFEFFu16 as i16;
+
+    for c in es[1536..2080 - 32].chunks(2) {
+        match i16::from_le_bytes(c.try_into().expect("chunk should be size 2")) {
+            0x0000i16 | NEG1_I16 => {} // These have the same encoding for i16 in little/big-endian
+            0x0001i16 | 0x0002i16 | NEG2_I16_GOOD => {
+                return None;
+            }
+            0x0100i16 | 0x0200i16 | NEG2_I16_BAD => {
+                // If they aren't in the expected range, we probably have a state with bad endian-ness.  So flip it.
+                #[cfg(not(hax))]
+                log::info!("spqr fixing bad encapsulate1 stored state endianness");
+                return Some(flip_endianness_of_encapsulation_state(es));
+            }
+            _ => {
+                // We're in a weird state, so just use what we had initially.
+                #[cfg(not(hax))]
+                log::warn!("spqr unable to fix encapsulate1 stored state endianness");
+                return None;
+            }
+        }
+    }
+    None
+}
+
+#[hax_lib::requires(es.len()%2 == 0 && es.len() == 2080)]
+#[hax_lib::ensures(|result| result.len() == es.len())]
+#[hax_lib::opaque]
+fn flip_endianness_of_encapsulation_state(es: &EncapsulationState) -> EncapsulationState {
+    assert!(es.len() % 2 == 0);
+    assert!(es.len() > 32);
+    let mut fixed_es = es.clone();
+    for i in (0..fixed_es.len() - 32).step_by(2) {
+        (fixed_es[i], fixed_es[i + 1]) = (fixed_es[i + 1], fixed_es[i])
+    }
+    fixed_es
+}
+
 /// Decapsulate ciphertext to get shared secret.
 #[hax_lib::requires(ct1.len() == 960 && ct2.len() == 128 && dk.len() == 2400)]
 #[hax_lib::ensures(|result| result.len() == 32)]

src/lib.rs

@@ -1066,4 +1066,108 @@ mod lib_test {
         ));
         Ok(())
     }
+
+    /// A test that runs a set number of steps, always sending one message A->B, then
+    /// one message B->A, with logging turned on, so we can watch how things work
+    /// in the most predictable case.
+    #[test]
+    fn lockstep_run_with_logging() -> Result<(), Error> {
+        let _ = env_logger::builder().is_test(true).try_init();
+        let mut rng = OsRng.unwrap_err();
+
+        let version = Version::V1;
+
+        let mut alex_pq_state = initial_state(Params {
+            version,
+            min_version: version,
+            direction: Direction::A2B,
+            auth_key: &[41u8; 32],
+            chain_params: ChainParams::default(),
+        })?;
+        let mut blake_pq_state = initial_state(Params {
+            version,
+            min_version: version,
+            direction: Direction::B2A,
+            auth_key: &[41u8; 32],
+            chain_params: ChainParams::default(),
+        })?;
+
+        for i in 0..30 {
+            log::info!("step {}", i);
+            // Now let's send some messages
+            let Send {
+                state,
+                msg,
+                key: alex_key,
+            } = send(&alex_pq_state, &mut rng)?;
+            alex_pq_state = state;
+            let Recv {
+                state,
+                key: blake_key,
+            } = recv(&blake_pq_state, &msg)?;
+            blake_pq_state = state;
+            assert_eq!(alex_key, blake_key);
+            let Send {
+                state,
+                msg,
+                key: blake_key,
+            } = send(&blake_pq_state, &mut rng)?;
+            blake_pq_state = state;
+            let Recv {
+                state,
+                key: alex_key,
+            } = recv(&alex_pq_state, &msg)?;
+            alex_pq_state = state;
+            assert_eq!(alex_key, blake_key);
+        }
+        log::info!("alex_state:  {}", hex::encode(alex_pq_state));
+        log::info!("blake_state: {}", hex::encode(blake_pq_state));
+        Ok(())
+    }
+
+    #[test]
+    fn regression_test_libcrux_issue_1275_from_generated_states() -> Result<(), Error> {
+        let _ = env_logger::builder().is_test(true).try_init();
+        let mut rng = OsRng.unwrap_err();
+
+        // These states are generated using the old "portable" logic for libcrux-ml-kem
+        // EncapsState serialization prior to libcrux:pr/1276, 30 steps into a lockstep
+        // protocol.  The send_ct side has already generated the encapsulation state
+        // and stored it locally, but hasn't yet called encapsulate2 on it.  This tests
+        // to make sure that the incremental_mlkem code path correctly notices and handles
+        // this eventuality.
+        let mut alex_pq_state = include_bytes!("issue1275_a_state.in").to_vec();
+        let mut blake_pq_state = include_bytes!("issue1275_b_state.in").to_vec();
+
+        // After 20 additional steps, we should be in epoch 2 successfully.  If
+        // we're unable to handle the bad state, one of these steps will fail.
+        for i in 30..50 {
+            log::info!("step {}", i);
+            let Send {
+                state,
+                msg,
+                key: alex_key,
+            } = send(&alex_pq_state, &mut rng)?;
+            alex_pq_state = state;
+            let Recv {
+                state,
+                key: blake_key,
+            } = recv(&blake_pq_state, &msg)?;
+            blake_pq_state = state;
+            assert_eq!(alex_key, blake_key);
+            let Send {
+                state,
+                msg,
+                key: blake_key,
+            } = send(&blake_pq_state, &mut rng)?;
+            blake_pq_state = state;
+            let Recv {
+                state,
+                key: alex_key,
+            } = recv(&alex_pq_state, &msg)?;
+            alex_pq_state = state;
+            assert_eq!(alex_key, blake_key);
+        }
+        Ok(())
+    }
 }