From 8ce98548252005f8bbea82901fb4bd66a869394e Mon Sep 17 00:00:00 2001
From: Vishal Mhatre <vishal.mhatre@gmail.com>
Date: Thu, 12 Oct 2023 16:33:50 -0700
Subject: [PATCH] [fix] Additional CFI checks

This fix addresses issue# https://github.com/chipsalliance/caliptra-sw/issues/920 and https://github.com/chipsalliance/caliptra-sw/issues/921
---
 cfi/lib/src/cfi.rs         |  9 ++++++++-
 cfi/lib/src/cfi_counter.rs |  2 +-
 drivers/src/soc_ifc.rs     |  4 ++++
 rom/dev/src/main.rs        | 14 +++++++++++---
 4 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/cfi/lib/src/cfi.rs b/cfi/lib/src/cfi.rs
index 98427d4ee5..8c977f9431 100644
--- a/cfi/lib/src/cfi.rs
+++ b/cfi/lib/src/cfi.rs
@@ -148,7 +148,7 @@ macro_rules! cfi_assert_macro {
         ///
         /// `a` - Left hand side
         /// `b` - Right hand side
-        #[inline(never)]
+        #[inline(always)]
         #[allow(unused)]
         pub fn $name<T>(lhs: T, rhs: T)
         where
@@ -159,6 +159,13 @@ macro_rules! cfi_assert_macro {
                 if !(lhs $op rhs) {
                     cfi_panic(CfiPanicInfo::$panic_info);
                 }
+
+                // Second check for glitch protection
+                CfiCounter::delay();
+                if !(cfi_launder(lhs) $op cfi_launder(rhs)) {
+                    cfi_panic(CfiPanicInfo::$panic_info);
+                }
+
             } else {
                 lhs $op rhs;
             }
diff --git a/cfi/lib/src/cfi_counter.rs b/cfi/lib/src/cfi_counter.rs
index d9ad7d1c8e..17156ab583 100644
--- a/cfi/lib/src/cfi_counter.rs
+++ b/cfi/lib/src/cfi_counter.rs
@@ -82,7 +82,7 @@ pub enum CfiCounter {}
 
 impl CfiCounter {
     /// Reset counter
-    #[inline(never)]
+    #[inline(always)]
     pub fn reset(trng: &mut caliptra_drivers::Trng) {
         prng().seed_from_trng(trng);
         Self::reset_internal();
diff --git a/drivers/src/soc_ifc.rs b/drivers/src/soc_ifc.rs
index f026816659..3e636d6e92 100644
--- a/drivers/src/soc_ifc.rs
+++ b/drivers/src/soc_ifc.rs
@@ -144,6 +144,10 @@ impl SocIfc {
         ((val >> 31) & 1) != 0
     }
 
+    pub fn hw_config_internal_trng(&mut self) -> bool {
+        self.soc_ifc.regs().cptra_hw_config().read().i_trng_en()
+    }
+
     /// Enable or disable WDT1
     ///
     /// # Arguments
diff --git a/rom/dev/src/main.rs b/rom/dev/src/main.rs
index 1ff51c3bf5..77aee8d255 100644
--- a/rom/dev/src/main.rs
+++ b/rom/dev/src/main.rs
@@ -16,13 +16,13 @@ Abstract:
 #![cfg_attr(feature = "fake-rom", allow(unused_imports))]
 
 use crate::{lock::lock_registers, print::HexBytes};
-use caliptra_cfi_lib::CfiCounter;
+use caliptra_cfi_lib::{cfi_assert, cfi_assert_eq, CfiCounter};
 use caliptra_registers::soc_ifc::SocIfcReg;
 use core::hint::black_box;
 
 use caliptra_drivers::{
     cprintln, report_fw_error_fatal, report_fw_error_non_fatal, CaliptraError, Ecc384, Hmac384,
-    KeyVault, Mailbox, ResetReason, Sha256, Sha384, Sha384Acc, ShaAccLockState, SocIfc,
+    KeyVault, Mailbox, ResetReason, Sha256, Sha384, Sha384Acc, ShaAccLockState, SocIfc, Trng,
 };
 use caliptra_error::CaliptraResult;
 use caliptra_image_types::RomInfo;
@@ -69,11 +69,19 @@ pub extern "C" fn rom_entry() -> ! {
 
     if !cfg!(feature = "no-cfi") {
         cprintln!("[state] CFI Enabled");
-        CfiCounter::reset(&mut env.trng);
+        for _ in 0..=2 {
+            CfiCounter::reset(&mut env.trng);
+        }
     } else {
         cprintln!("[state] CFI Disabled");
     }
 
+    // Check if TRNG is correctly sourced as per hw config.
+    match env.trng {
+        Trng::Internal(_) => cfi_assert!(env.soc_ifc.hw_config_internal_trng()),
+        Trng::External(_) => cfi_assert!(!env.soc_ifc.hw_config_internal_trng()),
+    }
+
     let _lifecyle = match env.soc_ifc.lifecycle() {
         caliptra_drivers::Lifecycle::Unprovisioned => "Unprovisioned",
         caliptra_drivers::Lifecycle::Manufacturing => "Manufacturing",