Verify TAs

litebox_common_optee/Cargo.toml

@@ -9,7 +9,12 @@ elf = { version = "0.8.0", default-features = false }
 litebox = { path = "../litebox/", version = "0.1.0" }
 litebox_common_linux = { path = "../litebox_common_linux/", version = "0.1.0" }
 num_enum = { version = "0.7.3", default-features = false }
+rsa = { version = "0.9.10", default-features = false, features = ["pem"], optional = true }
+sha2 = { version = "0.10.9", default-features = false, features = ["oid"], optional = true }
 zerocopy = { version = "0.8", features = ["derive"] }
 
+[features]
+signed-ta-rsa = ["dep:rsa", "dep:sha2"]
+
 [lints]
 workspace = true

litebox_common_optee/src/lib.rs

@@ -9,6 +9,8 @@
 extern crate alloc;
 
 use alloc::boxed::Box;
+#[cfg(feature = "signed-ta-rsa")]
+use alloc::vec::Vec;
 use core::mem::size_of;
 use litebox::platform::RawConstPointer as _;
 use litebox::utils::TruncateExt;
@@ -999,6 +1001,7 @@ impl TeeMemoryAccessRights {
 
 const TEE_ALG_AES_CTR: u32 = 0x1000_0210;
 const TEE_ALG_AES_GCM: u32 = 0x4000_0810;
+const TEE_ALG_RSASSA_PKCS1_PSS_MGF1_SHA256: u32 = 0x7041_4930;
 const TEE_ALG_RSASSA_PKCS1_V1_5_SHA256: u32 = 0x7000_4830;
 const TEE_ALG_RSASSA_PKCS1_V1_5_SHA512: u32 = 0x7000_6830;
 const TEE_ALG_HMAC_SHA256: u32 = 0x3000_0004;
@@ -1014,6 +1017,7 @@ const TEE_ALG_ILLEGAL_VALUE: u32 = 0xefff_ffff;
 pub enum TeeAlgorithm {
     AesCtr = TEE_ALG_AES_CTR,
     AesGcm = TEE_ALG_AES_GCM,
+    RsaPssSha256 = TEE_ALG_RSASSA_PKCS1_PSS_MGF1_SHA256,
     RsaPkcs1Sha256 = TEE_ALG_RSASSA_PKCS1_V1_5_SHA256,
     RsaPkcs1Sha512 = TEE_ALG_RSASSA_PKCS1_V1_5_SHA512,
     HmacSha256 = TEE_ALG_HMAC_SHA256,
@@ -1055,9 +1059,9 @@ impl From<TeeAlgorithm> for TeeAlgorithmClass {
         match algo {
             TeeAlgorithm::AesCtr | TeeAlgorithm::AesGcm => TeeAlgorithmClass::Cipher,
             TeeAlgorithm::HmacSha256 | TeeAlgorithm::HmacSha512 => TeeAlgorithmClass::Mac,
-            TeeAlgorithm::RsaPkcs1Sha256 | TeeAlgorithm::RsaPkcs1Sha512 => {
-                TeeAlgorithmClass::AsymmetricSignature
-            }
+            TeeAlgorithm::RsaPkcs1Sha256
+            | TeeAlgorithm::RsaPkcs1Sha512
+            | TeeAlgorithm::RsaPssSha256 => TeeAlgorithmClass::AsymmetricSignature,
             _ => TeeAlgorithmClass::Unknown,
         }
     }
@@ -2646,3 +2650,133 @@ mod tests {
         assert!(OpteeRpcArgs::from_header_and_raw_params(&header, &[]).is_err());
     }
 }
+
+/// `Shdr` from `optee_os/core/include/signed_hdr.h`
+#[cfg(feature = "signed-ta-rsa")]
+#[derive(Clone, Copy, FromBytes, IntoBytes, Immutable, KnownLayout)]
+#[repr(C)]
+pub struct Shdr {
+    pub magic: u32,
+    pub img_type: u32,
+    pub img_size: u32,
+    pub algo: u32,
+    pub hash_size: u16,
+    pub sig_size: u16,
+}
+
+/// `SHDR_MAGIC` from `optee_os/core/include/signed_hdr.h`
+#[cfg(feature = "signed-ta-rsa")]
+const SHDR_MAGIC: u32 = 0x4f54_5348;
+
+/// From `optee_os/core/include/signed_hdr.h`
+/// struct shdr_bootstrap_ta {
+///     uint8_t uuid[sizeof(TEE_UUID)];
+///     uint32_t ta_version;
+/// };
+/// and from `optee_os/core/include/tee_api_types.h`
+/// typedef struct {
+///     uint32_t timeLow;
+///     uint16_t timeMid;
+///     uint16_t timeHiAndVersion;
+///     uint8_t clockSeqAndNode[8];
+/// } TEE_UUID;
+#[cfg(feature = "signed-ta-rsa")]
+const SHDR_UUID_LEN: usize = 16;
+#[cfg(feature = "signed-ta-rsa")]
+const SHDR_VERSION_LEN: usize = 4;
+
+/// An RSA public key used to verify signed `.ta` files
+#[cfg(feature = "signed-ta-rsa")]
+pub struct TaVerifyKey(rsa::RsaPublicKey);
+
+#[cfg(feature = "signed-ta-rsa")]
+impl TaVerifyKey {
+    pub fn from_der(der: &[u8]) -> Result<Self, &'static str> {
+        use rsa::pkcs8::DecodePublicKey;
+
+        rsa::RsaPublicKey::from_public_key_der(der)
+            .map(TaVerifyKey)
+            .map_err(|_| "Invalid RSA public key DER")
+    }
+}
+
+/// Parse and verify an OP-TEE signed .ta file.
+/// A signed `.ta` file has the following layout:
+/// shdr || hash || sig || uuid || ta_version || img
+/// Return the parsed TaHead and the TA ELF binary.
+#[cfg(feature = "signed-ta-rsa")]
+pub fn parse_and_verify_ta<'a>(
+    ta_data: &'a [u8],
+    verify_key: &TaVerifyKey,
+) -> Result<(TaHead, &'a [u8]), &'static str> {
+    let hdr_size = size_of::<Shdr>();
+    if ta_data.len() < hdr_size {
+        return Err("Invalid signed TA file");
+    }
+    let shdr =
+        Shdr::read_from_bytes(&ta_data[..hdr_size]).map_err(|_| "Invalid signed TA header")?;
+    if shdr.magic != SHDR_MAGIC {
+        return Err("Invalid TA magic");
+    }
+    let hash_size = shdr.hash_size as usize;
+    let sig_size = shdr.sig_size as usize;
+    let img_size = shdr.img_size as usize;
+
+    let hash_offset = hdr_size;
+    let sig_offset = hash_offset
+        .checked_add(hash_size)
+        .ok_or("Invalid signed TA file")?;
+    let uuid_offset = sig_offset
+        .checked_add(sig_size)
+        .ok_or("Invalid signed TA file")?;
+    let version_offset = uuid_offset
+        .checked_add(SHDR_UUID_LEN)
+        .ok_or("Invalid signed TA file")?;
+    let img_offset = version_offset
+        .checked_add(SHDR_VERSION_LEN)
+        .ok_or("Invalid signed TA file")?;
+    let end = img_offset
+        .checked_add(img_size)
+        .ok_or("Invalid signed TA file")?;
+    if end > ta_data.len() {
+        return Err("Invalid signed TA file");
+    }
+    let sig = &ta_data[sig_offset..uuid_offset];
+    let uuid_and_version = &ta_data[uuid_offset..img_offset];
+    let img = &ta_data[img_offset..end];
+
+    // The signed message is shdr || uuid || ta_version || img.
+    let mut signed_message = Vec::with_capacity(hdr_size + uuid_and_version.len() + img.len());
+    signed_message.extend_from_slice(&ta_data[..hdr_size]);
+    signed_message.extend_from_slice(uuid_and_version);
+    signed_message.extend_from_slice(img);
+    verify_shdr_signature(&signed_message, sig, shdr.algo, &verify_key.0)?;
+    let ta_head = parse_ta_head(img).ok_or("Invalid TA ELF binary")?;
+
+    Ok((ta_head, img))
+}
+
+/// Verify a signed `.ta` file's signature against the given RSA public key.
+/// TODO: Support more signature algorithms if needed.
+#[cfg(feature = "signed-ta-rsa")]
+fn verify_shdr_signature(
+    message: &[u8],
+    signature: &[u8],
+    algo: u32,
+    rsa_pub_key: &rsa::RsaPublicKey,
+) -> Result<(), &'static str> {
+    use rsa::signature::Verifier;
+    use sha2::Sha256;
+
+    match algo {
+        TEE_ALG_RSASSA_PKCS1_PSS_MGF1_SHA256 => {
+            let verifying_key = rsa::pss::VerifyingKey::<Sha256>::new(rsa_pub_key.clone());
+            let sig =
+                rsa::pss::Signature::try_from(signature).map_err(|_| "Invalid PSS signature")?;
+            verifying_key
+                .verify(message, &sig)
+                .map_err(|_| "Signature verification failed")
+        }
+        _ => Err("Unsupported TA signature algorithm"),
+    }
+}

