From 65a42402ef8fcfefd2fdfd1da66b9e13dc08e96e Mon Sep 17 00:00:00 2001
From: Sean McGrail <mcgrails@amazon.com>
Date: Tue, 20 Feb 2024 22:44:51 +0000
Subject: [PATCH] Add FIPS indicator test coverage for RSA key-generation
 functions

---
 aws-lc-rs/src/fips.rs           |   4 ++
 aws-lc-rs/src/rsa.rs            |   3 +
 aws-lc-rs/src/rsa/tests/fips.rs | 111 ++++++++++++++++++++++++++++++++
 3 files changed, 118 insertions(+)
 create mode 100644 aws-lc-rs/src/rsa/tests/fips.rs

diff --git a/aws-lc-rs/src/fips.rs b/aws-lc-rs/src/fips.rs
index aa784c60b19..005f7ac1a2d 100644
--- a/aws-lc-rs/src/fips.rs
+++ b/aws-lc-rs/src/fips.rs
@@ -126,7 +126,11 @@ pub(crate) use indicator_check;
 macro_rules! check_fips_service_status {
     ($function:expr) => {{
         use $crate::fips::get_fips_service_status;
+        // Clear the current indicator status first by retrieving it
+        let _ = get_fips_service_status();
+        // do the expression
         let result = $function;
+        // Check indicator after expression
         get_fips_service_status().map(|()| result)
     }};
 }
diff --git a/aws-lc-rs/src/rsa.rs b/aws-lc-rs/src/rsa.rs
index a627332c2b2..5e055e088a1 100644
--- a/aws-lc-rs/src/rsa.rs
+++ b/aws-lc-rs/src/rsa.rs
@@ -77,6 +77,9 @@ pub(crate) use self::signature::RsaVerificationAlgorithmId;
 
 #[cfg(test)]
 mod tests {
+    #[cfg(feature = "fips")]
+    mod fips;
+
     #[cfg(feature = "ring-io")]
     #[test]
     fn test_rsa() {
diff --git a/aws-lc-rs/src/rsa/tests/fips.rs b/aws-lc-rs/src/rsa/tests/fips.rs
new file mode 100644
index 00000000000..6e96753b5be
--- /dev/null
+++ b/aws-lc-rs/src/rsa/tests/fips.rs
@@ -0,0 +1,111 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0 OR ISC
+
+#![cfg(debug_assertions)]
+
+use crate::{
+    fips::{assert_fips_status_indicator, FipsServiceStatus},
+    rsa::{KeyPair, KeySize, PrivateDecryptingKey},
+};
+
+macro_rules! generate_key {
+    ($name:ident, KeyPair, $size:expr) => {
+        #[test]
+        fn $name() {
+            // Using the non-fips generator will not set the indicator
+            let _ =
+                assert_fips_status_indicator!(KeyPair::generate($size), FipsServiceStatus::Unset)
+                    .expect("key generated");
+
+            // Using the fips generator should set the indicator
+            let _ = assert_fips_status_indicator!(
+                KeyPair::generate_fips($size),
+                FipsServiceStatus::Approved
+            )
+            .expect("key generated");
+        }
+    };
+    ($name:ident, PrivateDecryptingKey, $size:expr) => {
+        #[test]
+        fn $name() {
+            // Using the non-fips generator will not set the indicator
+            let _ = assert_fips_status_indicator!(
+                PrivateDecryptingKey::generate($size),
+                FipsServiceStatus::Unset
+            )
+            .expect("key generated");
+
+            // Using the fips generator should set the indicator
+            let _ = assert_fips_status_indicator!(
+                PrivateDecryptingKey::generate_fips($size),
+                FipsServiceStatus::Approved
+            )
+            .expect("key generated");
+        }
+    };
+    ($name:ident, KeyPair, $size:expr, false) => {
+        #[test]
+        fn $name() {
+            // Using the non-fips generator will not set the indicator
+            let _ =
+                assert_fips_status_indicator!(KeyPair::generate($size), FipsServiceStatus::Unset);
+
+            // Using the fips generator should set the indicator
+            let _ = assert_fips_status_indicator!(
+                KeyPair::generate_fips($size),
+                FipsServiceStatus::NonApproved
+            )
+            .expect_err("key size not allowed");
+        }
+    };
+    ($name:ident, PrivateDecryptingKey, $size:expr, false) => {
+        #[test]
+        fn $name() {
+            // Using the non-fips generator will not set the indicator
+            let _ = assert_fips_status_indicator!(
+                PrivateDecryptingKey::generate($size),
+                FipsServiceStatus::Unset
+            );
+
+            // Using the fips generator should set the indicator
+            let _ = assert_fips_status_indicator!(
+                PrivateDecryptingKey::generate_fips($size),
+                FipsServiceStatus::NonApproved
+            )
+            .expect_err("key size not allowed");
+        }
+    };
+}
+
+generate_key!(rsa2048_signing_generate_key, KeyPair, KeySize::Rsa2048);
+generate_key!(rsa3072_signing_generate_key, KeyPair, KeySize::Rsa3072);
+generate_key!(rsa4096_signing_generate_key, KeyPair, KeySize::Rsa4096);
+
+generate_key!(
+    rsa8192_signing_generate_key,
+    KeyPair,
+    KeySize::Rsa8192,
+    false
+);
+
+generate_key!(
+    rsa2048_encryption_generate_key,
+    PrivateDecryptingKey,
+    KeySize::Rsa2048
+);
+generate_key!(
+    rsa3072_encryption_generate_key,
+    PrivateDecryptingKey,
+    KeySize::Rsa3072
+);
+generate_key!(
+    rsa4096_encryption_signing_generate_key,
+    PrivateDecryptingKey,
+    KeySize::Rsa4096
+);
+generate_key!(
+    rsa8192_encryption_generate_key,
+    PrivateDecryptingKey,
+    KeySize::Rsa8192,
+    false
+);