Partial fix for celabshq/libcrux#1275

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

Author: gram-signal

Cargo.lock

@@ -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" }

Cargo.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"]

Cross.toml

@@ -69,13 +69,87 @@ 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"),
     );
     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/incremental_mlkem768.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(())
+    }
 }