ml-kem: fix potential kyberslash attack (#18)

Updates compress method to remove division of secret values by public value (q)
as described in the kyberslash attack

Author: supinie

.github/workflows/workspace.yml

@@ -21,7 +21,7 @@ jobs:
     timeout-minutes: 45
     steps:
       - uses: actions/checkout@v4
-      - uses: dtolnay/rust-toolchain@1.74.0
+      - uses: dtolnay/rust-toolchain@beta # TODO: use `1.79` after 2024-06-13
         with:
           components: clippy
       - run: cargo clippy -- -D warnings

Cargo.lock

@@ -28,6 +28,7 @@ sha3 = { version = "0.10.8", default-features = false }
 criterion = "0.5.1"
 hex = "0.4.3"
 hex-literal = "0.4.1"
+num-rational = { version = "0.4.2", default-features = false, features = ["num-bigint"] }
 rand = "0.8.5"
 crypto-common = { version = "0.1.6", features = ["rand_core"] }
 

ml-kem/Cargo.toml

@@ -17,8 +17,9 @@ pub struct FieldElement(pub Integer);
 impl FieldElement {
     pub const Q: Integer = 3329;
     pub const Q32: u32 = Self::Q as u32;
-    const Q64: u64 = Self::Q as u64;
+    pub const Q64: u64 = Self::Q as u64;
     const BARRETT_SHIFT: usize = 24;
+    #[allow(clippy::integer_division_remainder_used)]
     const BARRETT_MULTIPLIER: u64 = (1 << Self::BARRETT_SHIFT) / Self::Q64;
 
     // A fast modular reduction for small numbers `x < 2*q`
@@ -263,6 +264,7 @@ impl NttPolynomial {
 #[allow(clippy::cast_possible_truncation)]
 const ZETA_POW_BITREV: [FieldElement; 128] = {
     const ZETA: u64 = 17;
+    #[allow(clippy::integer_division_remainder_used)]
     const fn bitrev7(x: usize) -> usize {
         ((x >> 6) % 2)
             | (((x >> 5) % 2) << 1)
@@ -277,6 +279,7 @@ const ZETA_POW_BITREV: [FieldElement; 128] = {
     let mut pow = [FieldElement(0); 128];
     let mut i = 0;
     let mut curr = 1u64;
+    #[allow(clippy::integer_division_remainder_used)]
     while i < 128 {
         pow[i] = FieldElement(curr as u16);
         i += 1;
@@ -300,6 +303,7 @@ const GAMMA: [FieldElement; 128] = {
     let mut i = 0;
     while i < 128 {
         let zpr = ZETA_POW_BITREV[i].0 as u64;
+        #[allow(clippy::integer_division_remainder_used)]
         let g = (zpr * zpr * ZETA) % FieldElement::Q64;
         gamma[i] = FieldElement(g as u16);
         i += 1;

ml-kem/src/algebra.rs

@@ -6,6 +6,8 @@ use crate::util::Truncate;
 pub trait CompressionFactor: EncodingSize {
     const POW2_HALF: u32;
     const MASK: Integer;
+    const DIV_SHIFT: usize;
+    const DIV_MUL: u64;
 }
 
 impl<T> CompressionFactor for T
@@ -14,6 +16,9 @@ where
 {
     const POW2_HALF: u32 = 1 << (T::USIZE - 1);
     const MASK: Integer = ((1 as Integer) << T::USIZE) - 1;
+    const DIV_SHIFT: usize = 34;
+    #[allow(clippy::integer_division_remainder_used)]
+    const DIV_MUL: u64 = (1 << T::DIV_SHIFT) / FieldElement::Q64;
 }
 
 // Traits for objects that allow compression / decompression
@@ -25,25 +30,26 @@ pub trait Compress {
 impl Compress for FieldElement {
     // Equation 4.5: Compress_d(x) = round((2^d / q) x)
     //
-    // Here and in decompression, we leverage the following fact:
+    // Here and in decompression, we leverage the following facts:
     //
     //   round(a / b) = floor((a + b/2) / b)
+    //   a / q ~= (a * x) >> s where x >> s ~= 1/q
     fn compress<D: CompressionFactor>(&mut self) -> &Self {
-        const Q_HALF: u32 = (FieldElement::Q32 - 1) / 2;
-        let x = u32::from(self.0);
-        let y = ((x << D::USIZE) + Q_HALF) / FieldElement::Q32;
+        const Q_HALF: u64 = (FieldElement::Q64 + 1) >> 1;
+        let x = u64::from(self.0);
+        let y = ((((x << D::USIZE) + Q_HALF) * D::DIV_MUL) >> D::DIV_SHIFT).truncate();
         self.0 = y.truncate() & D::MASK;
         self
     }
-    // Equation 4.6: Decomporess_d(x) = round((q / 2^d) x)
+
+    // Equation 4.6: Decompress_d(x) = round((q / 2^d) x)
     fn decompress<D: CompressionFactor>(&mut self) -> &Self {
         let x = u32::from(self.0);
         let y = ((x * FieldElement::Q32) + D::POW2_HALF) >> D::USIZE;
         self.0 = y.truncate();
         self
     }
 }
-
 impl Compress for Polynomial {
     fn compress<D: CompressionFactor>(&mut self) -> &Self {
         for x in &mut self.0 {
@@ -84,37 +90,100 @@ impl<K: ArraySize> Compress for PolynomialVector<K> {
 pub(crate) mod test {
     use super::*;
     use hybrid_array::typenum::{U1, U10, U11, U12, U4, U5, U6};
+    use num_rational::Ratio;
+
+    fn rational_compress<D: CompressionFactor>(input: u16) -> u16 {
+        let fraction = Ratio::new(u32::from(input) * (1 << D::USIZE), FieldElement::Q32);
+        (fraction.round().to_integer() as u16) & D::MASK
+    }
+
+    fn rational_decompress<D: CompressionFactor>(input: u16) -> u16 {
+        let fraction = Ratio::new(u32::from(input) * FieldElement::Q32, 1 << D::USIZE);
+        fraction.round().to_integer() as u16
+    }
 
-    // Verify that the integer compression routine produces the same results as rounding with
-    // floats.
-    fn compression_known_answer_test<D: CompressionFactor>() {
-        let fq: f64 = FieldElement::Q as f64;
-        let f2d: f64 = 2.0_f64.powi(D::I32);
+    // Verify against inequality 4.7
+    #[allow(clippy::integer_division_remainder_used)]
+    fn compression_decompression_inequality<D: CompressionFactor>() {
+        const QI32: i32 = FieldElement::Q as i32;
+        let error_threshold = Ratio::new(FieldElement::Q, 1 << D::USIZE).to_integer() as i32;
 
         for x in 0..FieldElement::Q {
-            let fx = x as f64;
-            let mut x = FieldElement(x);
+            let mut y = FieldElement(x);
+            y.compress::<D>();
+            y.decompress::<D>();
+
+            let mut error = i32::from(y.0) - i32::from(x) + QI32;
+            if error > (QI32 - 1) / 2 {
+                error -= QI32;
+            }
+
+            assert!(
+                error.abs() <= error_threshold,
+                "Inequality failed for x = {x}: error = {}, error_threshold = {error_threshold}, D = {:?}",
+                error.abs(),
+                D::USIZE
+            );
+        }
+    }
 
-            // Verify equivalence of compression
-            x.compress::<D>();
-            let fcx = ((f2d / fq * fx).round() as Integer) % (1 << D::USIZE);
-            assert_eq!(x.0, fcx);
+    fn decompression_compression_equality<D: CompressionFactor>() {
+        for x in 0..(1 << D::USIZE) {
+            let mut y = FieldElement(x);
+            y.decompress::<D>();
+            y.compress::<D>();
 
-            // Verify equivalence of decompression
-            x.decompress::<D>();
-            let fdx = (fq / f2d * (fcx as f64)).round() as Integer;
-            assert_eq!(x.0, fdx);
+            assert_eq!(y.0, x, "failed for x: {}, D: {}", x, D::USIZE);
+        }
+    }
+
+    fn decompress_KAT<D: CompressionFactor>() {
+        for y in 0..(1 << D::USIZE) {
+            let x_expected = rational_decompress::<D>(y);
+            let mut x_actual = FieldElement(y);
+            x_actual.decompress::<D>();
+
+            assert_eq!(x_expected, x_actual.0);
+        }
+    }
+
+    fn compress_KAT<D: CompressionFactor>() {
+        for x in 0..FieldElement::Q {
+            let y_expected = rational_compress::<D>(x);
+            let mut y_actual = FieldElement(x);
+            y_actual.compress::<D>();
+
+            assert_eq!(y_expected, y_actual.0, "for x: {}, D: {}", x, D::USIZE);
         }
     }
 
+    fn compress_decompress_properties<D: CompressionFactor>() {
+        compression_decompression_inequality::<D>();
+        decompression_compression_equality::<D>();
+    }
+
+    fn compress_decompress_KATs<D: CompressionFactor>() {
+        decompress_KAT::<D>();
+        compress_KAT::<D>();
+    }
+
     #[test]
-    fn compress_decompress() {
-        compression_known_answer_test::<U1>();
-        compression_known_answer_test::<U4>();
-        compression_known_answer_test::<U5>();
-        compression_known_answer_test::<U6>();
-        compression_known_answer_test::<U10>();
-        compression_known_answer_test::<U11>();
-        compression_known_answer_test::<U12>();
+    fn decompress_compress() {
+        compress_decompress_properties::<U1>();
+        compress_decompress_properties::<U4>();
+        compress_decompress_properties::<U5>();
+        compress_decompress_properties::<U6>();
+        compress_decompress_properties::<U10>();
+        compress_decompress_properties::<U11>();
+        // preservation under decompression first only holds for d < 12
+        compression_decompression_inequality::<U12>();
+
+        compress_decompress_KATs::<U1>();
+        compress_decompress_KATs::<U4>();
+        compress_decompress_KATs::<U5>();
+        compress_decompress_KATs::<U6>();
+        compress_decompress_KATs::<U10>();
+        compress_decompress_KATs::<U11>();
+        compress_decompress_KATs::<U12>();
     }
 }

ml-kem/src/compress.rs

@@ -165,11 +165,13 @@ pub(crate) mod test {
         D: ArraySize + Rem<N>,
         Mod<D, N>: Zero,
     {
+        #[allow(clippy::integer_division_remainder_used)]
         fn repeat(&self) -> Array<T, D> {
             Array::from_fn(|i| self[i % N::USIZE].clone())
         }
     }
 
+    #[allow(clippy::integer_division_remainder_used)]
     fn byte_codec_test<D>(decoded: DecodedValue, encoded: EncodedPolynomial<D>)
     where
         D: EncodingSize,
@@ -247,6 +249,7 @@ pub(crate) mod test {
         byte_codec_test::<U12>(decoded, encoded);
     }
 
+    #[allow(clippy::integer_division_remainder_used)]
     #[test]
     fn byte_codec_12_mod() {
         // DecodeBytes_12 is required to reduce mod q

ml-kem/src/encode.rs

@@ -6,6 +6,7 @@
     html_favicon_url = "https://raw.githubusercontent.com/RustCrypto/meta/master/logo.svg"
 )]
 #![warn(clippy::pedantic)] // Be pedantic by default
+#![warn(clippy::integer_division_remainder_used)] // Be judicious about using `/` and `%`
 #![allow(non_snake_case)] // Allow notation matching the spec
 #![allow(clippy::clone_on_copy)] // Be explicit about moving data
 #![deny(missing_docs)] // Require all public interfaces to be documented

ml-kem/src/lib.rs

@@ -117,6 +117,7 @@ where
     let mut x = 0usize;
     while x < max {
         let mut y = 0usize;
+        #[allow(clippy::integer_division_remainder_used)]
         while y < max {
             let x_ones = x.count_ones() as u16;
             let y_ones = y.count_ones() as u16;