ML-DSA-65 reference implementation after adding SCT protection.

This reverts commit 75b3d35.
formosa-crypto · Nov 29, 2024 · 6023050 · 6023050
1 parent 75b3d35
commit 6023050
Show file tree

Hide file tree

Showing 11 changed files with 107 additions and 6 deletions.
diff --git a/ref/ml_dsa_65/arithmetic/rounding.jinc b/ref/ml_dsa_65/arithmetic/rounding.jinc
@@ -65,33 +65,45 @@ namespace coefficient {
         return a0_32, a1;
     }
 
-    fn make_hint(reg u32 a0, reg u32 a1) -> reg u32 {
+    fn make_hint(reg u32 a0, reg u32 a1, reg u64 msf) -> reg u32, #msf reg u64 {
         reg u32 result;
         inline bool condition;
 
         result = 0;
 
         condition = a0 >s GAMMA2;
         if (condition) {
+            msf = #update_msf(condition, msf);
+
             result = 1;
         } else {
+            msf = #update_msf(!condition, msf);
+
             condition = a0 <s -GAMMA2;
             if (condition) {
+                msf = #update_msf(condition, msf);
                 result = 1;
             } else {
+                msf = #update_msf(!condition, msf);
+
                 condition = a0 == -GAMMA2;
                 if (condition) {
+                    msf = #update_msf(condition, msf);
+
                     condition = a1 != 0;
                     if (condition) {
+                        msf = #update_msf(condition, msf);
                         result = 1;
                     } else {
+                        msf = #update_msf(!condition, msf);
                     }
                 } else {
+                    msf = #update_msf(!condition, msf);
                 }
             }
         }
 
-        return result;
+        return result, msf;
     }
 
     inline
@@ -135,7 +147,7 @@ namespace coefficient {
         return result;
     }
 
-    fn use_hint(reg u32 a hint) -> reg u32 {
+    fn use_hint(reg u32 a hint, reg u64 msf) -> reg u32, #msf reg u64 {
         reg u32 a0;
         reg u32 a1;
 
@@ -145,17 +157,22 @@ namespace coefficient {
 
         b = hint != 0;
         if (b) {
+            msf = #update_msf(b, msf);
+
             b = a0 >s 0;
             if (b) {
+                msf = #update_msf(b, msf);
                 a1 += 1;
                 a1 &= 15;
             } else {
+                msf = #update_msf(!b, msf);
                 a1 -= 1;
                 a1 &= 15;
             }
         } else {
+            msf = #update_msf(!b, msf);
         }
 
-        return a1;
+        return a1, msf;
     }
 }
