module-lattice: harden reduction; ml-kem: expand cargo checkct

.github/workflows/ml-kem.yml

@@ -87,6 +87,7 @@ jobs:
         with:
           directory: ml-kem
           timeout: 120
+          jobs: 1
 
   cross:
     needs: set-msrv

ml-kem/checkct/.cargo/config.toml

@@ -1,10 +1,10 @@
 [build]
 target = [
-    # TODO(tarcieri): investigate codegen problems on commented out platforms
-    #"riscv32imac-unknown-none-elf",
-    #"thumbv6m-none-eabi",
+    # TODO(tarcieri): full encapsulation exceeds analysis budget on `riscv32`
+    # "riscv32imac-unknown-none-elf",
+    "thumbv6m-none-eabi",
     "thumbv7em-none-eabihf",
-    #"x86_64-unknown-linux-gnu",
+    "x86_64-unknown-linux-gnu",
 ]
 
 [target.x86_64-unknown-linux-gnu]

ml-kem/checkct/driver/Cargo.toml

@@ -6,6 +6,7 @@ edition = "2024"
 
 [dependencies]
 checkct_macros = { path = "../checkct_macros" }
+module-lattice = { path = "../../../module-lattice", default-features = false }
 ml-kem = { path = "../..", default-features = false }
 # PrivateRng and PublicRng implement the traits of all these rand_core versions,
 # so that they can be passed to libraries depending on any of them.

ml-kem/checkct/driver/src/driver.rs

@@ -1,14 +1,42 @@
 use crate::rng::{PrivateRng, PublicRng};
 use checkct_macros::checkct;
 use core::hint::black_box;
-use ml_kem::{B32, EncapsulationKey, MlKem512};
+use ml_kem::{B32, EncapsulationKey};
+
+macro_rules! check_encapsulation {
+    ($name:ident, $kem:ident, $key_len:expr) => {
+        #[checkct]
+        pub fn $name() {
+            // The generated BINSEC script references both RNG hooks.
+            black_box(PublicRng.next_u32());
+
+            // One accepted public key per parameter set. Keeping rho concrete
+            // avoids symbolic SHAKE/rejection sampling during public-key expansion.
+            let ek = EncapsulationKey::<ml_kem::$kem>::new(&[0u8; $key_len].into()).unwrap();
+            let mut m = B32::default();
+            PrivateRng.fill_bytes(&mut m);
+            black_box(ek.encapsulate_deterministic(&m));
+        }
+    };
+}
+
+check_encapsulation!(encapsulate_deterministic, MlKem512, 800);
+check_encapsulation!(encapsulate_768, MlKem768, 1184);
+check_encapsulation!(encapsulate_1024, MlKem1024, 1568);
+
+module_lattice::define_field!(KyberField, u16, u32, u64, 3329);
+module_lattice::define_field!(DilithiumField, u32, u64, u128, 8_380_417);
+
+#[checkct]
+pub fn small_reduce_u16() {
+    use module_lattice::Field;
+    black_box(PublicRng.next_u32());
+    black_box(KyberField::small_reduce(PrivateRng.next_u32() as u16));
+}
 
 #[checkct]
-pub fn encapsulate_deterministic() {
-    // The generated BINSEC script references both RNG hooks unconditionally.
+pub fn small_reduce_u32() {
+    use module_lattice::Field;
     black_box(PublicRng.next_u32());
-    let ek = EncapsulationKey::<MlKem512>::new(&[0u8; 800].into()).unwrap();
-    let mut m = B32::default();
-    PrivateRng.fill_bytes(&mut m);
-    black_box(ek.encapsulate_deterministic(&m));
+    black_box(DilithiumField::small_reduce(PrivateRng.next_u32()));
 }

module-lattice/Cargo.toml

@@ -16,19 +16,19 @@ categories = ["cryptography", "no-std"]
 keywords = ["crypto", "kyber", "lattice", "post-quantum"]
 
 [dependencies]
-array = { version = "0.4.8", package = "hybrid-array", features = ["extra-sizes"] }
+array = { version = "0.4.8", package = "hybrid-array", features = ["ctutils", "extra-sizes"] }
+ctutils = "0.4.7"
 num-traits = { version = "0.2", default-features = false }
 
 # optional dependencies
-ctutils = { version = "0.4.5", optional = true }
 zeroize = { version = "1.8.1", optional = true, default-features = false }
 
 [dev-dependencies]
 getrandom = { version = "0.4", features = ["sys_rng"] }
 
 [features]
 alloc = []
-ctutils = ["dep:ctutils", "array/ctutils"]
+ctutils = [] # TODO(tarcieri): remove in next breaking release
 zeroize = ["array/zeroize", "dep:zeroize"]
 
 [lints]

module-lattice/src/algebra.rs

@@ -72,13 +72,8 @@ macro_rules! define_field {
             const BARRETT_MULTIPLIER: Self::LongLong = (1 << Self::BARRETT_SHIFT) / Self::QLL;
 
             fn small_reduce(x: Self::Int) -> Self::Int {
-                // Branchless conditional subtraction: if x >= Q, subtract Q; else
-                // leave x alone. Compilers already emit `csel` here at O2, but the
-                // explicit mask form removes the dependency on optimizer choices
-                // and keeps the generated assembly free of secret-dependent control
-                // flow at every optimization level.
-                let mask = ((x >= Self::Q) as $int).wrapping_neg();
-                x - (Self::Q & mask)
+                use $crate::ctutils::{CtLt, CtSelect};
+                x - Self::Q.ct_select(&0, x.ct_lt(&Self::Q))
             }
 
             fn barrett_reduce(x: Self::Long) -> Self::Int {

module-lattice/src/lib.rs

@@ -32,5 +32,4 @@ pub use encoding::{
 pub use maybe_box::MaybeBox;
 pub use truncate::Truncate;
 
-#[cfg(feature = "ctutils")]
 pub use ctutils;

module-lattice/tests/algebra.rs

@@ -14,9 +14,19 @@ module_lattice::define_field!(DilithiumField, u32, u64, u128, 8_380_417);
 
 #[test]
 fn small_reduce() {
+    for x in 0..=u16::MAX {
+        let expected = if x >= 3329 { x - 3329 } else { x };
+        assert_eq!(KyberField::small_reduce(x), expected);
+    }
+
     assert_eq!(KyberField::small_reduce(3328), 3328);
     assert_eq!(KyberField::small_reduce(3329), 0);
 
+    for x in [0, 1, 8_380_416, 8_380_417, 16_760_833, u32::MAX] {
+        let expected = if x >= 8_380_417 { x - 8_380_417 } else { x };
+        assert_eq!(DilithiumField::small_reduce(x), expected);
+    }
+
     assert_eq!(DilithiumField::small_reduce(8_380_416), 8_380_416);
     assert_eq!(DilithiumField::small_reduce(8_380_417), 0);
 }