Make Xoshiro128 a normal rust type, and remove call to Trng::assume_i…

…nit() The latter fixes a serious bug with the csrng not being initialized before we seed the xorshiro prng.
chipsalliance · Aug 11, 2023 · ddcba60 · ddcba60
1 parent 5354762
commit ddcba60
Show file tree

Hide file tree

Showing 7 changed files with 85 additions and 212 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/cfi/lib/Cargo.toml b/cfi/lib/Cargo.toml
@@ -13,6 +13,7 @@ caliptra_common = { workspace = true, default-features = false }
 caliptra-drivers.workspace = true
 caliptra-registers.workspace = true
 ufmt.workspace = true
+zerocopy.workspace = true
 
 [dev-dependencies]
 caliptra-cfi-derive.workspace = true

diff --git a/cfi/lib/src/cfi_ctr.rs b/cfi/lib/src/cfi_ctr.rs
@@ -17,7 +17,7 @@ References:
 --*/
 
 use crate::cfi::{cfi_panic, CfiPanicInfo};
-use crate::xoshiro::{Xoshiro128, Xoshiro128Reg};
+use crate::xoshiro::Xoshiro128;
 #[cfg(not(feature = "cfi-test"))]
 use caliptra_common::memory_layout::{CFI_MASK_ORG, CFI_VAL_ORG};
 use core::default::Default;
@@ -67,14 +67,33 @@ impl Default for CfiInt {
     }
 }
 
