Run AES KATs from central harness

drivers/src/aes.rs

@@ -193,17 +193,14 @@ impl Aes {
     /// Create a new AES driver.
     ///
     /// Runs the non-GCM KATs (ECB, CBC, CTR, CMAC) at construction time.
-    /// GCM and CMAC-KDF KATs are run by `AesGcm::new()` instead.
+    /// GCM and CMAC-KDF KATs are run by the central KAT harness.
     pub fn new(aes: AesReg, aes_clp: AesClpReg) -> CaliptraResult<Self> {
         if cfg!(feature = "rom") {
             panic!("Do not use in ROM!");
         }
         let mut aes = Self { aes, aes_clp };
         if cfg!(feature = "runtime") {
-            crate::kats::execute_ecb_kat(&mut aes)?;
-            crate::kats::execute_cbc_kat(&mut aes)?;
-            crate::kats::execute_ctr_kat(&mut aes)?;
-            crate::kats::execute_cmac_kat(&mut aes)?;
+            aes.run_kats()?;
         }
         Ok(aes)
     }
@@ -212,6 +209,15 @@ impl Aes {
         Self { aes, aes_clp }
     }
 
+    /// Run the non-GCM KATs (ECB, CBC, CTR, CMAC).
+    pub fn run_kats(&mut self) -> CaliptraResult<()> {
+        crate::kats::execute_ecb_kat(self)?;
+        crate::kats::execute_cbc_kat(self)?;
+        crate::kats::execute_ctr_kat(self)?;
+        crate::kats::execute_cmac_kat(self)?;
+        Ok(())
+    }
+
     /// Seed the AES Trivium stream cipher primitive with fresh entropy.
     ///
     /// After reset, the Trivium primitive is initialized to a netlist constant
@@ -1455,36 +1461,34 @@ impl Aes {
     }
 }
 
-/// AES-GCM driver with compile-time KAT guarantee.
+/// AES-GCM driver.
 ///
