From 8c492889e669ffcbf7a17c0aad22147dcd333fc1 Mon Sep 17 00:00:00 2001
From: Kor Nielsen <kor@google.com>
Date: Tue, 29 Aug 2023 21:26:40 -0700
Subject: [PATCH] Fix 1600 byte ROM size regression in #684

When callers from separate crates call a large generic function like
verify_lms_signature_cfi(), rustc 1.70 may build multiple versions
(depending on optimizer heuristics), even when all the generic
parameters are identical. This is bad, as it can bloat the binary and
the second copy violates the FIPS requirements that the same machine
code be used for the KAT as the actual implementation. To defend against
it, we provide a non-generic function that production firmware should
call instead.
---
 drivers/src/lms.rs                            | 51 +++++++++++++++++--
 drivers/test-fw/src/bin/lms_32_tests.rs       |  6 +--
 drivers/test-fw/src/bin/negative_tests_lms.rs |  2 +-
 3 files changed, 51 insertions(+), 8 deletions(-)

diff --git a/drivers/src/lms.rs b/drivers/src/lms.rs
index 3791f27590..5e524c96e3 100644
--- a/drivers/src/lms.rs
+++ b/drivers/src/lms.rs
@@ -379,12 +379,12 @@ impl Lms {
 
     ///  Note: Use this function only if glitch protection is not needed.
     ///        If glitch protection is needed, use `verify_lms_signature_cfi` instead.
-    pub fn verify_lms_signature<const N: usize, const P: usize, const H: usize>(
+    pub fn verify_lms_signature(
         &self,
         sha256_driver: &mut Sha256,
         input_string: &[u8],
-        lms_public_key: &LmsPublicKey<N>,
-        lms_sig: &LmsSignature<N, P, H>,
+        lms_public_key: &LmsPublicKey<6>,
+        lms_sig: &LmsSignature<6, 51, 15>,
     ) -> CaliptraResult<LmsResult> {
         let mut candidate_key =
             self.verify_lms_signature_cfi(sha256_driver, input_string, lms_public_key, lms_sig)?;
@@ -397,7 +397,50 @@ impl Lms {
         result
     }
 
-    pub fn verify_lms_signature_cfi<const N: usize, const P: usize, const H: usize>(
+    ///  Note: Use this function only if glitch protection is not needed.
+    ///        If glitch protection is needed, use `verify_lms_signature_cfi_generic` instead.
+    pub fn verify_lms_signature_generic<const N: usize, const P: usize, const H: usize>(
+        &self,
+        sha256_driver: &mut Sha256,
+        input_string: &[u8],
+        lms_public_key: &LmsPublicKey<N>,
+        lms_sig: &LmsSignature<N, P, H>,
+    ) -> CaliptraResult<LmsResult> {
+        let mut candidate_key = self.verify_lms_signature_cfi_generic(
+            sha256_driver,
+            input_string,
+            lms_public_key,
+            lms_sig,
+        )?;
+        let result = if candidate_key != HashValue::from(lms_public_key.digest) {
+            Ok(LmsResult::SigVerifyFailed)
+        } else {
+            Ok(LmsResult::Success)
+        };
+        candidate_key.0.fill(0);
+        result
+    }
+
+    // When callers from separate crates call a function like
+    // verify_lms_signature_cfi_generic(), Rustc 1.70
+    // may build multiple versions (depending on optimizer heuristics), even when all the
+    // generic parameters are identical. This is bad, as it can bloat the binary and the
+    // second copy violates the FIPS requirements that the same machine code be used for the
+    // KAT as the actual implementation. To defend against it, we provide this non-generic
+    // function that production firmware should call instead.
+    #[inline(never)]
+    pub fn verify_lms_signature_cfi(
+        &self,
+        sha256_driver: &mut Sha256,
+        input_string: &[u8],
+        lms_public_key: &LmsPublicKey<6>,
+        lms_sig: &LmsSignature<6, 51, 15>,
+    ) -> CaliptraResult<HashValue<6>> {
+        self.verify_lms_signature_cfi_generic(sha256_driver, input_string, lms_public_key, lms_sig)
+    }
+
+    #[inline(always)]
+    pub fn verify_lms_signature_cfi_generic<const N: usize, const P: usize, const H: usize>(
         &self,
         sha256_driver: &mut Sha256,
         input_string: &[u8],
diff --git a/drivers/test-fw/src/bin/lms_32_tests.rs b/drivers/test-fw/src/bin/lms_32_tests.rs
index 35ab8d4052..c67cdf9b9d 100644
--- a/drivers/test-fw/src/bin/lms_32_tests.rs
+++ b/drivers/test-fw/src/bin/lms_32_tests.rs
@@ -551,12 +551,12 @@ fn test_lms_lower_32() {
     };
 
     let final_result = Lms::default()
-        .verify_lms_signature(&mut sha256, &MESSAGE, &LMS_PUBLIC_KEY, &FINAL_LMS_SIG)
+        .verify_lms_signature_generic(&mut sha256, &MESSAGE, &LMS_PUBLIC_KEY, &FINAL_LMS_SIG)
         .unwrap();
     assert_eq!(final_result, LmsResult::Success);
 
     let candidate_key = Lms::default()
-        .verify_lms_signature_cfi(&mut sha256, &MESSAGE, &LMS_PUBLIC_KEY, &FINAL_LMS_SIG)
+        .verify_lms_signature_cfi_generic(&mut sha256, &MESSAGE, &LMS_PUBLIC_KEY, &FINAL_LMS_SIG)
         .unwrap();
     assert_eq!(candidate_key, HashValue::from(LMS_PUBLIC_KEY.digest));
 }
@@ -815,7 +815,7 @@ fn test_hss_upper_32() {
     };
 
     let result = Lms::default()
-        .verify_lms_signature(
+        .verify_lms_signature_generic(
             &mut sha256,
             &PUBLIC_BUFFER,
             &HSS_PUBLIC_KEY,
diff --git a/drivers/test-fw/src/bin/negative_tests_lms.rs b/drivers/test-fw/src/bin/negative_tests_lms.rs
index 11437ebf4a..432e1c874a 100644
--- a/drivers/test-fw/src/bin/negative_tests_lms.rs
+++ b/drivers/test-fw/src/bin/negative_tests_lms.rs
@@ -436,7 +436,7 @@ fn test_failures_lms_24() {
     );
 
     assert_eq!(
-        Lms::default().verify_lms_signature(
+        Lms::default().verify_lms_signature_generic(
             &mut sha256,
             &MESSAGE,
             &LmsPublicKey {