litebox_runner_lvbs/Cargo.toml

@@ -23,6 +23,7 @@ x86_64 = { version = "0.15.2", default-features = false, features = ["instructio
 [features]
 devbox = ["litebox_platform_lvbs/devbox"]
 preemption_test_quantum = ["litebox_platform_lvbs/preemption_test_quantum"]
+signed-ta-rsa = ["litebox_shim_optee/signed-ta-rsa"]
 
 [lints]
 workspace = true

litebox_runner_lvbs/src/lib.rs

@@ -824,13 +824,14 @@ fn open_session_new_instance(
     ta_req_info: &litebox_shim_optee::msg_handler::TaRequestInfo<PAGE_SIZE>,
 ) -> Result<(), OpteeSmcReturnCode> {
     let shim = litebox_shim_optee::OpteeShimBuilder::new(platform, session_manager()).build();
-    if shim.get_ta_bin(&ta_uuid).is_none() {
+    let Some(ta_source) = shim.get_ta_source(&ta_uuid) else {
         msg_args.session = 0;
         msg_args.ret = TeeResult::ItemNotFound;
         msg_args.ret_origin = TeeOrigin::Tee;
         write_non_ta_msg_args_to_normal_world(platform, msg_args, msg_args_phys_addr)?;
         return Ok(());
-    }
+    };
+    debug_serial_println!("Loading TA: uuid={:?}, source={:?}", ta_uuid, ta_source);
 
     // Token is declared before `task_pt_guard` so it drops AFTER it.
     // Marker only releases once CR3 is back to base. See
@@ -1414,15 +1415,12 @@ const LDELF_BINARY: &[u8] = &[0u8; 0];
 const TA_BINARY: &[u8] = &[0u8; 0];
 const TA_BINARIES: &[&[u8]] = &[TA_BINARY];
 
-/// Register a TA binary embedded in the runner image.
+/// Register a raw or signed TA binary embedded in the runner image.
 fn register_embedded_ta(
     shim: &litebox_shim_optee::OpteeShim<Platform>,
     ta_binary: &'static [u8],
 ) -> bool {
-    let Some(ta_head) = litebox_common_optee::parse_ta_head(ta_binary) else {
-        return false;
-    };
-    shim.store_ta_bin(&ta_head.uuid, ta_binary)
+    shim.store_embedded_ta(ta_binary)
 }
 
 /// Register all TA binaries embedded in the runner image.

litebox_runner_optee_on_linux_userland/src/lib.rs

@@ -114,7 +114,7 @@ fn run_ta_with_default_commands(
     let ta_uuid = litebox_common_optee::parse_ta_head(ta_bin)
         .expect("Failed to parse TA header from ta_bin")
         .uuid;
-    assert!(shim.store_ta_bin(&ta_uuid, ta_bin));
+    assert!(shim.store_ta_bin(&ta_uuid, ta_bin, litebox_shim_optee::TaSource::BuiltIn,));
     for func_id in [UteeEntryFunc::OpenSession, UteeEntryFunc::CloseSession] {
         let params = [const { UteeParamOwned::None }; UteeParamOwned::TEE_NUM_PARAMS];
 

litebox_runner_optee_on_linux_userland/src/tests.rs

@@ -29,7 +29,7 @@ pub fn run_ta_with_test_commands(
     };
     let ta_head =
         litebox_common_optee::parse_ta_head(ta_bin).expect("Failed to parse TA header from ta_bin");
-    assert!(shim.store_ta_bin(&ta_head.uuid, ta_bin));
+    assert!(shim.store_ta_bin(&ta_head.uuid, ta_bin, litebox_shim_optee::TaSource::BuiltIn,));
     let mut ta_info: Option<LoadedProgram<Platform>> = None;
     // The active session id for the TA. Set at OpenSession and reused for the
     // subsequent InvokeCommand entries on the same persistent session.

litebox_shim_optee/Cargo.toml

@@ -23,6 +23,9 @@ zerocopy = { version = "0.8", default-features = false, features = ["derive"] }
 zeroize = { version = "1.8", default-features = false, features = ["alloc"] }
 p384 = { version = "0.13.1", default-features = false, features = ["arithmetic", "ecdsa"] }
 
+[features]
+signed-ta-rsa = ["litebox_common_optee/signed-ta-rsa"]
+
 [lints]
 workspace = true