-/// This struct wraps an `Aes` driver and ensures that the GCM KAT
-/// is executed before any GCM operations can be performed.
-/// This provides compile-time guarantees similar to the `Sha1` driver.
+/// This struct wraps an `Aes` driver for GCM and CMAC operations.
+/// GCM and CMAC-KDF KATs are exposed through `run_kats()` and invoked by the
+/// central KAT harness.
 pub struct AesGcm {
     aes: Aes,
 }
 
 #[allow(clippy::too_many_arguments)]
 impl AesGcm {
-    /// Create a new AesGcm driver after running the GCM KAT.
+    /// Create a new AesGcm driver after seeding AES entropy.
     ///
-    /// Seeds the AES Trivium stream cipher with fresh entropy from the TRNG,
-    /// then runs the GCM and CMAC-KDF KATs.
+    /// Seeds the AES Trivium stream cipher with fresh entropy from the TRNG.
+    /// GCM and CMAC-KDF KATs are run by `run_kats()`.
     ///
     /// # Arguments
     ///
     /// * `aes` - AES register block
     /// * `aes_clp` - AES CLP register block
-    /// * `trng` - TRNG driver (used for entropy seeding and KATs)
+    /// * `trng` - TRNG driver used for entropy seeding
     ///
     /// # Returns
     ///
-    /// * `CaliptraResult<Self>` - The AesGcm driver if KATs pass
+    /// * `CaliptraResult<Self>` - The AesGcm driver if entropy seeding succeeds
     pub fn new(aes: AesReg, aes_clp: AesClpReg, trng: &mut Trng) -> CaliptraResult<Self> {
         let mut aes = Aes::new_gcm(aes, aes_clp);
         aes.seed_entropy_if(trng)?;
-        crate::kats::execute_gcm_kat(&mut aes, trng)?;
-        crate::kats::execute_cmackdf_kat(&mut aes)?;
         Ok(Self { aes })
     }
 
@@ -1528,8 +1532,7 @@ impl AesGcm {
 
     /// Run the GCM and CMAC-KDF KATs.
     ///
-    /// This is used by FIPS SELF_TEST command to re-run the KATs on demand.
-    /// The KATs are also run at `AesGcm::new()` construction time.
+    /// This is used by the central KAT harness.
     pub fn run_kats(&mut self, trng: &mut Trng) -> CaliptraResult<()> {
         crate::kats::execute_gcm_kat(&mut self.aes, trng)?;
         crate::kats::execute_cmackdf_kat(&mut self.aes)?;

kat/Cargo.toml

@@ -11,7 +11,7 @@ doctest = false
 
 [features]
 default = []
-# When building for ROM, skip AES KATs since ROM uses AesGcm::new() which runs its own KAT
+# When building for ROM, use AesGcm for the GCM and CMAC-KDF KATs.
 rom = ["caliptra-drivers/rom"]
 fips-test-hooks = ["caliptra-drivers/fips-test-hooks"]
 

kat/src/lib.rs

@@ -102,10 +102,7 @@ pub fn execute_kat(env: &mut KatsEnv<'_, '_>) -> CaliptraResult<InitializedDrive
     #[cfg(not(feature = "rom"))]
     {
         caliptra_drivers::kats::execute_cmackdf_kat(env.aes)?;
-        caliptra_drivers::kats::execute_ecb_kat(env.aes)?;
-        caliptra_drivers::kats::execute_cbc_kat(env.aes)?;
-        caliptra_drivers::kats::execute_ctr_kat(env.aes)?;
-        caliptra_drivers::kats::execute_cmac_kat(env.aes)?;
+        env.aes.run_kats()?;
         caliptra_drivers::kats::execute_gcm_kat(env.aes, env.trng)?;
         if let Some(mlkem1024) = env.mlkem1024.as_mut() {
             caliptra_drivers::kats::execute_mlkem1024_kat(mlkem1024)?;

rom/dev/src/main.rs

@@ -76,8 +76,7 @@ pub extern "C" fn rom_entry() -> ! {
         Err(e) => handle_fatal_error(e.into()),
     };
 
-    // Initialize CFI before creating the rest of the environment
-    // (AesGcm::new runs KATs which have CFI annotations)
+    // Initialize CFI before creating the rest of the environment and running KATs.
     if cfg!(feature = "cfi") {
         let mut entropy_gen = || {
             trng.generate4()
@@ -92,7 +91,7 @@ pub extern "C" fn rom_entry() -> ! {
     // before any code that might use memcpy/memset (like AES KATs).
     report_boot_status(RomBootStatus::CfiInitialized.into());
 
-    // Now create the rest of the environment (includes WDT start and AES KATs)
+    // Now create the rest of the environment and start the WDT.
     let mut env = match unsafe { rom_env::RomEnv::new_from_registers(trng) } {
         Ok(env) => env,
         Err(e) => handle_fatal_error(e.into()),

rom/dev/src/rom_env.rs

@@ -78,7 +78,7 @@ pub struct RomEnv {
     /// Dma engine
     pub dma: Dma,
 
-    /// AES-GCM engine (KATs run at construction after WDT starts)
+    /// AES-GCM engine
     pub aes_gcm: AesGcm,
 }
 
@@ -98,13 +98,12 @@ impl RomEnv {
     /// Create the ROM environment. CFI must be initialized before calling this.
     /// Takes ownership of the pre-created TRNG.
     pub unsafe fn new_from_registers(mut trng: Trng) -> CaliptraResult<Self> {
-        // Create SocIfc early so we can start the WDT before running AES KATs
+        // Create SocIfc early so we can start the WDT before running KATs
         let mut soc_ifc = SocIfc::new(SocIfcReg::new());
 
         // Start the Watchdog Timer before running KATs
         crate::wdt::start_wdt(&mut soc_ifc);
 
-        // Create AesGcm which runs the GCM and CMAC-KDF KATs at construction time
         let aes_gcm = AesGcm::new(AesReg::new(), AesClpReg::new(), &mut trng)?;
         Ok(Self {
             doe: DeobfuscationEngine::new(DoeReg::new()),