Add support for OP-TEE key derivation

dev_tests/src/ratchet.rs

@@ -36,8 +36,8 @@ fn ratchet_globals() -> Result<()> {
             ("dev_bench/", 1),
             ("litebox/", 9),
             ("litebox_platform_linux_kernel/", 5),
-            ("litebox_platform_linux_userland/", 5),
-            ("litebox_platform_lvbs/", 20),
+            ("litebox_platform_linux_userland/", 6),
+            ("litebox_platform_lvbs/", 21),
             ("litebox_platform_multiplex/", 1),
             ("litebox_platform_windows_userland/", 7),
             ("litebox_runner_linux_userland/", 1),

litebox/src/platform/mod.rs

@@ -670,3 +670,38 @@ pub trait CrngProvider {
     /// failures.
     fn fill_bytes_crng(&self, buf: &mut [u8]);
 }
+
+/// A provider of the Unique Platform Key (UPK).
+///
+/// The UPK is a platform-wide secret used to derive security-critical keys
+/// (e.g., TA unique keys, secure storage keys). It serves as the root of trust
+/// for cryptographic operations within the trusted execution environment.
+///
+/// Trusted hardware devices provide similar funtionalities under various names:
+/// - OP-TEE's Hardware Unique Key (HUK)
+/// - DICE's Unique Device Secret (UDS)
+/// - TPM 2.0's Primary Seeds (e.g., Storage Primary Seed)
+///
+/// Ideally, the platform features a persistent, hardware-backed key (e.g., fused
+/// OTP, PUF-derived). The key should be unique per device, inaccessible outside
+/// the trusted environment, and never directly exposed—only used to derive other
+/// keys.
+///
+/// If hardware support is unavailable, the platform may generate an ephemeral key
+/// from a boot nonce and CRNG, valid only for the current boot session.
+pub trait UniquePlatformKeyProvider {
+    /// Returns a reference to the Unique Platform Key.
+    ///
+    /// # Errors
+    ///
+    /// Returns [`UniquePlatformKeyError`] if the UPK is not supported or not
+    /// initialized.
+    fn unique_platform_key(&self) -> Result<&[u8], UniquePlatformKeyError> {
+        Err(UniquePlatformKeyError)
+    }
+}
+
+/// Error returned when unique platform key is not supported on the platform.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
+#[error("unique platform key is not supported on this platform")]
+pub struct UniquePlatformKeyError;

litebox_common_optee/src/lib.rs

@@ -447,6 +447,11 @@ pub enum TeeParamType {
 impl UteeParams {
     pub const TEE_NUM_PARAMS: usize = TEE_NUM_PARAMS;
 
+    /// Returns `true` if every parameter slot matches the expected type.
+    pub fn has_types(&self, expected: [TeeParamType; 4]) -> bool {
+        (0..TEE_NUM_PARAMS).all(|i| self.get_type(i).is_ok_and(|t| t == expected[i]))
+    }
+
     pub fn get_type(&self, index: usize) -> Result<TeeParamType, Errno> {
         let type_byte = match index {
             0 => self.types.type_0(),
@@ -596,6 +601,22 @@ impl TeeUuid {
         bytes[8..16].copy_from_slice(&data[1].to_be_bytes());
         Self::from_bytes(bytes)
     }
+
+    /// Converts the UUID to a 16-byte array with little-endian encoding for numeric fields.
+    ///
+    /// The byte layout is:
+    /// - bytes[0..4]: `time_low` (little-endian u32)
+    /// - bytes[4..6]: `time_mid` (little-endian u16)
+    /// - bytes[6..8]: `time_hi_and_version` (little-endian u16)
+    /// - bytes[8..16]: `clock_seq_and_node` (8 bytes, direct copy)
+    pub fn to_le_bytes(self) -> [u8; 16] {
+        let mut bytes = [0u8; 16];
+        bytes[0..4].copy_from_slice(&self.time_low.to_le_bytes());
+        bytes[4..6].copy_from_slice(&self.time_mid.to_le_bytes());
+        bytes[6..8].copy_from_slice(&self.time_hi_and_version.to_le_bytes());
+        bytes[8..16].copy_from_slice(&self.clock_seq_and_node);
+        bytes
+    }
 }
 
 /// `TEE_Identity` from `optee_os/lib/libutee/include/tee_api_types.h`.
@@ -1701,6 +1722,26 @@ impl From<OpteeSmcReturnCode> for litebox_common_linux::errno::Errno {
     }
 }
 
+/// HUK subkey usage identifiers, matching OP-TEE's `enum huk_subkey_usage`.
+#[derive(Clone, Copy)]
+#[repr(u32)]
+pub enum HukSubkeyUsage {
+    /// Secure Storage Key
+    Ssk = 0,
+    /// RPMB key
+    Rpmb = 1,
+    /// TA unique key
+    UniqueTa = 2,
+    /// Die ID
+    DieId = 3,
+}
+
+/// Hardware Unique Key length in bytes.
+pub const HUK_LEN: usize = 32;
+
+/// Maximum length of an HUK subkey in bytes.
+pub const HUK_SUBKEY_MAX_LEN: usize = 32;
+
 #[cfg(test)]
 mod tests {
     use super::*;

litebox_platform_linux_userland/src/lib.rs

@@ -2206,12 +2206,49 @@ unsafe fn interrupt_signal_handler(
     set_signal_return(context, interrupt_callback, 0, 0, 0, 0);
 }
 
+#[cfg(all(target_arch = "x86_64", feature = "optee_syscall"))]
+/// Length of the Unique Platform Key in bytes.
+pub const UPK_LEN: usize = 32;
+
+#[cfg(all(target_arch = "x86_64", feature = "optee_syscall"))]
+static UPK_ONCE: std::sync::OnceLock<[u8; UPK_LEN]> = std::sync::OnceLock::new();
+
+/// Sets the Unique Platform Key (UPK) for this platform.
+///
+/// This should be called once during platform initialization with a key derived
+/// from hardware or a boot nonce.
+///
+/// # Panics
+/// Panics if `key` length does not match `UPK_LEN`.
+#[cfg(all(target_arch = "x86_64", feature = "optee_syscall"))]
+pub fn set_unique_platform_key(key: &[u8]) {
+    assert_eq!(key.len(), UPK_LEN, "Unique Platform Key length mismatch");
+    UPK_ONCE.get_or_init(|| {
+        let mut upk = [0u8; UPK_LEN];
+        upk.copy_from_slice(key);
+        upk
+    });
+}
+
 impl litebox::platform::CrngProvider for LinuxUserland {
     fn fill_bytes_crng(&self, buf: &mut [u8]) {
         getrandom::fill(buf).expect("getrandom failed");
     }
 }
 
+#[cfg(all(target_arch = "x86_64", feature = "optee_syscall"))]
+impl litebox::platform::UniquePlatformKeyProvider for LinuxUserland {
+    fn unique_platform_key(&self) -> Result<&[u8], litebox::platform::UniquePlatformKeyError> {
+        UPK_ONCE
+            .get()
+            .map(<[u8; UPK_LEN]>::as_slice)
+            .ok_or(litebox::platform::UniquePlatformKeyError)
+    }
+}
+
+#[cfg(not(all(target_arch = "x86_64", feature = "optee_syscall")))]
+impl litebox::platform::UniquePlatformKeyProvider for LinuxUserland {}
+
 /// Dummy `VmapManager`.
 ///
 /// In general, userland platforms do not support `vmap` and `vunmap` (which are kernel functions).

litebox_platform_lvbs/src/host/lvbs_impl.rs

@@ -96,6 +96,42 @@ impl litebox::platform::CrngProvider for LvbsLinuxKernel {
     }
 }
 
+#[cfg(feature = "optee_syscall")]
+/// Length of the Unique Platform Key in bytes.
+pub const UPK_LEN: usize = 32;
+#[cfg(feature = "optee_syscall")]
+static UPK_ONCE: spin::Once<[u8; UPK_LEN]> = spin::Once::new();
+
+/// Sets the Unique Platform Key (UPK) for this platform.
+///
+/// This should be called once during platform initialization with a key derived
+/// from hardware or a boot nonce.
+///
+/// # Panics
+/// Panics if `key` length does not match `UPK_LEN`.
+#[cfg(feature = "optee_syscall")]
+pub fn set_unique_platform_key(key: &[u8]) {
+    assert_eq!(key.len(), UPK_LEN, "Unique Platform Key length mismatch");
+    UPK_ONCE.call_once(|| {
+        let mut upk = [0u8; UPK_LEN];
+        upk.copy_from_slice(key);
+        upk
+    });
+}
+
+#[cfg(feature = "optee_syscall")]
+impl litebox::platform::UniquePlatformKeyProvider for LvbsLinuxKernel {
+    fn unique_platform_key(&self) -> Result<&[u8], litebox::platform::UniquePlatformKeyError> {
+        UPK_ONCE
+            .get()
+            .map(<[u8; UPK_LEN]>::as_slice)
+            .ok_or(litebox::platform::UniquePlatformKeyError)
+    }
+}
+
+#[cfg(not(feature = "optee_syscall"))]
+impl litebox::platform::UniquePlatformKeyProvider for LvbsLinuxKernel {}
+
 pub struct HostLvbsInterface;
 
 impl HostLvbsInterface {}

litebox_platform_lvbs/src/host/mod.rs

@@ -8,6 +8,8 @@ pub mod lvbs_impl;
 pub mod per_cpu_variables;
 
 pub use lvbs_impl::LvbsLinuxKernel;
+#[cfg(feature = "optee_syscall")]
+pub use lvbs_impl::{UPK_LEN, set_unique_platform_key};
 
 #[cfg(test)]
 pub mod mock;

litebox_platform_lvbs/src/lib.rs

@@ -1296,6 +1296,8 @@ unsafe extern "C" fn switch_to_user(_ctx: &litebox_common_linux::PtRegs) -> ! {
     );
 }
 
+pub use crate::host::{UPK_LEN, set_unique_platform_key};
+
 // Note on user page table management:
 // The legacy platform code creates a new page table to load a program in a separate
 // address space and destroys it when the program terminates. This is why the old syscall

litebox_runner_lvbs/src/lib.rs

@@ -8,6 +8,7 @@ extern crate alloc;
 use core::{ops::Neg, panic::PanicInfo};
 use litebox::{
     mm::linux::PAGE_SIZE,
+    platform::CrngProvider,
     utils::{ReinterpretSignedExt, TruncateExt},
 };
 use litebox_common_linux::errno::Errno;
@@ -100,6 +101,13 @@ pub fn init() -> Option<&'static Platform> {
             );
 
             allocate_per_cpu_variables();
+
+            // Set the Unique Platform Key
+            //
+            // *** This is PoC. Once an HVCI/Heki call is ready, we SHALL set UPK via that function ***
+            let mut upk = [0u8; litebox_platform_lvbs::UPK_LEN];
+            platform.fill_bytes_crng(&mut upk);
+            litebox_platform_lvbs::set_unique_platform_key(&upk);
         } else {
             panic!("Failed to get memory info");
         }