diff --git a/ref/ml_dsa_65/encoding/signature.jinc b/ref/ml_dsa_65/encoding/signature.jinc
@@ -21,6 +21,8 @@ namespace signature {
         reg u64 i j hints_written hint_offset;
         reg u32 hint_coefficient;
 
+        reg u64 msf;
+
         inline int k;
         inline bool condition;
 
@@ -42,27 +44,37 @@ namespace signature {
 
 
         // Write out hints
+        msf = #init_msf();
 
         hints_written = 0;
         i = 0;
         while {condition = i < ROWS_IN_MATRIX_A;} (condition) {
+            msf = #update_msf(condition, msf);
+
             j = 0;
             while {condition = j < COEFFICIENTS_IN_POLYNOMIAL;} (condition) {
+                msf = #update_msf(condition, msf);
+
                 hint_offset = i;
                 hint_offset <<= LOG_2_COEFFICIENTS_IN_POLYNOMIAL;
                 hint_offset += j;
 
                 hint_coefficient = hint[hint_offset];
+                hint_coefficient = #protect_32(hint_coefficient, msf);
 
                 condition = hint_coefficient != 0;
                 if (condition) {
+                    msf = #update_msf(condition, msf);
+
                     signature[START_OF_HINT + hints_written] = (8u)j;
                     hints_written += 1;
                 } else {
+                    msf = #update_msf(!condition, msf);
                 }
 
                 j += 1;
             }
+            msf = #update_msf(!condition, msf);
 
             signature[START_OF_HINT + MAX_ONES_IN_HINT + i] = (8u) hints_written;
             i += 1;
@@ -105,6 +117,7 @@ namespace signature {
             }
 
             current_true_hints_seen = (64u)hint_encoded[MAX_ONES_IN_HINT + encoded_offset];
+            _ = #init_msf();
 
             if (current_true_hints_seen < previous_true_hints_seen) {
                 ill_formed_hint = 1;
@@ -121,9 +134,11 @@ namespace signature {
                     } (done == 0) {
                         // Coefficients are ordered for strong unforgeability
                         hint_at_j = (64u)hint_encoded[j];
+                        _ = #init_msf();
 
                         if (j > previous_true_hints_seen) {
                             hint_at_j_minus_one = (64u)hint_encoded[j - 1];
+                            _ = #init_msf();
 
                             if (hint_at_j <= hint_at_j_minus_one) {
                                 ill_formed_hint = 1;
@@ -155,6 +170,7 @@ namespace signature {
             done |= ill_formed_hint;
         } (done == 0) {
             hint = (64u)hint_encoded[encoded_offset];
+            _ = #init_msf();
 
             if(hint != 0) {
                 ill_formed_hint = 1;

diff --git a/ref/ml_dsa_65/keygen.jinc b/ref/ml_dsa_65/keygen.jinc
@@ -87,6 +87,7 @@ fn keygen_internal(reg ptr u8[32] randomness)
 
     s1 = row_vector::ntt(s1);
 
+    _ = #init_msf();
     matrix_A = sample::matrix_A(seed_for_matrix_A);
 
     t = row_vector::multiply_with_matrix_A(matrix_A, s1);

diff --git a/ref/ml_dsa_65/ml_dsa.jazz b/ref/ml_dsa_65/ml_dsa.jazz
@@ -33,6 +33,8 @@ export fn ml_dsa_65_keygen(
     reg u8 byte;
     reg u64 i;
 
+    _ = #init_msf();
+
     () = #spill(verification_key_out, signing_key_out);
     verification_key, signing_key = keygen_internal(randomness);
     () = #unspill(verification_key_out, signing_key_out);
@@ -74,6 +76,8 @@ export fn ml_dsa_65_sign(
     #flexible reg u8 byte;
     reg u64 i = 0;
 
+    _ = #init_msf();
+
     i = 0;
     while(i < SIGNING_KEY_SIZE) {
         byte = (u8)[signing_key + i];
@@ -112,6 +116,8 @@ export fn ml_dsa_65_verify(
     #public reg u8 byte;
     reg u64 i;
 
+    _ = #init_msf();
+
     i = 0;
     while(i < VERIFICATION_KEY_SIZE) {
         #declassify byte = (u8)[verification_key + i];

diff --git a/ref/ml_dsa_65/polynomial.jinc b/ref/ml_dsa_65/polynomial.jinc
@@ -137,16 +137,22 @@ namespace polynomial {
         reg u32 weight a0 a1 hint;
         reg u64 i;
 
+        reg u64 msf = #init_msf();
+
         inline bool condition;
 
         weight = 0;
         i = 0;
         while {condition = i < COEFFICIENTS_IN_POLYNOMIAL;} (condition) {
+            msf = #update_msf(condition, msf);
+
             a0 = f0[i];
+            a0 = #protect_32(a0, msf);
 
             a1 = f1[i];
+            a1 = #protect_32(a1, msf);
 
-            hint = coefficient::make_hint(a0, a1);
+            hint, msf = coefficient::make_hint(a0, a1, msf);
             hints[i] = hint;
             weight += hint;
             i += 1;
@@ -196,15 +202,21 @@ namespace polynomial {
         reg u64 i;
         reg u32 a h;
 
+        reg u64 msf = #init_msf();
+
         inline bool b;
 
         i = 0;
         while { b = i < COEFFICIENTS_IN_POLYNOMIAL; } (b) {
+            msf = #update_msf(b, msf);
+
             h = hints[i];
+            h = #protect_32(h, msf);
 
             a = commitment[i];
+            a = #protect_32(a, msf);
 
-            a = coefficient::use_hint(a, h);
+            a, msf = coefficient::use_hint(a, h, msf);
 
             commitment[i] = a;
 

diff --git a/ref/ml_dsa_65/sample/challenge.jinc b/ref/ml_dsa_65/sample/challenge.jinc
@@ -72,6 +72,7 @@ namespace sample {
                     xof_offset = 0;
                 }
                 sample_at = (64u) state[u8 xof_offset];
+                _ = #init_msf();
 
                 xof_offset += 1;
             } (sample_at > i)

diff --git a/ref/ml_dsa_65/sample/error_vectors.jinc b/ref/ml_dsa_65/sample/error_vectors.jinc
@@ -45,44 +45,63 @@ fn rejection_sample_less_than_eta(
 
     inline bool b;
 
+    reg u64 msf = #init_msf();
+
     xof_offset = 0;
 
     while { b = xof_offset < SHAKE256_RATE; } (b) {
+        msf = #update_msf(b, msf);
+
         b = sampled < COEFFICIENTS_IN_POLYNOMIAL;
         if (b) {
+            msf = #update_msf(b, msf);
+
             byte = state[u8 xof_offset];
 
             #declassify try_coefficient = (32u)byte;
             try_coefficient &= 0x0F;
 
+            try_coefficient = #protect_32(try_coefficient, msf);
+
             b = try_coefficient < 9;
             if (b) {
+                msf = #update_msf(b, msf);
+
                 temp = 4;
                 temp -= try_coefficient;
                 error[sampled] = temp;
 
                 sampled += 1;
             } else {
+                msf = #update_msf(!b, msf);
             }
 
             b = sampled < COEFFICIENTS_IN_POLYNOMIAL;
             if (b) {
+                msf = #update_msf(b, msf);
+
                 #declassify try_coefficient = (32u)byte;
                 try_coefficient >>= 4;
+                try_coefficient = #protect_32(try_coefficient, msf);
 
                 b = try_coefficient < 9;
                 if (b) {
+                    msf = #update_msf(b, msf);
+
                     temp = 4;
                     temp -= try_coefficient;
                     error[sampled] = temp;
 
                     sampled += 1;
                 } else {
+                    msf = #update_msf(!b, msf);
                 }
             } else {
+                msf = #update_msf(!b, msf);
             }
 
         } else {
+            msf = #update_msf(!b, msf);
         }
 
         xof_offset += 1;

diff --git a/ref/ml_dsa_65/sample/mask.jinc b/ref/ml_dsa_65/sample/mask.jinc
@@ -56,6 +56,8 @@ fn sample_mask_polynomial(
         state = _keccakf1600_ref1(state);
         () = #unspill(mask, mask_encoded_offset);
 
+        _ = #init_msf();
+
         i = 0;
         while(i < SHAKE256_RATE) {
             byte = state[u8 i];

diff --git a/ref/ml_dsa_65/sample/matrix_A.jinc b/ref/ml_dsa_65/sample/matrix_A.jinc
@@ -42,11 +42,17 @@ fn rejection_sample_less_than_field_modulus(
 
     inline int k;
 
+    reg u64 msf = #init_msf();
+
     xof_offset = 0;
 
     while { b = xof_offset < SHAKE128_RATE; } (b) {
+        msf = #update_msf(b, msf);
+
         b = sampled < COEFFICIENTS_IN_POLYNOMIAL;
         if (b) {
+            msf = #update_msf(b, msf);
+
             try_coefficient = 0;
             for k = 0 to 3 {
                 #declassify byte = (32u)state[u8 xof_offset + k];
@@ -55,15 +61,19 @@ fn rejection_sample_less_than_field_modulus(
                 try_coefficient |= byte;
             }
             try_coefficient &= 0x00_7F_FF_FF;
+            try_coefficient = #protect_32(try_coefficient, msf);
 
 
             b = try_coefficient < MODULUS;
             if (b) {
+                msf = #update_msf(b, msf);
                 result[sampled] = try_coefficient;
                 sampled += 1;
             } else {
+                msf = #update_msf(!b, msf);
             }
         } else {
+            msf = #update_msf(!b, msf);
         }
 
         xof_offset += 3;

diff --git a/ref/ml_dsa_65/sign.jinc b/ref/ml_dsa_65/sign.jinc
@@ -90,6 +90,7 @@ fn sign_internal(
 
     #declassify seed_for_matrix_A = signing_key[0:32];
 
+    _ = #init_msf();
     matrix_A = sample::matrix_A(seed_for_matrix_A);
     () = #unspill(randomness);
 
@@ -144,22 +145,29 @@ fn sign_internal(
         #declassify signer_response = row_vector::reduce32(signer_response);
         #declassify infinity_norm_check_result = row_vector::check_infinity_norm(signer_response, (int) (GAMMA1 - BETA));
 
+        _ = #init_msf();
+
         if(infinity_norm_check_result == 0) {
             cs2 = column_vector::multiply_by_polynomial(s2, verifier_challenge);
             cs2 = column_vector::invert_ntt_montgomery(cs2);
             w0_minus_cs2 = column_vector::subtract(w0, cs2);
             #declassify w0_minus_cs2 = column_vector::reduce32(w0_minus_cs2);
             #declassify infinity_norm_check_result = column_vector::check_infinity_norm(w0_minus_cs2, (int) (GAMMA2 - BETA));
 
+            _ = #init_msf();
+
             if(infinity_norm_check_result == 0) {
                 ct0 = column_vector::multiply_by_polynomial(t0, verifier_challenge);
                 ct0 = column_vector::invert_ntt_montgomery(ct0);
                 #declassify ct0 = column_vector::reduce32(ct0);
                 #declassify infinity_norm_check_result = column_vector::check_infinity_norm(ct0, GAMMA2);
+
+                _ = #init_msf();
                 if(infinity_norm_check_result == 0) {
                     w0_minus_cs2_plus_ct0 = column_vector::add(w0_minus_cs2, ct0);
 
                     ones_in_hint, hint = column_vector::make_hint(w0_minus_cs2_plus_ct0, w1);
+                    _ = #init_msf();
                     if(ones_in_hint <= MAX_ONES_IN_HINT) {
                         exit_rejection_sampling_loop = 1;
                     }