+fn prng() -> &'static Xoshiro128 {
+    use caliptra_drivers::memory_layout::CFI_XO_S0_ORG;
+    if cfg!(feature = "cfi-test") {
+        static mut STATE: Xoshiro128 = Xoshiro128::new_unseeded();
+        unsafe { &STATE }
+    } else {
+        unsafe { &*(CFI_XO_S0_ORG as *const Xoshiro128) }
+    }
+}
+
 /// CFI counter
 pub enum CfiCounter {}
 
 impl CfiCounter {
     /// Reset counter
     #[inline(never)]
-    pub fn reset() {
-        Xoshiro128::new(Xoshiro128Reg);
+    pub fn reset(trng: &mut caliptra_drivers::Trng) {
+        prng().seed_from_trng(trng);
+        Self::reset_internal();
+    }
+
+    #[cfg(feature = "cfi-test")]
+    pub fn reset_for_test() {
+        Self::reset_internal()
+    }
+
+    fn reset_internal() {
         Self::write(CfiInt::default());
     }
 
@@ -152,7 +171,7 @@ impl CfiCounter {
     pub fn delay() {
         #[cfg(all(target_arch = "riscv32", feature = "cfi", feature = "cfi-counter"))]
         unsafe {
-            let cycles = Xoshiro128::assume_init(Xoshiro128Reg).next() % 256;
+            let cycles = prng().next() % 256;
             let real_cyc = 1 + cycles / 2;
             core::arch::asm!(
                 "1:",

diff --git a/cfi/lib/src/lib.rs b/cfi/lib/src/lib.rs
@@ -17,4 +17,4 @@ mod xoshiro;
 
 pub use cfi::*;
 pub use cfi_ctr::{CfiCounter, CfiInt};
-pub use xoshiro::{Xoshiro128, Xoshiro128Reg};
+pub use xoshiro::Xoshiro128;
diff --git a/cfi/lib/src/xoshiro.rs b/cfi/lib/src/xoshiro.rs
@@ -15,141 +15,25 @@ References:
 
 --*/
 
-use crate::{cfi_panic, CfiPanicInfo};
-#[cfg(not(feature = "cfi-test"))]
-use caliptra_common::memory_layout::{CFI_XO_S0_ORG, CFI_XO_S1_ORG, CFI_XO_S2_ORG, CFI_XO_S3_ORG};
-
-#[cfg(feature = "cfi-test")]
-static mut S0: u32 = 0u32;
-#[cfg(feature = "cfi-test")]
-static mut S1: u32 = 0u32;
-#[cfg(feature = "cfi-test")]
-static mut S2: u32 = 0u32;
-#[cfg(feature = "cfi-test")]
-static mut S3: u32 = 0u32;
-
-pub struct Xoshiro128Reg;
-
-impl Xoshiro128Reg {
-    pub fn s0(&self) -> u32 {
-        unsafe {
-            #[cfg(feature = "cfi-test")]
-            {
-                core::ptr::read_volatile(&S0 as *const u32)
-            }
-
-            #[cfg(not(feature = "cfi-test"))]
-            {
-                core::ptr::read_volatile(CFI_XO_S0_ORG as *const u32)
-            }
-        }
-    }
-
-    pub fn set_s0(&self, val: u32) {
-        unsafe {
-            #[cfg(feature = "cfi-test")]
-            {
-                core::ptr::write_volatile(&mut S0 as *mut u32, val);
-            }
-
-            #[cfg(not(feature = "cfi-test"))]
-            {
-                core::ptr::write_volatile(CFI_XO_S0_ORG as *mut u32, val);
-            }
-        }
-    }
-
-    pub fn s1(&self) -> u32 {
-        unsafe {
-            #[cfg(feature = "cfi-test")]
-            {
-                core::ptr::read_volatile(&S1 as *const u32)
-            }
-
-            #[cfg(not(feature = "cfi-test"))]
-            {
-                core::ptr::read_volatile(CFI_XO_S1_ORG as *const u32)
-            }
-        }
-    }
-
-    pub fn set_s1(&self, val: u32) {
-        unsafe {
-            #[cfg(feature = "cfi-test")]
-            {
-                core::ptr::write_volatile(&mut S1 as *mut u32, val);
-            }
-
-            #[cfg(not(feature = "cfi-test"))]
-            {
-                core::ptr::write_volatile(CFI_XO_S1_ORG as *mut u32, val);
-            }
-        }
-    }
-
-    pub fn s2(&self) -> u32 {
-        unsafe {
-            #[cfg(feature = "cfi-test")]
-            {
-                core::ptr::read_volatile(&S2 as *const u32)
-            }
-
-            #[cfg(not(feature = "cfi-test"))]
-            {
-                core::ptr::read_volatile(CFI_XO_S2_ORG as *const u32)
-            }
-        }
-    }
-
-    pub fn set_s2(&self, val: u32) {
-        unsafe {
-            #[cfg(feature = "cfi-test")]
-            {
-                core::ptr::write_volatile(&mut S2 as *mut u32, val);
-            }
-
-            #[cfg(not(feature = "cfi-test"))]
-            {
-                core::ptr::write_volatile(CFI_XO_S2_ORG as *mut u32, val);
-            }
-        }
-    }
-
-    pub fn s3(&self) -> u32 {
-        unsafe {
-            #[cfg(feature = "cfi-test")]
-            {
-                core::ptr::read_volatile(&S3 as *const u32)
-            }
-
-            #[cfg(not(feature = "cfi-test"))]
-            {
-                core::ptr::read_volatile(CFI_XO_S3_ORG as *const u32)
-            }
-        }
-    }
-
-    pub fn set_s3(&self, val: u32) {
-        unsafe {
-            #[cfg(feature = "cfi-test")]
-            {
-                core::ptr::write_volatile(&mut S3 as *mut u32, val);
-            }
+use core::cell::Cell;
 
-            #[cfg(not(feature = "cfi-test"))]
-            {
-                core::ptr::write_volatile(CFI_XO_S3_ORG as *mut u32, val);
-            }
-        }
-    }
-}
+use crate::{cfi_panic, CfiPanicInfo};
 
 /// Provides an implementation of the xoshiro128** algorithm. This implementation is used
 /// on 32-bit when no seed is specified and an instance of the base Random class is constructed.
 /// As such, we are free to implement however we see fit, without back compat concerns around
 /// the sequence of numbers generated or what methods call what other methods.
+#[repr(C)]
 pub struct Xoshiro128 {
-    reg: Xoshiro128Reg,
+    // It is critical for safety that every bit-pattern of this struct
+    // is valid (no padding, no enums, no references), similar to the requirements for
+    // zerocopy::FromBytes.
+    // TODO: unsafe_impl!(T: FromBytes, Cell<T> => FromBytes for Cell<T>) in
+    // zerocopy upstream to check this at compile time
+    s0: Cell<u32>,
+    s1: Cell<u32>,
+    s2: Cell<u32>,
+    s3: Cell<u32>,
 }
 
 impl Xoshiro128 {
@@ -162,40 +46,36 @@ impl Xoshiro128 {
     /// # Returns
     ///
     /// * Instance of the xoshiro128** algorithm
-    pub(crate) fn new(reg: Xoshiro128Reg) -> Self {
-        if !cfg!(feature = "cfi-test") {
-            let mut trng = unsafe {
-                caliptra_drivers::Trng::assume_initialized(
-                    caliptra_registers::csrng::CsrngReg::new(),
-                    caliptra_registers::entropy_src::EntropySrcReg::new(),
-                    caliptra_registers::soc_ifc_trng::SocIfcTrngReg::new(),
-                    &caliptra_registers::soc_ifc::SocIfcReg::new(),
-                )
-            };
-
-            loop {
-                if let Ok(entropy) = trng.generate() {
-                    reg.set_s0(entropy.0[0]);
-                    reg.set_s1(entropy.0[1]);
-                    reg.set_s2(entropy.0[2]);
-                    reg.set_s3(entropy.0[3]);
-                } else {
-                    cfi_panic(CfiPanicInfo::TrngError)
-                }
-
-                // Atlease one value must be non-zero
-                if reg.s0() | reg.s1() | reg.s2() | reg.s3() != 0 {
-                    break;
-                }
-            }
-        }
+    pub(crate) const fn new_unseeded() -> Self {
+        Self::new_with_seed(0, 0, 0, 0)
+    }
 
-        Self { reg }
+    /// Create a new instance of the xoshiro128** algorithm.
+    pub const fn new_with_seed(s0: u32, s1: u32, s2: u32, s3: u32) -> Self {
+        Self {
+            s0: Cell::new(s0),
+            s1: Cell::new(s1),
+            s2: Cell::new(s2),
+            s3: Cell::new(s3),
+        }
     }
 
-    /// Create a new instance of the xoshiro128** algorithm from a register
-    pub fn assume_init(reg: Xoshiro128Reg) -> Self {
-        Self { reg }
+    pub fn seed_from_trng(&self, trng: &mut caliptra_drivers::Trng) {
+        loop {
+            if let Ok(entropy) = trng.generate() {
+                self.s0.set(entropy.0[0]);
+                self.s1.set(entropy.0[1]);
+                self.s2.set(entropy.0[2]);
+                self.s3.set(entropy.0[3]);
+            } else {
+                cfi_panic(CfiPanicInfo::TrngError)
+            }
+
+            // Atlease one value must be non-zero
+            if self.s0.get() | self.s1.get() | self.s2.get() | self.s3.get() != 0 {
+                break;
+            }
+        }
     }
 
     /// Get the next random number
@@ -209,10 +89,10 @@ impl Xoshiro128 {
         //     worldwide. This software is distributed without any warranty.
         //
         //     See <http://creativecommons.org/publicdomain/zero/1.0/>.
-        let mut s0 = self.reg.s0();
-        let mut s1 = self.reg.s1();
-        let mut s2 = self.reg.s2();
-        let mut s3 = self.reg.s3();
+        let mut s0 = self.s0.get();
+        let mut s1 = self.s1.get();
+        let mut s2 = self.s2.get();
+        let mut s3 = self.s3.get();
 
         let result = u32::wrapping_mul(u32::wrapping_mul(s1, 5).rotate_left(7), 9);
         let t = s1 << 9;
@@ -225,10 +105,10 @@ impl Xoshiro128 {
         s2 ^= t;
         s3 = s3.rotate_left(11);
 
-        self.reg.set_s0(s0);
-        self.reg.set_s1(s1);
-        self.reg.set_s2(s2);
-        self.reg.set_s3(s3);
+        self.s0.set(s0);
+        self.s1.set(s1);
+        self.s2.set(s2);
+        self.s3.set(s3);
 
         result
     }