OP-TEE TA unique key derivation

litebox_common_optee/src/lib.rs

@@ -482,6 +482,11 @@ pub enum TeeParamType {
 impl UteeParams {
     pub const TEE_NUM_PARAMS: usize = TEE_NUM_PARAMS;
 
+    /// Return `true` if every parameter matches the expected type.
+    pub fn has_types(&self, expected: [TeeParamType; Self::TEE_NUM_PARAMS]) -> bool {
+        (0..Self::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(),
@@ -619,6 +624,16 @@ impl TeeUuid {
         bytes[8..16].copy_from_slice(&data[1].to_le_bytes());
         Self::from_bytes(bytes)
     }
+
+    /// Converts the UUID to a 16-byte array with little-endian encoding.
+    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
+    }
 }
 
 /// TA flags from `optee_os/lib/libutee/include/user_ta_header.h`.
@@ -2306,6 +2321,27 @@ pub fn parse_ta_head(elf_data: &[u8]) -> Option<TaHead> {
     None
 }
 
+/// Hardware Unique Key (HUK) subkey usage identifiers based on OP-TEE's `enum huk_subkey_usage`.
+#[derive(Clone, Copy)]
+#[repr(u32)]
+pub enum HukSubkeyUsage {
+    /// RPMB key
+    Rpmb = 0,
+    /// Secure Storage Key
+    Ssk = 1,
+    /// Die ID
+    DieId = 2,
+    /// TA unique key
+    UniqueTa = 3,
+    /// TA encryption key
+    TaEnc = 4,
+    /// SCP03 set of encryption keys
+    Se050 = 5,
+}
+
+/// Maximum length of an HUK subkey in bytes.
+pub const HUK_SUBKEY_MAX_LEN: usize = 32;
+
 #[cfg(test)]
 mod tests {
     use super::*;

litebox_platform_lvbs/src/host/lvbs_impl.rs

@@ -136,6 +136,25 @@ pub(crate) fn set_platform_root_key(key: &[u8]) {
     });
 }
 
+impl litebox::platform::DerivedKeyProvider for LvbsLinuxKernel {
+    fn derive_key<E>(
+        &self,
+        shim_kdf: Option<fn(&[u8], litebox::platform::KDFParams) -> Result<(), E>>,
+        params: litebox::platform::KDFParams,
+    ) -> Result<(), litebox::platform::DerivedKeyError<E>> {
+        let Some(prk) = PRK_ONCE.get() else {
+            return Err(litebox::platform::DerivedKeyError::UnsupportedRebootPersistentKey);
+        };
+        match shim_kdf {
+            None => {
+                // LVBS platform doesn't have its own KDF implementation.
+                Err(litebox::platform::DerivedKeyError::ShimKDFRequired)
+            }
+            Some(shim_kdf) => Ok(shim_kdf(prk, params)?),
+        }
+    }
+}
+
 pub struct HostLvbsInterface;
 
 impl HostLvbsInterface {}

litebox_runner_optee_on_linux_userland/src/lib.rs

@@ -85,6 +85,8 @@ pub fn run(cli_args: CliArgs) -> Result<()> {
     let _litebox = shim_builder.litebox();
     let shim = shim_builder.build();
 
+    platform.initialize_boot_specific_kdf_support();
+
     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

@@ -14,11 +14,14 @@ litebox_common_linux = { path = "../litebox_common_linux/", version = "0.1.0" }
 litebox_common_optee = { path = "../litebox_common_optee/", version = "0.1.0" }
 litebox_platform_multiplex = { path = "../litebox_platform_multiplex/", version = "0.1.0", default-features = false }
 litebox_util_log = { version = "0.1.0", path = "../litebox_util_log" }
+hmac = { version = "0.12", default-features = false }
 num_enum = { version = "0.7.3", default-features = false }
 once_cell = { version = "1.20.2", default-features = false, features = ["alloc", "race"] }
+sha2 = { version = "0.10", default-features = false }
 spin = { version = "0.10.0", default-features = false, features = ["spin_mutex", "once"] }
 thiserror = { version = "2.0.6", default-features = false }
 zerocopy = { version = "0.8", default-features = false, features = ["derive"] }
+zeroize = { version = "1.8", default-features = false, features = ["alloc"] }
 
 [features]
 default = ["platform_lvbs"]

litebox_shim_optee/src/syscalls/pta.rs

@@ -5,12 +5,19 @@
 //! the functions of built-in TAs.
 
 use crate::{Task, UserConstPtr, UserMutPtr};
-use litebox::{
-    platform::{RawConstPointer as _, RawMutPointer as _},
-    utils::TruncateExt,
+use alloc::vec;
+use alloc::vec::Vec;
+use hmac::{Hmac, Mac};
+use litebox::platform::{
+    DerivedKeyError, DerivedKeyProvider, KDFParams, RawConstPointer as _, RawMutPointer as _,
+};
+use litebox::utils::TruncateExt;
+use litebox_common_optee::{
+    HUK_SUBKEY_MAX_LEN, HukSubkeyUsage, TeeParamType, TeeResult, TeeUuid, UteeParams,
 };
-use litebox_common_optee::{TeeParamType, TeeResult, TeeUuid, UteeParams};
 use num_enum::TryFromPrimitive;
+use sha2::Sha256;
+use zeroize::{Zeroize, Zeroizing};
 
 pub const PTA_SYSTEM_UUID: TeeUuid = TeeUuid {
     time_low: 0x3a2f_8978,
@@ -34,6 +41,13 @@ const PTA_SYSTEM_DLSYM: u32 = 11;
 const PTA_SYSTEM_GET_TPM_EVENT_LOG: u32 = 12;
 const PTA_SYSTEM_SUPP_PLUGIN_INVOKE: u32 = 13;
 
+/// Minimum size of a derived key in bytes.
+const TA_DERIVED_KEY_MIN_SIZE: usize = 16;
+/// Maximum size of a derived key in bytes.
+const TA_DERIVED_KEY_MAX_SIZE: usize = 32;
+/// Maximum size of extra data for key derivation in bytes.
+const TA_DERIVED_EXTRA_DATA_MAX_SIZE: usize = 1024;
+
 /// `PTA_SYSTEM_*` command ID from `optee_os/lib/libutee/include/pta_system.h`
 #[derive(Clone, Copy, TryFromPrimitive)]
 #[repr(u32)]
@@ -72,6 +86,8 @@ pub fn is_pta_session(ta_sess_id: u32) -> bool {
     ta_sess_id == crate::SessionIdPool::get_pta_session_id()
 }
 
+type HmacSha256 = Hmac<Sha256>;
+
 impl Task {
     /// Handle a command of the system PTA.
     pub fn handle_system_pta_command(
@@ -81,58 +97,7 @@ impl Task {
     ) -> Result<(), TeeResult> {
         #[allow(clippy::single_match_else)]
         match PtaSystemCommandId::try_from(cmd_id).map_err(|_| TeeResult::BadParameters)? {
-            PtaSystemCommandId::DeriveTaUniqueKey => {
-                if params
-                    .get_type(0)
-                    .is_ok_and(|t| t == TeeParamType::MemrefInput)
-                    && params
-                        .get_type(1)
-                        .is_ok_and(|t| t == TeeParamType::MemrefOutput)
-                    && params.get_type(2).is_ok_and(|t| t == TeeParamType::None)
-                    && params.get_type(3).is_ok_and(|t| t == TeeParamType::None)
-                    && let Ok(Some(input)) =
-                        params.get_values(0).map_err(|_| TeeResult::BadParameters)
-                    && let Ok(Some(output)) =
-                        params.get_values(1).map_err(|_| TeeResult::BadParameters)
-                {
-                    // TODO: revisit buffer size limits based on OP-TEE spec and deployment constraints
-                    let input_len =
-                        usize::try_from(input.1).map_err(|_| TeeResult::BadParameters)?;
-                    if input_len > crate::MAX_KERNEL_BUF_SIZE {
-                        return Err(TeeResult::BadParameters);
-                    }
-                    let input_addr: usize = input.0.truncate();
-                    let input_ptr = UserConstPtr::<u8>::from_usize(input_addr);
-                    let _extra_data = input_ptr
-                        .to_owned_slice(input_len)
-                        .ok_or(TeeResult::BadParameters)?;
-
-                    let output_len =
-                        usize::try_from(output.1).map_err(|_| TeeResult::BadParameters)?;
-                    if output_len > crate::MAX_KERNEL_BUF_SIZE {
-                        return Err(TeeResult::BadParameters);
-                    }
-                    let output_addr: usize = output.0.truncate();
-                    let output_ptr = UserMutPtr::<u8>::from_usize(output_addr);
-
-                    // TODO: derive a TA unique key using the hardware unique key (HUK), TA's UUID, and `extra_data`
-                    litebox_util_log::debug!(
-                        ptr:% = format_args!("{:#x}", output_addr),
-                        size:% = output_len;
-                        "derive key into secure memory"
-                    );
-                    // TODO: replace below with a secure key derivation function
-                    let mut key_buf = alloc::vec![0u8; output_len];
-                    self.sys_cryp_random_number_generate(&mut key_buf)?;
-                    output_ptr
-                        .copy_from_slice(0, &key_buf)
-                        .ok_or(TeeResult::BadParameters)?;
-
-                    Ok(())
-                } else {
-                    Err(TeeResult::BadParameters)
-                }
-            }
+            PtaSystemCommandId::DeriveTaUniqueKey => self.derive_ta_unique_key(params),
             _ => {
                 #[cfg(debug_assertions)]
                 todo!("support other system PTA commands {cmd_id}");
@@ -141,4 +106,116 @@ impl Task {
             }
         }
     }
+
+    /// Derives a unique key for a TA using HUK.
+    ///
+    /// This follows the OP-TEE `system_derive_ta_unique_key` implementation from
+    /// `core/pta/system.c`.
+    fn derive_ta_unique_key(&self, params: &UteeParams) -> Result<(), TeeResult> {
+        use TeeParamType::{MemrefInput, MemrefOutput, None};
+
+        if !params.has_types([MemrefInput, MemrefOutput, None, None]) {
+            return Err(TeeResult::BadParameters);
+        }
+
+        let (extra_data_addr, extra_data_size_u64) = params
+            .get_values(0)
+            .map_err(|_| TeeResult::BadParameters)?
+            .ok_or(TeeResult::BadParameters)?;
+        let extra_data_size: usize = extra_data_size_u64.truncate();
+
+        let (subkey_addr, subkey_size_u64) = params
+            .get_values(1)
+            .map_err(|_| TeeResult::BadParameters)?
+            .ok_or(TeeResult::BadParameters)?;
+        let subkey_size: usize = subkey_size_u64.truncate();
+
+        if extra_data_size > TA_DERIVED_EXTRA_DATA_MAX_SIZE
+            || !(TA_DERIVED_KEY_MIN_SIZE..=TA_DERIVED_KEY_MAX_SIZE).contains(&subkey_size)
+            || subkey_addr == 0
+        {
+            return Err(TeeResult::BadParameters);
+        }
+
+        let extra_data = if extra_data_size == 0 {
+            Vec::new().into_boxed_slice()
+        } else {
+            let extra_data_ptr = UserConstPtr::<u8>::from_usize(extra_data_addr.truncate());
+            extra_data_ptr
+                .to_owned_slice(extra_data_size)
+                .ok_or(TeeResult::BadParameters)?
+        };
+
+        // In LiteBox for LVBS, TA's address space is always in the secure world.
+        let subkey_ptr = UserMutPtr::<u8>::from_usize(subkey_addr.truncate());
+
+        // subkey = KDF(huk, usage || ta_uuid || extra_data)
+        let ta_uuid_bytes = self.ta_app_id.to_le_bytes();
+        let mut subkey_buf = Zeroizing::new(vec![0u8; subkey_size]);
+        self.huk_subkey_derive(
+            HukSubkeyUsage::UniqueTa,
+            &[&ta_uuid_bytes, &extra_data],
+            &mut subkey_buf,
+        )
+        .and_then(|()| {
+            subkey_ptr
+                .copy_from_slice(0, &subkey_buf)
+                .ok_or(TeeResult::AccessDenied)
+        })
+    }
+
+    /// Derive a subkey using HUK and constant data.
+    ///
+    /// This follows the OP-TEE `huk_subkey_derive` interface from `core/kernel/huk_subkey.c`.
+    fn huk_subkey_derive(
+        &self,
+        usage: HukSubkeyUsage,
+        const_data: &[&[u8]],
+        subkey: &mut [u8],
+    ) -> Result<(), TeeResult> {
+        let subkey_len = subkey.len();
+        if subkey_len > HUK_SUBKEY_MAX_LEN {
+            return Err(TeeResult::BadParameters);
+        }
+
+        let kdf_context_len =
+            core::mem::size_of::<u32>() + const_data.iter().map(|chunk| chunk.len()).sum::<usize>();
+        let mut kdf_context = Zeroizing::new(Vec::with_capacity(kdf_context_len));
+        kdf_context.extend_from_slice(&(usage as u32).to_le_bytes());
+        for chunk in const_data {
+            kdf_context.extend_from_slice(chunk);
+        }
+        let kdf_params = KDFParams {
+            context: kdf_context.as_slice(),
+            output: subkey,
+        };
+
+        self.global
+            .platform
+            .derive_key(Some(huk_subkey_derive_inner), kdf_params)
+            .map_err(|err| match err {
+                DerivedKeyError::ShimKDFRequired
+                | DerivedKeyError::UnsupportedRebootPersistentKey => TeeResult::NotSupported,
+                DerivedKeyError::ShimKDFError(err) => err,
+            })?;
+
+        Ok(())
+    }
+}
+
+/// A KDF callback that derives a subkey from `huk` and `params.context` to be passed to
+/// the underlying platform implementation of `derive_key`.
+fn huk_subkey_derive_inner(huk: &[u8], params: KDFParams<'_>) -> Result<(), TeeResult> {
+    let subkey_len = params.output.len();
+    if subkey_len > HUK_SUBKEY_MAX_LEN {
+        return Err(TeeResult::BadParameters);
+    }
+
+    let mut hmac = HmacSha256::new_from_slice(huk).map_err(|_| TeeResult::BadParameters)?;
+    hmac.update(params.context);
+
+    let mut hmac_bytes = hmac.finalize().into_bytes();
+    params.output.copy_from_slice(&hmac_bytes[..subkey_len]);
+    hmac_bytes.zeroize();
+    Ok(())
 }