litebox_runner_optee_on_linux_userland/src/lib.rs

@@ -3,6 +3,7 @@
 
 use anyhow::Result;
 use clap::Parser;
+use litebox::platform::CrngProvider;
 use litebox_common_optee::{TeeUuid, UteeEntryFunc, UteeParamOwned};
 use litebox_platform_multiplex::Platform;
 use std::path::PathBuf;
@@ -97,6 +98,11 @@ pub fn run(cli_args: CliArgs) -> Result<()> {
         InterceptionBackend::Rewriter => {}
     }
 
+    // For now, we use a random UPK for this runner. We can get one via command line if needed.
+    let mut upk = [0u8; litebox_platform_linux_userland::UPK_LEN];
+    platform.fill_bytes_crng(&mut upk);
+    litebox_platform_linux_userland::set_unique_platform_key(&upk);
+
     if cli_args.command_sequence.is_empty() {
         run_ta_with_default_commands(&shim, ldelf_data.as_slice(), prog_data.as_slice());
     } else {

litebox_shim_optee/Cargo.toml

@@ -6,6 +6,9 @@ edition = "2024"
 [dependencies]
 aes = { version = "0.7", default-features = false }
 arrayvec = { version = "0.7.6", default-features = false }
+hmac = { version = "0.12", default-features = false }
+sha2 = { version = "0.10", default-features = false }
+zeroize = { version = "1.8", default-features = false, features = ["alloc"] }
 bitflags = "2.9.0"
 cfg-if = "1.0.0"
 ctr = { version = "0.8", default-features = false }
@@ -37,4 +40,5 @@ platform_lvbs = ["litebox_platform_multiplex/platform_lvbs_with_optee_syscall"]
 workspace = true
 
 [dev-dependencies]
+litebox_platform_linux_userland = { path = "../litebox_platform_linux_userland/", version = "0.1.0", features = ["optee_syscall"] }
 litebox_platform_multiplex = { path = "../litebox_platform_multiplex/", version = "0.1.0", default-features = false, features = ["platform_linux_userland_with_optee_syscall"] }

litebox_shim_optee/src/lib.rs

@@ -205,6 +205,7 @@ impl OpteeShim {
                 thread: ThreadState::new(),
                 session_id: self.0.session_id_pool.allocate(),
                 ta_app_id: ta_uuid,
+                ta_svn: 0, // TODO: Initialize from TA binary or OP-TEE Secure Storage
                 client_identity: client.unwrap_or(TeeIdentity {
                     login: TeeLogin::User,
                     uuid: TeeUuid::default(),
@@ -1141,6 +1142,8 @@ struct Task {
     session_id: u32,
     /// TA UUID
     ta_app_id: TeeUuid,
+    /// TA Secure Version Number (SVN)
+    ta_svn: u32,
     /// Client identity (VTL0 process or another TA)
     client_identity: TeeIdentity,
     /// TEE cryptography state map
@@ -1258,6 +1261,7 @@ mod test_utils {
                 thread: ThreadState::new(),
                 session_id: self.session_id_pool.allocate(),
                 ta_app_id: TeeUuid::default(),
+                ta_svn: 0,
                 client_identity: TeeIdentity {
                     login: TeeLogin::User,
                     uuid: TeeUuid::default(),