From 87cbadf0a6d5ade5824d9c146633a5229bbb678a Mon Sep 17 00:00:00 2001
From: Jayant Gandhi <jayg@microsoft.com>
Date: Tue, 30 Jun 2026 17:39:30 +0000
Subject: [PATCH 1/4] feat(fw): implement GetUnwrappingKey DDI handler
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Within an open session, return the partition's RSA-2048 unwrapping
public key (raw wire `n_le || e_le`) plus its vault key id, for the
host to RSA-AES key-wrap a payload for RsaUnwrap.

- Lazy RSA-2048 keygen via a new HsmPartitionManager::
  part_ensure_unwrapping_key: the std PAL generates on first use
  (keygen is too slow to run at partition enable); hardware PALs
  generate in the background from part-init and report
  PendingKeyGeneration until ready, prompting the host to retry.
- Derive the public key from the vault-stored private key on demand
  (no separate pub cached, matching the reference firmware) via a new
  HsmRsa::rsa_priv_pub_key that converts the PAL's vault representation
  (raw on hardware, DER in the std/OpenSSL PAL) into wire form; the
  std driver owns the big-endian->little-endian flip.
- Serialize the public key straight into the response frame's reserved
  slot (query/alloc/use) — no intermediate buffer or copy.
- Centralize the BE<->LE reverse_copy helper in drivers/mod.rs, shared
  by the ecc and rsa drivers.
- Fix ras_gen_keypair -> rsa_gen_keypair typo.

masked_key is emitted empty for now (firmware-side masking pending the
unmask path). Uno PAL methods are stubbed (PendingKeyGeneration /
UnsupportedCmd).

Adds get_unwrapping_key smoke tests (happy path, stable-across-calls,
no-session). Full smoke 53/53; std PAL unit 157/157.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 ddi/mbor/types/tests/azihsm_ddi_tests.rs      |   1 +
 .../integration/get_unwrapping_key_smoke.rs   | 106 ++++++++++++++++++
 .../ddi/mbor/types/src/get_unwrapping_key.rs  |   2 +-
 .../lib/src/ddi/mbor/get_unwrapping_key.rs    |  75 +++++++++++++
 fw/core/lib/src/ddi/mbor/mod.rs               |   3 +
 fw/pal/traits/src/crypto/rsa.rs               |  35 +++++-
 fw/pal/traits/src/part.rs                     |  26 +++++
 fw/plat/std/pal/src/drivers/ecc.rs            |  24 +---
 fw/plat/std/pal/src/drivers/mod.rs            |  18 +++
 fw/plat/std/pal/src/drivers/rsa.rs            |  46 ++++++++
 fw/plat/std/pal/src/part.rs                   |  36 ++++++
 fw/plat/std/pal/src/rsa.rs                    |  19 +++-
 fw/plat/uno/fw/pal/src/crypto/rsa.rs          |  13 ++-
 fw/plat/uno/fw/pal/src/part.rs                |   9 ++
 14 files changed, 386 insertions(+), 27 deletions(-)
 create mode 100644 ddi/mbor/types/tests/integration/get_unwrapping_key_smoke.rs
 create mode 100644 fw/core/lib/src/ddi/mbor/get_unwrapping_key.rs

diff --git a/ddi/mbor/types/tests/azihsm_ddi_tests.rs b/ddi/mbor/types/tests/azihsm_ddi_tests.rs
index 84565e3a9..d568db90a 100644
--- a/ddi/mbor/types/tests/azihsm_ddi_tests.rs
+++ b/ddi/mbor/types/tests/azihsm_ddi_tests.rs
@@ -42,6 +42,7 @@ mod integration {
     pub mod get_session_encryption_key;
     pub mod get_session_encryption_key_smoke;
     pub mod get_unwrapping_key;
+    pub mod get_unwrapping_key_smoke;
     pub mod hkdf_smoke;
     pub mod hmac;
     pub mod hmac_smoke;
diff --git a/ddi/mbor/types/tests/integration/get_unwrapping_key_smoke.rs b/ddi/mbor/types/tests/integration/get_unwrapping_key_smoke.rs
new file mode 100644
index 000000000..6111f4ce4
--- /dev/null
+++ b/ddi/mbor/types/tests/integration/get_unwrapping_key_smoke.rs
@@ -0,0 +1,106 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+//! GetUnwrappingKey smoke tests for the emu backend.
+//!
+//! - Happy path: with an open session, the command returns an RSA-2048
+//!   public key and a non-zero key id.  We intentionally do *not*
+//!   assert on `masked_key` — the firmware emits an empty placeholder
+//!   until vault-key masking is wired up.
+//! - Stability: two calls on the same partition return the same key id
+//!   and public-key bytes — the key is partition-cached and lazily
+//!   generated on first use.
+//! - Without a session: rejected with `FileHandleSessionIdDoesNotMatch`
+//!   by the host-side dev validator before the request reaches firmware.
+
+#![cfg(test)]
+
+use azihsm_ddi::*;
+use azihsm_ddi_mbor_types::*;
+use test_with_tracing::test;
+
+use super::common::*;
+
+#[test]
+fn test_get_unwrapping_key_smoke() {
+    ddi_dev_test(
+        common_setup,
+        common_cleanup,
+        |dev, _ddi, _path, session_id| {
+            let resp = helper_get_unwrapping_key(
+                dev,
+                Some(session_id),
+                Some(DdiApiRev { major: 1, minor: 0 }),
+            )
+            .expect("get_unwrapping_key should succeed");
+
+            assert_eq!(resp.hdr.op, DdiOp::GetUnwrappingKey);
+            assert_eq!(resp.hdr.status, DdiStatus::Success);
+            assert_eq!(resp.hdr.sess_id, Some(session_id));
+
+            assert_ne!(resp.data.key_id, 0, "unwrap key_id must be non-zero");
+            assert!(
+                !resp.data.pub_key.der.is_empty(),
+                "unwrap pub_key must be non-empty"
+            );
+            assert_eq!(
+                resp.data.pub_key.key_kind,
+                DdiKeyType::Rsa2kPublic,
+                "unwrap key must be RSA-2048"
+            );
+        },
+    );
+}
+
+#[test]
+fn test_get_unwrapping_key_stable_across_calls_smoke() {
+    ddi_dev_test(
+        common_setup,
+        common_cleanup,
+        |dev, _ddi, _path, session_id| {
+            let r1 = helper_get_unwrapping_key(
+                dev,
+                Some(session_id),
+                Some(DdiApiRev { major: 1, minor: 0 }),
+            )
+            .expect("first get_unwrapping_key");
+            let r2 = helper_get_unwrapping_key(
+                dev,
+                Some(session_id),
+                Some(DdiApiRev { major: 1, minor: 0 }),
+            )
+            .expect("second get_unwrapping_key");
+
+            assert_eq!(
+                r1.data.key_id, r2.data.key_id,
+                "repeat call must return the same key id"
+            );
+            assert_eq!(
+                r1.data.pub_key.der.as_slice(),
+                r2.data.pub_key.der.as_slice(),
+                "repeat call must return the same pub_key bytes"
+            );
+        },
+    );
+}
+
+#[test]
+fn test_get_unwrapping_key_no_session_smoke() {
+    ddi_dev_test(
+        common_setup,
+        common_cleanup,
+        |dev, _ddi, _path, _session_id| {
+            let err = helper_get_unwrapping_key(dev, None, Some(DdiApiRev { major: 1, minor: 0 }))
+                .expect_err("must be rejected without a session");
+
+            assert!(
+                matches!(
+                    err,
+                    DdiError::DdiStatus(DdiStatus::FileHandleSessionIdDoesNotMatch)
+                ),
+                "expected FileHandleSessionIdDoesNotMatch, got {:?}",
+                err
+            );
+        },
+    );
+}
diff --git a/fw/core/ddi/mbor/types/src/get_unwrapping_key.rs b/fw/core/ddi/mbor/types/src/get_unwrapping_key.rs
index f5c6c9287..32c391009 100644
--- a/fw/core/ddi/mbor/types/src/get_unwrapping_key.rs
+++ b/fw/core/ddi/mbor/types/src/get_unwrapping_key.rs
@@ -14,7 +14,7 @@ pub struct DdiGetUnwrappingKeyReq {}
 pub struct DdiGetUnwrappingKeyResp<'a> {
     #[ddi(id = 1)]
     pub key_id: u16,
-    #[ddi(id = 2)]
+    #[ddi(id = 2, frame)]
     pub pub_key: DdiPublicKey<'a>,
     #[ddi(id = 3, max_len = 1024)]
     pub masked_key: &'a [u8],
diff --git a/fw/core/lib/src/ddi/mbor/get_unwrapping_key.rs b/fw/core/lib/src/ddi/mbor/get_unwrapping_key.rs
new file mode 100644
index 000000000..aff09a601
--- /dev/null
+++ b/fw/core/lib/src/ddi/mbor/get_unwrapping_key.rs
@@ -0,0 +1,75 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+//! DDI GetUnwrappingKey command handler.
+//!
+//! Within an open session, return the partition's RSA-2048 *unwrapping*
+//! key — the public key (raw wire `n_le ‖ e_le`) plus its vault key id,
+//! used by the host to RSA-AES key-wrap a payload for
+//! [`RsaUnwrap`](super::DdiOp::RsaUnwrap).
+//!
+//! The unwrapping key is materialised by the PAL and cached for the
+//! lifetime of the partition.  RSA key generation is expensive, so it
+//! is never done at partition enable: the std (emulator) PAL generates
+//! it lazily on the first call, while hardware PALs generate it in the
+//! background from partition init and report `PendingKeyGeneration`
+//! until it is ready, prompting the host to retry.
+
+use azihsm_fw_ddi_mbor_types::get_unwrapping_key::DdiGetUnwrappingKeyReq;
+use azihsm_fw_ddi_mbor_types::get_unwrapping_key::DdiGetUnwrappingKeyResp;
+use azihsm_fw_ddi_mbor_types::DdiPublicKeyFrameParams;
+
+use super::*;
+
+/// Handle `DdiGetUnwrappingKeyCmd`.
+pub(crate) async fn get_unwrapping_key<'p, P: HsmPal>(
+    pal: &'p P,
+    io: &impl HsmIo,
+    decoder: &mut DdiDecoder<'_>,
+    hdr: &DdiReqHdr,
+) -> HsmResult<&'p DmaBuf> {
+    let _body: DdiGetUnwrappingKeyReq = decoder.decode_data()?;
+    let sess_id = hdr.sess_id.ok_or(HsmError::SessionExpected)?;
+
+    // Ensure the partition's RSA-2048 unwrapping key exists (generated
+    // lazily by the std PAL, or in the background on hardware) and get
+    // its vault key id.  On hardware this may return
+    // `PendingKeyGeneration` while generation is in flight, which the
+    // host retries.
+    let key_id = pal.part_ensure_unwrapping_key(io).await?;
+
+    // Query the wire length of the unwrapping public key derived from
+    // the vault-stored private key — no separate public key is cached
+    // (matches the reference firmware).  The actual serialization writes
+    // straight into the reserved response slot below.
+    let priv_key = pal.vault_key(io, key_id)?;
+    let pub_len = pal.rsa_priv_pub_key(io, priv_key, None)?;
+
+    // `masked_key` is the host's opaque re-import blob; firmware-side
+    // masking is pending the corresponding unmask path — emit an empty
+    // placeholder for now so the response is wire-valid.
+    let (resp, layout) = pal.dma_alloc_var_with(io, |buf| {
+        let mut encoder = super::encode_resp_hdr(
+            &super::success_hdr_sess(hdr, DdiOp::GetUnwrappingKey, sess_id),
+            buf,
+        )?;
+        let layout = DdiGetUnwrappingKeyResp::reserve(
+            &mut encoder,
+            u16::from(key_id),
+            DdiPublicKeyFrameParams {
+                raw_len: pub_len,
+                key_kind: DdiKeyType::Rsa2kPublic,
+            },
+            0, /* masked_key length — empty placeholder */
+        )?;
+        Ok((encoder.position(), layout))
+    })?;
+
+    // Serialize the wire-format public key directly into the frame's
+    // reserved slot — the PAL converts its vault representation into the
+    // wire form (incl. any big-endian↔little-endian flip).
+    let frame = DdiGetUnwrappingKeyResp::from_layout(resp, &layout);
+    pal.rsa_priv_pub_key(io, priv_key, Some(frame.pub_key.raw))?;
+
+    Ok(resp)
+}
diff --git a/fw/core/lib/src/ddi/mbor/mod.rs b/fw/core/lib/src/ddi/mbor/mod.rs
index 7e880fc81..f3d76ee30 100644
--- a/fw/core/lib/src/ddi/mbor/mod.rs
+++ b/fw/core/lib/src/ddi/mbor/mod.rs
@@ -18,6 +18,7 @@ pub(crate) mod get_device_info;
 pub(crate) mod get_establish_cred_encryption_key;
 pub(crate) mod get_sealed_bk3;
 pub(crate) mod get_session_encryption_key;
+pub(crate) mod get_unwrapping_key;
 pub(crate) mod hkdf_derive;
 pub(crate) mod hmac;
 pub(crate) mod init_bk3;
@@ -48,6 +49,7 @@ pub(crate) use get_device_info::*;
 pub(crate) use get_establish_cred_encryption_key::*;
 pub(crate) use get_sealed_bk3::*;
 pub(crate) use get_session_encryption_key::*;
+pub(crate) use get_unwrapping_key::*;
 pub(crate) use hkdf_derive::*;
 pub(crate) use hmac::*;
 pub(crate) use init_bk3::*;
@@ -124,6 +126,7 @@ pub(crate) async fn dispatch<'p, P: HsmPal>(
             get_establish_cred_encryption_key(pal, io, decoder, hdr).await
         }
         DdiOp::GetSessionEncryptionKey => get_session_encryption_key(pal, io, decoder, hdr).await,
+        DdiOp::GetUnwrappingKey => get_unwrapping_key(pal, io, decoder, hdr).await,
         DdiOp::GetSealedBk3 => get_sealed_bk3(pal, io, decoder, hdr),
         DdiOp::SetSealedBk3 => set_sealed_bk3(pal, io, decoder, hdr),
         DdiOp::InitBk3 => init_bk3(pal, io, decoder, hdr).await,
diff --git a/fw/pal/traits/src/crypto/rsa.rs b/fw/pal/traits/src/crypto/rsa.rs
index f34c5b2fc..e57862347 100644
--- a/fw/pal/traits/src/crypto/rsa.rs
+++ b/fw/pal/traits/src/crypto/rsa.rs
@@ -191,7 +191,7 @@ pub trait HsmRsa {
     /// - `Err(HsmError::InvalidArg)` — buffer-size mismatch.
     /// - `Err(HsmError)` — PKA / RNG failure or PCT failed (the key
     ///   pair is rejected).
-    async fn ras_gen_keypair(
+    async fn rsa_gen_keypair(
         &self,
         io: &impl HsmIo,
         key_size: HsmRsaKey,
@@ -257,6 +257,39 @@ pub trait HsmRsa {
         y: &mut DmaBuf,
     ) -> Result<(), HsmError>;
 
+    /// Derive the raw wire-format public key (`n_le || e_le`) from a
+    /// vault-stored RSA private key.
+    ///
+    /// Used to recover the public key of a vault-stored private key
+    /// (e.g. the partition unwrapping key) without persisting a
+    /// separate copy.
+    ///
+    /// `priv_key` is in the PAL's own vault representation — raw key
+    /// material on real PKA hardware, DER in the std/OpenSSL PAL — and
+    /// this method converts it into the wire form: the little-endian
+    /// modulus followed by a fixed 4-byte little-endian exponent (the
+    /// raw layout the host's `DdiDerPublicKey` post-decode turns into
+    /// DER).  The PAL owns the whole vault-format → wire-format
+    /// conversion, including any big-endian↔little-endian flip (real
+    /// PKA is little-endian native; an OpenSSL backend holds big-endian
+    /// components and reverses them).
+    ///
+    /// Follows the query/alloc/use convention: pass `pub_out = None` to
+    /// query the wire length the caller must allocate, then
+    /// `pub_out = Some(buf)` to serialize.
+    ///
+    /// # Returns
+    /// - `Ok(len)` — wire byte length (`modulus_len + 4`); in `Some`
+    ///   mode those bytes are written into `pub_out`.
+    /// - `Err(HsmError::InvalidArg)` — `priv_key` is not valid vault
+    ///   format or `pub_out` is too small.
+    fn rsa_priv_pub_key(
+        &self,
+        io: &impl HsmIo,
+        priv_key: &DmaBuf,
+        pub_out: Option<&mut DmaBuf>,
+    ) -> HsmResult<usize>;
+
     /// PKCS#1 v1.5 encrypt (EME-PKCS1-v1_5).
     ///
     /// Pads `message` with random non-zero bytes per RFC 8017 §7.2.1
diff --git a/fw/pal/traits/src/part.rs b/fw/pal/traits/src/part.rs
index c31e39ed1..aa11175c9 100644
--- a/fw/pal/traits/src/part.rs
+++ b/fw/pal/traits/src/part.rs
@@ -414,6 +414,32 @@ pub trait HsmPartitionManager {
     /// [`HsmResult<()>`] — `Ok(())` on success.  Idempotent on an
     /// already-absent slot (also returns `Ok(())`).
     fn part_prop_clear(&self, io: &impl HsmIo, id: PartPropId) -> HsmResult<()>;
+
+    /// Obtain the partition's RSA-2048 unwrapping key, returning its
+    /// vault key id once the key is available.
+    ///
+    /// The key is materialised once per partition and cached: the
+    /// private key lives in the vault (with `internal`/`local`/`unwrap`
+    /// attributes), its id in
+    /// [`RSA_UNWRAPPING_KEY_ID`](PartPropId::RSA_UNWRAPPING_KEY_ID), and
+    /// the wire-format public key in
+    /// [`RSA_UNWRAPPING_PUB_KEY`](PartPropId::RSA_UNWRAPPING_PUB_KEY).
+    ///
+    /// RSA key generation is expensive, so PALs never generate it at
+    /// partition enable.  Implementations differ in *when* they
+    /// generate:
+    /// - The std (emulator) PAL generates lazily on the first call and
+    ///   returns the id synchronously.
+    /// - Hardware PALs kick generation off in the background at
+    ///   partition init; while it is still in flight they return
+    ///   [`HsmError::PendingKeyGeneration`] so the host retries until
+    ///   the key is ready.
+    ///
+    /// # Errors
+    /// - [`HsmError::PendingKeyGeneration`] — generation is still in
+    ///   progress; the host should retry.
+    /// - Propagated PKA / vault failures.
+    async fn part_ensure_unwrapping_key(&self, io: &impl HsmIo) -> HsmResult<HsmKeyId>;
 }
 
 /// Length of the per-partition `BK_BOOT` boot-key material in bytes.
diff --git a/fw/plat/std/pal/src/drivers/ecc.rs b/fw/plat/std/pal/src/drivers/ecc.rs
index 3a4f901b5..b53c15b24 100644
--- a/fw/plat/std/pal/src/drivers/ecc.rs
+++ b/fw/plat/std/pal/src/drivers/ecc.rs
@@ -41,31 +41,9 @@ use azihsm_crypto::SignOp;
 use azihsm_crypto::VerifyOp;
 use azihsm_fw_hsm_pal_traits::*;
 
+use super::reverse_copy;
 use crate::worker::WorkerPool;
 
-// On-wire-format helpers
-//
-// The std PAL accepts inputs and produces outputs in the wire-LE
-// convention that real PKA hardware emits natively (`x_le || y_le`
-// for public keys, `r_le || s_le` for signatures, with P-521
-// components padded from 66 to 68 bytes per word).  The single
-// `reverse_copy` helper centralizes the BE↔LE byte reversal so the
-// `_le`-suffixed driver methods can stay free of byte-shuffling
-// boilerplate.
-
-/// Reverse-copy `src` into `dst[..src.len()]`.  Used for BE↔LE
-/// component swaps; the direction is symmetric.  Trailing pad bytes
-/// in `dst` (e.g. the 2-byte pad after a 66-byte P-521 coordinate
-/// inside a 68-byte word) are left untouched — callers that need
-/// them zeroed should pre-`fill(0)` the full `dst` slot.
-fn reverse_copy(dst: &mut [u8], src: &[u8]) {
-    let len = src.len();
-    debug_assert!(dst.len() >= len);
-    for (d, s) in dst[..len].iter_mut().zip(src.iter().rev()) {
-        *d = *s;
-    }
-}
-
 /// Std ECC driver — software ECC via OpenSSL with async worker dispatch.
 pub struct StdEcc {
     pool: WorkerPool,
diff --git a/fw/plat/std/pal/src/drivers/mod.rs b/fw/plat/std/pal/src/drivers/mod.rs
index bb333cf1a..d0f90f58f 100644
--- a/fw/plat/std/pal/src/drivers/mod.rs
+++ b/fw/plat/std/pal/src/drivers/mod.rs
@@ -51,3 +51,21 @@ pub mod oic;
 pub mod rsa;
 pub mod session;
 pub mod vault;
+
+/// Reverse-copy `src` into `dst[..src.len()]`.
+///
+/// Centralizes the big-endian↔little-endian byte reversal shared by the
+/// crypto drivers: their wire convention is little-endian (what real PKA
+/// hardware emits natively), while the OpenSSL backend produces
+/// big-endian components.  The reversal is symmetric, so the same helper
+/// serves both directions.  Trailing pad bytes in `dst` beyond
+/// `src.len()` (e.g. the 2-byte pad after a 66-byte P-521 coordinate
+/// inside a 68-byte word) are left untouched — callers that need them
+/// zeroed should pre-`fill(0)` the full `dst` slot.
+pub(crate) fn reverse_copy(dst: &mut [u8], src: &[u8]) {
+    let len = src.len();
+    debug_assert!(dst.len() >= len);
+    for (d, s) in dst[..len].iter_mut().zip(src.iter().rev()) {
+        *d = *s;
+    }
+}
diff --git a/fw/plat/std/pal/src/drivers/rsa.rs b/fw/plat/std/pal/src/drivers/rsa.rs
index aca73b05b..bec613c5e 100644
--- a/fw/plat/std/pal/src/drivers/rsa.rs
+++ b/fw/plat/std/pal/src/drivers/rsa.rs
@@ -52,12 +52,58 @@ use azihsm_crypto::EncryptOp;
 use azihsm_crypto::KeyGenerationOp;
 use azihsm_crypto::PrivateKey;
 use azihsm_crypto::RsaEncryptAlgo;
+use azihsm_crypto::RsaKeyOp;
 use azihsm_crypto::RsaPrivateKey;
 use azihsm_crypto::RsaPublicKey;
 use azihsm_fw_hsm_pal_traits::*;
 
 use crate::worker::WorkerPool;
 
+/// Fixed width (bytes) of the public exponent `e` in the raw wire
+/// public-key encoding (`n_le || e_le`).
+pub const RSA_PUB_EXP_WIRE_LEN: usize = 4;
+
+/// Serialize an RSA public-key handle into the raw wire form
+/// `n_le || e_le` — the little-endian modulus followed by a fixed
+/// 4-byte little-endian exponent, the raw layout the host's
+/// `DdiDerPublicKey` post-decode turns into DER.
+///
+/// The big-endian↔little-endian flip lives here in the driver (real
+/// PKA hardware is little-endian native; the OpenSSL backend produces
+/// big-endian components, reversed below).
+///
+/// Follows the query/alloc/use convention: pass `out = None` to query
+/// the wire length (`n_len + RSA_PUB_EXP_WIRE_LEN`) the caller must
+/// allocate, then `out = Some(buf)` to serialize.  Returns the wire
+/// length in both modes.
+pub fn rsa_pub_wire(pubk: &RsaPublicKey, out: Option<&mut [u8]>) -> HsmResult<usize> {
+    let n_len = pubk.n(None).map_err(|_| HsmError::RsaGenerateError)?;
+    let wire_len = n_len + RSA_PUB_EXP_WIRE_LEN;
+    // Query mode: report the buffer size the caller must allocate.
+    let Some(out) = out else {
+        return Ok(wire_len);
+    };
+    if out.len() < wire_len {
+        return Err(HsmError::RsaInvalidKeyLength);
+    }
+    // Modulus: big-endian -> little-endian into the leading `n_len` bytes.
+    let mut n_be = [0u8; 512];
+    pubk.n(Some(&mut n_be[..n_len]))
+        .map_err(|_| HsmError::RsaGenerateError)?;
+    super::reverse_copy(&mut out[..n_len], &n_be[..n_len]);
+    // Exponent: right-align big-endian in the fixed wire field, then
+    // reverse the whole field to little-endian.
+    let e_len = pubk.e(None).map_err(|_| HsmError::RsaGenerateError)?;
+    if e_len > RSA_PUB_EXP_WIRE_LEN {
+        return Err(HsmError::RsaGenerateError);
+    }
+    let mut e_be = [0u8; RSA_PUB_EXP_WIRE_LEN];
+    pubk.e(Some(&mut e_be[RSA_PUB_EXP_WIRE_LEN - e_len..]))
+        .map_err(|_| HsmError::RsaGenerateError)?;
+    super::reverse_copy(&mut out[n_len..wire_len], &e_be);
+    Ok(wire_len)
+}
+
 /// Std RSA driver — software RSA via OpenSSL with async worker dispatch.
 ///
 /// Created once during PAL initialization and shared across all IO tasks.
diff --git a/fw/plat/std/pal/src/part.rs b/fw/plat/std/pal/src/part.rs
index da2105ca4..3bca49023 100644
--- a/fw/plat/std/pal/src/part.rs
+++ b/fw/plat/std/pal/src/part.rs
@@ -453,6 +453,42 @@ impl HsmPartitionManager for StdHsmPal {
     fn part_prop_clear(&self, io: &impl HsmIo, id: PartPropId) -> HsmResult<()> {
         self.prop_clear(io, id)
     }
+
+    async fn part_ensure_unwrapping_key(&self, io: &impl HsmIo) -> HsmResult<HsmKeyId> {
+        let pid = io.pid();
+        // Fast path: already materialised this enable cycle.
+        if let Some(kid) = self.active_part(pid)?.unwrapping_key_id {
+            return Ok(kid);
+        }
+
+        // Generate the RSA-2048 key pair without holding a borrow into
+        // the partition table across the keygen await.  Only the private
+        // key is persisted in the vault (as DER — the std PAL's
+        // representation; real firmware stores raw key material); the
+        // public key is later derived from it on demand (matching the
+        // reference firmware), so nothing extra is cached in partition
+        // state.
+        let (pk, _pubk) = self.rsa.gen_keypair(2048).await?;
+        let priv_len = pk.to_bytes(None).map_err(|_| HsmError::RsaToDerError)?;
+        let mut priv_buf = vec![0u8; priv_len];
+        pk.to_bytes(Some(&mut priv_buf[..priv_len]))
+            .map_err(|_| HsmError::RsaToDerError)?;
+        let attrs = HsmVaultKeyAttrs::new()
+            .with_internal(true)
+            .with_local(true)
+            .with_unwrap(true);
+
+        // Store the private key in the vault and record its id.
+        let entry = self.active_part_mut(pid)?;
+        let kid = entry.vault.create(
+            &priv_buf[..priv_len],
+            HsmVaultKeyKind::Rsa2kPrivate,
+            None,
+            attrs,
+        )?;
+        entry.unwrapping_key_id = Some(kid);
+        Ok(kid)
+    }
 }
 
 // ---------------------------------------------------------------------------
diff --git a/fw/plat/std/pal/src/rsa.rs b/fw/plat/std/pal/src/rsa.rs
index 3399b1090..f6131edd4 100644
--- a/fw/plat/std/pal/src/rsa.rs
+++ b/fw/plat/std/pal/src/rsa.rs
@@ -13,6 +13,7 @@
 
 use azihsm_crypto::ExportableKey;
 use azihsm_crypto::ImportableKey;
+use azihsm_crypto::PrivateKey;
 use azihsm_crypto::RsaPrivateKey;
 use azihsm_crypto::RsaPublicKey;
 
@@ -23,7 +24,7 @@ fn key_size_bits(key_size: HsmRsaKey) -> usize {
 }
 
 impl HsmRsa for StdHsmPal {
-    async fn ras_gen_keypair(
+    async fn rsa_gen_keypair(
         &self,
         _io: &impl HsmIo,
         key_size: HsmRsaKey,
@@ -74,6 +75,22 @@ impl HsmRsa for StdHsmPal {
         self.rsa.mod_exp_pub(&pub_key, x, y).await
     }
 
+    fn rsa_priv_pub_key(
+        &self,
+        _io: &impl HsmIo,
+        priv_key: &DmaBuf,
+        pub_out: Option<&mut DmaBuf>,
+    ) -> HsmResult<usize> {
+        // The std PAL's vault representation is DER: parse it, derive
+        // the public key, and emit the raw wire form `n_le || e_le`.
+        // This is the vault-format → wire-format conversion; the BE->LE
+        // flip lives in the driver.  In query mode (`pub_out == None`)
+        // only the wire length is computed and returned.
+        let pk = RsaPrivateKey::from_bytes(priv_key).map_err(|_| HsmError::InvalidArg)?;
+        let pubk = pk.public_key().map_err(|_| HsmError::RsaGenerateError)?;
+        crate::drivers::rsa::rsa_pub_wire(&pubk, pub_out.map(|b| &mut **b))
+    }
+
     async fn rsa_pkcs1_encrypt<'a>(
         &self,
         _io: &impl HsmIo,
diff --git a/fw/plat/uno/fw/pal/src/crypto/rsa.rs b/fw/plat/uno/fw/pal/src/crypto/rsa.rs
index 53d084e3c..01c98c061 100644
--- a/fw/plat/uno/fw/pal/src/crypto/rsa.rs
+++ b/fw/plat/uno/fw/pal/src/crypto/rsa.rs
@@ -94,7 +94,7 @@ impl UnoHsmPal {
 }
 
 impl HsmRsa for UnoHsmPal {
-    async fn ras_gen_keypair(
+    async fn rsa_gen_keypair(
         &self,
         _io: &impl HsmIo,
         _key_size: HsmRsaKey,
@@ -129,6 +129,17 @@ impl HsmRsa for UnoHsmPal {
         self.pka.rsa_mod_exp_pub(upka_key_type, key, x, y).await
     }
 
+    fn rsa_priv_pub_key(
+        &self,
+        _io: &impl HsmIo,
+        _priv_key: &DmaBuf,
+        _pub_out: Option<&mut DmaBuf>,
+    ) -> HsmResult<usize> {
+        // TODO: convert the raw vault-format RSA private key into the
+        // wire public key on Uno PKA (GetUnwrappingKey / RsaUnwrap).
+        Err(HsmError::UnsupportedCmd)
+    }
+
     // ── PKCS#1 v1.5 encryption ─────────────────────────────────────
 
     async fn rsa_pkcs1_encrypt<'a>(
diff --git a/fw/plat/uno/fw/pal/src/part.rs b/fw/plat/uno/fw/pal/src/part.rs
index 160d1ec48..ff9e6a3ae 100644
--- a/fw/plat/uno/fw/pal/src/part.rs
+++ b/fw/plat/uno/fw/pal/src/part.rs
@@ -684,4 +684,13 @@ impl HsmPartitionManager for UnoHsmPal {
         }
         Ok(())
     }
+
+    async fn part_ensure_unwrapping_key(&self, _io: &impl HsmIo) -> HsmResult<HsmKeyId> {
+        // On Uno the RSA-2048 unwrapping key is generated in the
+        // background, started at partition init, and can take a while.
+        // Until that upka-backed generation (and its completion check)
+        // is wired up here, report it as still in progress so the host
+        // retries — mirroring the reference firmware's behaviour.
+        Err(HsmError::PendingKeyGeneration)
+    }
 }

From d9a6e3861730180e32ad0af6f0344effbaf09997 Mon Sep 17 00:00:00 2001
From: Jayant Gandhi <jayg@microsoft.com>
Date: Tue, 30 Jun 2026 18:20:20 +0000
Subject: [PATCH 2/4] fix(fw): address GetUnwrappingKey review feedback
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- part_ensure_unwrapping_key (std): close a TOCTOU race across the
  keygen await — re-check unwrapping_key_id while holding the mutable
  entry borrow before creating the vault key, so two concurrent callers
  on one partition can't both insert (leaking one) and return different
  ids. The re-check -> create -> store window is await-free, hence
  atomic on the single-threaded executor.
- rsa_pub_wire (std driver): bound the modulus length against the fixed
  512-byte scratch buffer and error out, so a malformed/oversized vault
  key can never panic across the PAL trust boundary.
- Docs: drop the stale PartPropId::RSA_UNWRAPPING_PUB_KEY reference (no
  such property — the public key is derived on demand, not cached), and
  reword the unwrapping-key caching as the partition's current enabled
  incarnation (clear_enabled_state deletes it on disable/free).
- Smoke test: drop the `key_id != 0` assertion — key_id is an opaque
  handle and the sim/mock backend (which CI runs this suite under)
  legitimately assigns 0; the stability test already covers id
  consistency.

Validation: get_unwrapping_key smoke 3/3 under both mock and emu; full
mock smoke 45/45; std unit 157/157; host + Uno build clean.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 .../integration/get_unwrapping_key_smoke.rs      | 11 ++++++-----
 fw/core/lib/src/ddi/mbor/get_unwrapping_key.rs   | 13 ++++++++-----
 fw/pal/traits/src/part.rs                        | 16 ++++++++++------
 fw/plat/std/pal/src/drivers/rsa.rs               | 12 ++++++++++--
 fw/plat/std/pal/src/part.rs                      | 12 +++++++++++-
 5 files changed, 45 insertions(+), 19 deletions(-)

diff --git a/ddi/mbor/types/tests/integration/get_unwrapping_key_smoke.rs b/ddi/mbor/types/tests/integration/get_unwrapping_key_smoke.rs
index 6111f4ce4..cb6d8932c 100644
--- a/ddi/mbor/types/tests/integration/get_unwrapping_key_smoke.rs
+++ b/ddi/mbor/types/tests/integration/get_unwrapping_key_smoke.rs
@@ -1,12 +1,14 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-//! GetUnwrappingKey smoke tests for the emu backend.
+//! GetUnwrappingKey smoke tests.
 //!
 //! - Happy path: with an open session, the command returns an RSA-2048
-//!   public key and a non-zero key id.  We intentionally do *not*
-//!   assert on `masked_key` — the firmware emits an empty placeholder
-//!   until vault-key masking is wired up.
+//!   public key and a vault key id.  We intentionally do *not* assert
+//!   on `masked_key` — the firmware emits an empty placeholder until
+//!   vault-key masking is wired up — nor on the exact `key_id` value,
+//!   which is an opaque handle (the sim backend legitimately assigns
+//!   `0`).
 //! - Stability: two calls on the same partition return the same key id
 //!   and public-key bytes — the key is partition-cached and lazily
 //!   generated on first use.
@@ -38,7 +40,6 @@ fn test_get_unwrapping_key_smoke() {
             assert_eq!(resp.hdr.status, DdiStatus::Success);
             assert_eq!(resp.hdr.sess_id, Some(session_id));
 
-            assert_ne!(resp.data.key_id, 0, "unwrap key_id must be non-zero");
             assert!(
                 !resp.data.pub_key.der.is_empty(),
                 "unwrap pub_key must be non-empty"
diff --git a/fw/core/lib/src/ddi/mbor/get_unwrapping_key.rs b/fw/core/lib/src/ddi/mbor/get_unwrapping_key.rs
index aff09a601..680e11a69 100644
--- a/fw/core/lib/src/ddi/mbor/get_unwrapping_key.rs
+++ b/fw/core/lib/src/ddi/mbor/get_unwrapping_key.rs
@@ -9,11 +9,14 @@
 //! [`RsaUnwrap`](super::DdiOp::RsaUnwrap).
 //!
 //! The unwrapping key is materialised by the PAL and cached for the
-//! lifetime of the partition.  RSA key generation is expensive, so it
-//! is never done at partition enable: the std (emulator) PAL generates
-//! it lazily on the first call, while hardware PALs generate it in the
-//! background from partition init and report `PendingKeyGeneration`
-//! until it is ready, prompting the host to retry.
+//! partition's current enabled incarnation — it is dropped (and its
+//! vault slot deleted) when the partition is disabled or freed, and
+//! regenerated on the next request after a re-enable.  RSA key
+//! generation is expensive, so it is never done at partition enable:
+//! the std (emulator) PAL generates it lazily on the first call, while
+//! hardware PALs generate it in the background from partition init and
+//! report `PendingKeyGeneration` until it is ready, prompting the host
+//! to retry.
 
 use azihsm_fw_ddi_mbor_types::get_unwrapping_key::DdiGetUnwrappingKeyReq;
 use azihsm_fw_ddi_mbor_types::get_unwrapping_key::DdiGetUnwrappingKeyResp;
diff --git a/fw/pal/traits/src/part.rs b/fw/pal/traits/src/part.rs
index aa11175c9..7c41c0f89 100644
--- a/fw/pal/traits/src/part.rs
+++ b/fw/pal/traits/src/part.rs
@@ -418,12 +418,16 @@ pub trait HsmPartitionManager {
     /// Obtain the partition's RSA-2048 unwrapping key, returning its
     /// vault key id once the key is available.
     ///
-    /// The key is materialised once per partition and cached: the
-    /// private key lives in the vault (with `internal`/`local`/`unwrap`
-    /// attributes), its id in
-    /// [`RSA_UNWRAPPING_KEY_ID`](PartPropId::RSA_UNWRAPPING_KEY_ID), and
-    /// the wire-format public key in
-    /// [`RSA_UNWRAPPING_PUB_KEY`](PartPropId::RSA_UNWRAPPING_PUB_KEY).
+    /// The key is materialised once per enabled partition incarnation:
+    /// the private key lives in the vault (with `internal`/`local`/
+    /// `unwrap` attributes) and its id in
+    /// [`RSA_UNWRAPPING_KEY_ID`](PartPropId::RSA_UNWRAPPING_KEY_ID).
+    /// It is dropped when the partition is disabled or freed and
+    /// regenerated on the next request after a re-enable.  No public key
+    /// is cached — callers derive it on demand from the vault private
+    /// key via
+    /// [`HsmRsa::rsa_priv_pub_key`](crate::HsmRsa::rsa_priv_pub_key),
+    /// matching the reference firmware.
     ///
     /// RSA key generation is expensive, so PALs never generate it at
     /// partition enable.  Implementations differ in *when* they
diff --git a/fw/plat/std/pal/src/drivers/rsa.rs b/fw/plat/std/pal/src/drivers/rsa.rs
index bec613c5e..58a4944b8 100644
--- a/fw/plat/std/pal/src/drivers/rsa.rs
+++ b/fw/plat/std/pal/src/drivers/rsa.rs
@@ -86,8 +86,16 @@ pub fn rsa_pub_wire(pubk: &RsaPublicKey, out: Option<&mut [u8]>) -> HsmResult<us
     if out.len() < wire_len {
         return Err(HsmError::RsaInvalidKeyLength);
     }
-    // Modulus: big-endian -> little-endian into the leading `n_len` bytes.
-    let mut n_be = [0u8; 512];
+    // Modulus: big-endian -> little-endian into the leading `n_len`
+    // bytes.  `pubk` is derived from untrusted vault data, so guard the
+    // fixed-size scratch buffer against an out-of-range modulus length
+    // (e.g. a larger DER-imported modulus) rather than panicking on the
+    // slice.
+    const MAX_MODULUS_LEN: usize = 512; // RSA-4096
+    if n_len > MAX_MODULUS_LEN {
+        return Err(HsmError::RsaInvalidKeyLength);
+    }
+    let mut n_be = [0u8; MAX_MODULUS_LEN];
     pubk.n(Some(&mut n_be[..n_len]))
         .map_err(|_| HsmError::RsaGenerateError)?;
     super::reverse_copy(&mut out[..n_len], &n_be[..n_len]);
diff --git a/fw/plat/std/pal/src/part.rs b/fw/plat/std/pal/src/part.rs
index 3bca49023..ecd272a39 100644
--- a/fw/plat/std/pal/src/part.rs
+++ b/fw/plat/std/pal/src/part.rs
@@ -478,8 +478,18 @@ impl HsmPartitionManager for StdHsmPal {
             .with_local(true)
             .with_unwrap(true);
 
-        // Store the private key in the vault and record its id.
+        // Store the private key in the vault and record its id.  The
+        // `active_part_mut` borrow through to setting `unwrapping_key_id`
+        // is await-free, so on the single-threaded executor it is atomic.
+        // Re-check here: a concurrent caller may have generated and
+        // stored the key while this task was awaiting keygen above.  If
+        // so, drop our freshly generated key (it was never inserted into
+        // the vault) and return the existing id, so the partition only
+        // ever materialises a single unwrapping key.
         let entry = self.active_part_mut(pid)?;
+        if let Some(kid) = entry.unwrapping_key_id {
+            return Ok(kid);
+        }
         let kid = entry.vault.create(
             &priv_buf[..priv_len],
             HsmVaultKeyKind::Rsa2kPrivate,

From 6bce42cb48fcdd137fef41ff5bc6b53da47ebf3d Mon Sep 17 00:00:00 2001
From: Jayant Gandhi <jayg@microsoft.com>
Date: Tue, 30 Jun 2026 23:29:46 +0000
Subject: [PATCH 3/4] fix(fw): drive GetUnwrappingKey off the partition
 property model
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Address PR feedback: follow the existing property model instead of
adding a new PAL trait method for the unwrapping key.

- Remove HsmPartitionManager::part_ensure_unwrapping_key. The handler
  now just reads the partition's RSA_UNWRAPPING_KEY_ID property via
  part_unwrapping_key_id; an absent id (PartPropNotFound) means
  generation is still pending, surfaced as PendingKeyGeneration so the
  host retries. This keeps the core handler PAL-agnostic — real
  hardware generates the key in the background at part-init and only
  ever exposes it through the property.
- std PAL: provision the key lazily behind the property read. A new
  private provision_unwrapping_key runs inside part_prop_get_u16 for
  RSA_UNWRAPPING_KEY_ID, generating the RSA-2048 key synchronously on
  first read so only a GetUnwrappingKey request pays the keygen cost
  (never partition enable — avoids a ~19x smoke-test slowdown). The
  routine is await-free, so on the single-threaded executor it runs
  atomically: concurrent reads cannot race to generate two keys (first
  write wins, the rest reuse it).
- Uno PAL: drop the part_ensure_unwrapping_key stub; its getter stays
  uniform with every other key-id property (PartPropNotFound when
  unset), which the handler maps to PendingKeyGeneration.

Validation: get_unwrapping_key smoke 3/3 under both mock and emu; std
unit 157/157; host + Uno build clean; fmt/copyright clean.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 .../lib/src/ddi/mbor/get_unwrapping_key.rs    | 39 +++++---
 fw/pal/traits/src/part.rs                     | 30 ------
 fw/plat/std/pal/src/part.rs                   | 98 ++++++++++---------
 fw/plat/uno/fw/pal/src/part.rs                |  9 --
 4 files changed, 76 insertions(+), 100 deletions(-)

diff --git a/fw/core/lib/src/ddi/mbor/get_unwrapping_key.rs b/fw/core/lib/src/ddi/mbor/get_unwrapping_key.rs
index 680e11a69..57eb693e4 100644
--- a/fw/core/lib/src/ddi/mbor/get_unwrapping_key.rs
+++ b/fw/core/lib/src/ddi/mbor/get_unwrapping_key.rs
@@ -8,15 +8,18 @@
 //! used by the host to RSA-AES key-wrap a payload for
 //! [`RsaUnwrap`](super::DdiOp::RsaUnwrap).
 //!
-//! The unwrapping key is materialised by the PAL and cached for the
-//! partition's current enabled incarnation — it is dropped (and its
-//! vault slot deleted) when the partition is disabled or freed, and
-//! regenerated on the next request after a re-enable.  RSA key
-//! generation is expensive, so it is never done at partition enable:
-//! the std (emulator) PAL generates it lazily on the first call, while
-//! hardware PALs generate it in the background from partition init and
-//! report `PendingKeyGeneration` until it is ready, prompting the host
-//! to retry.
+//! The unwrapping key id lives in the partition's
+//! [`RSA_UNWRAPPING_KEY_ID`](crate::part_state::part_unwrapping_key_id)
+//! property; this handler simply reads it.  RSA key generation is
+//! expensive, so it is never done at partition enable — each PAL
+//! materialises the key behind the property read instead: the std
+//! (emulator) PAL generates it lazily and synchronously on first read,
+//! while hardware PALs generate it in the background from partition
+//! init and leave the property unset until ready.  An absent id
+//! therefore means generation is still pending, which this handler
+//! surfaces as `PendingKeyGeneration` so the host retries.  No public
+//! key is cached: it is derived from the vault
+//! private key on demand (matching the reference firmware).
 
 use azihsm_fw_ddi_mbor_types::get_unwrapping_key::DdiGetUnwrappingKeyReq;
 use azihsm_fw_ddi_mbor_types::get_unwrapping_key::DdiGetUnwrappingKeyResp;
@@ -34,12 +37,18 @@ pub(crate) async fn get_unwrapping_key<'p, P: HsmPal>(
     let _body: DdiGetUnwrappingKeyReq = decoder.decode_data()?;
     let sess_id = hdr.sess_id.ok_or(HsmError::SessionExpected)?;
 
-    // Ensure the partition's RSA-2048 unwrapping key exists (generated
-    // lazily by the std PAL, or in the background on hardware) and get
-    // its vault key id.  On hardware this may return
-    // `PendingKeyGeneration` while generation is in flight, which the
-    // host retries.
-    let key_id = pal.part_ensure_unwrapping_key(io).await?;
+    // Read the partition's RSA-2048 unwrapping key id from its property.
+    // The PAL materialises the key behind this read: the std PAL
+    // generates it synchronously on first read (so it never reports the
+    // id as absent), while hardware PALs generate it in the background
+    // and leave the property unset until ready.  An absent id therefore
+    // means generation is still pending — surface it as such so the host
+    // retries.
+    let key_id = match crate::part_state::part_unwrapping_key_id(pal, io) {
+        Ok(id) => id,
+        Err(HsmError::PartPropNotFound) => return Err(HsmError::PendingKeyGeneration),
+        Err(e) => return Err(e),
+    };
 
     // Query the wire length of the unwrapping public key derived from
     // the vault-stored private key — no separate public key is cached
diff --git a/fw/pal/traits/src/part.rs b/fw/pal/traits/src/part.rs
index 7c41c0f89..c31e39ed1 100644
--- a/fw/pal/traits/src/part.rs
+++ b/fw/pal/traits/src/part.rs
@@ -414,36 +414,6 @@ pub trait HsmPartitionManager {
     /// [`HsmResult<()>`] — `Ok(())` on success.  Idempotent on an
     /// already-absent slot (also returns `Ok(())`).
     fn part_prop_clear(&self, io: &impl HsmIo, id: PartPropId) -> HsmResult<()>;
-
-    /// Obtain the partition's RSA-2048 unwrapping key, returning its
-    /// vault key id once the key is available.
-    ///
-    /// The key is materialised once per enabled partition incarnation:
-    /// the private key lives in the vault (with `internal`/`local`/
-    /// `unwrap` attributes) and its id in
-    /// [`RSA_UNWRAPPING_KEY_ID`](PartPropId::RSA_UNWRAPPING_KEY_ID).
-    /// It is dropped when the partition is disabled or freed and
-    /// regenerated on the next request after a re-enable.  No public key
-    /// is cached — callers derive it on demand from the vault private
-    /// key via
-    /// [`HsmRsa::rsa_priv_pub_key`](crate::HsmRsa::rsa_priv_pub_key),
-    /// matching the reference firmware.
-    ///
-    /// RSA key generation is expensive, so PALs never generate it at
-    /// partition enable.  Implementations differ in *when* they
-    /// generate:
-    /// - The std (emulator) PAL generates lazily on the first call and
-    ///   returns the id synchronously.
-    /// - Hardware PALs kick generation off in the background at
-    ///   partition init; while it is still in flight they return
-    ///   [`HsmError::PendingKeyGeneration`] so the host retries until
-    ///   the key is ready.
-    ///
-    /// # Errors
-    /// - [`HsmError::PendingKeyGeneration`] — generation is still in
-    ///   progress; the host should retry.
-    /// - Propagated PKA / vault failures.
-    async fn part_ensure_unwrapping_key(&self, io: &impl HsmIo) -> HsmResult<HsmKeyId>;
 }
 
 /// Length of the per-partition `BK_BOOT` boot-key material in bytes.
diff --git a/fw/plat/std/pal/src/part.rs b/fw/plat/std/pal/src/part.rs
index ecd272a39..259d6e944 100644
--- a/fw/plat/std/pal/src/part.rs
+++ b/fw/plat/std/pal/src/part.rs
@@ -419,6 +419,15 @@ impl HsmPartitionManager for StdHsmPal {
     }
 
     fn part_prop_get_u16(&self, io: &impl HsmIo, id: PartPropId) -> HsmResult<u16> {
+        // The partition's RSA-2048 unwrapping key is provisioned lazily
+        // the first time its id is read.  Real PKA hardware generates it
+        // in the background from partition init (reporting
+        // `PendingKeyGeneration` until ready); the emulator generates it
+        // synchronously here so that only a GetUnwrappingKey request pays
+        // the (expensive) keygen cost, never partition enable.
+        if id == PartPropId::RSA_UNWRAPPING_KEY_ID {
+            self.provision_unwrapping_key(io.pid())?;
+        }
         self.prop_get_u16(io, id)
     }
 
@@ -453,52 +462,6 @@ impl HsmPartitionManager for StdHsmPal {
     fn part_prop_clear(&self, io: &impl HsmIo, id: PartPropId) -> HsmResult<()> {
         self.prop_clear(io, id)
     }
-
-    async fn part_ensure_unwrapping_key(&self, io: &impl HsmIo) -> HsmResult<HsmKeyId> {
-        let pid = io.pid();
-        // Fast path: already materialised this enable cycle.
-        if let Some(kid) = self.active_part(pid)?.unwrapping_key_id {
-            return Ok(kid);
-        }
-
-        // Generate the RSA-2048 key pair without holding a borrow into
-        // the partition table across the keygen await.  Only the private
-        // key is persisted in the vault (as DER — the std PAL's
-        // representation; real firmware stores raw key material); the
-        // public key is later derived from it on demand (matching the
-        // reference firmware), so nothing extra is cached in partition
-        // state.
-        let (pk, _pubk) = self.rsa.gen_keypair(2048).await?;
-        let priv_len = pk.to_bytes(None).map_err(|_| HsmError::RsaToDerError)?;
-        let mut priv_buf = vec![0u8; priv_len];
-        pk.to_bytes(Some(&mut priv_buf[..priv_len]))
-            .map_err(|_| HsmError::RsaToDerError)?;
-        let attrs = HsmVaultKeyAttrs::new()
-            .with_internal(true)
-            .with_local(true)
-            .with_unwrap(true);
-
-        // Store the private key in the vault and record its id.  The
-        // `active_part_mut` borrow through to setting `unwrapping_key_id`
-        // is await-free, so on the single-threaded executor it is atomic.
-        // Re-check here: a concurrent caller may have generated and
-        // stored the key while this task was awaiting keygen above.  If
-        // so, drop our freshly generated key (it was never inserted into
-        // the vault) and return the existing id, so the partition only
-        // ever materialises a single unwrapping key.
-        let entry = self.active_part_mut(pid)?;
-        if let Some(kid) = entry.unwrapping_key_id {
-            return Ok(kid);
-        }
-        let kid = entry.vault.create(
-            &priv_buf[..priv_len],
-            HsmVaultKeyKind::Rsa2kPrivate,
-            None,
-            attrs,
-        )?;
-        entry.unwrapping_key_id = Some(kid);
-        Ok(kid)
-    }
 }
 
 // ---------------------------------------------------------------------------
@@ -570,6 +533,49 @@ impl StdHsmPal {
         self.part_if_mut(pid, |s| s != PartState::Unallocated)
     }
 
+    /// Provision the partition's RSA-2048 unwrapping key on demand,
+    /// generating it if it does not yet exist (emulator only).
+    ///
+    /// Real PKA hardware generates this key in the background from
+    /// partition init and reports [`HsmError::PendingKeyGeneration`]
+    /// until it is ready.  The emulator instead generates it lazily and
+    /// synchronously the first time the
+    /// [`RSA_UNWRAPPING_KEY_ID`](PartPropId::RSA_UNWRAPPING_KEY_ID)
+    /// property is read, so only a GetUnwrappingKey request pays the
+    /// keygen cost (never partition enable).  The whole routine is
+    /// `await`-free, so on the single-threaded executor it runs
+    /// atomically — concurrent reads cannot race to generate two keys.
+    ///
+    /// Only the private key is persisted (as DER — the std PAL's
+    /// representation; real firmware stores raw key material); the public
+    /// key is derived from it on demand (matching the reference
+    /// firmware).
+    fn provision_unwrapping_key(&self, pid: HsmPartId) -> HsmResult<()> {
+        if self.active_part(pid)?.unwrapping_key_id.is_some() {
+            return Ok(());
+        }
+
+        let pk = RsaPrivateKey::generate(256).map_err(|_| HsmError::RsaGenerateError)?;
+        let priv_len = pk.to_bytes(None).map_err(|_| HsmError::RsaToDerError)?;
+        let mut priv_buf = vec![0u8; priv_len];
+        pk.to_bytes(Some(&mut priv_buf[..priv_len]))
+            .map_err(|_| HsmError::RsaToDerError)?;
+        let attrs = HsmVaultKeyAttrs::new()
+            .with_internal(true)
+            .with_local(true)
+            .with_unwrap(true);
+
+        let entry = self.active_part_mut(pid)?;
+        let kid = entry.vault.create(
+            &priv_buf[..priv_len],
+            HsmVaultKeyKind::Rsa2kPrivate,
+            None,
+            attrs,
+        )?;
+        entry.unwrapping_key_id = Some(kid);
+        Ok(())
+    }
+
     /// Borrow a partition that is actively serving host traffic.
     ///
     /// "Serving" means [`PartState::Enabled`] or
diff --git a/fw/plat/uno/fw/pal/src/part.rs b/fw/plat/uno/fw/pal/src/part.rs
index ff9e6a3ae..160d1ec48 100644
--- a/fw/plat/uno/fw/pal/src/part.rs
+++ b/fw/plat/uno/fw/pal/src/part.rs
@@ -684,13 +684,4 @@ impl HsmPartitionManager for UnoHsmPal {
         }
         Ok(())
     }
-
-    async fn part_ensure_unwrapping_key(&self, _io: &impl HsmIo) -> HsmResult<HsmKeyId> {
-        // On Uno the RSA-2048 unwrapping key is generated in the
-        // background, started at partition init, and can take a while.
-        // Until that upka-backed generation (and its completion check)
-        // is wired up here, report it as still in progress so the host
-        // retries — mirroring the reference firmware's behaviour.
-        Err(HsmError::PendingKeyGeneration)
-    }
 }

From fc3f54b01b3679276787edb93af7f93d3412baa2 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Wed, 1 Jul 2026 22:26:16 +0000
Subject: [PATCH 4/4] Apply remaining changes

Co-authored-by: jaygmsft <22506014+jaygmsft@users.noreply.github.com>
---
 .../lib/src/ddi/mbor/get_unwrapping_key.rs    |  5 +++-
 fw/pal/traits/src/crypto/rsa.rs               |  7 ++++-
 fw/plat/std/pal/src/drivers/rsa.rs            |  2 +-
 fw/plat/std/pal/src/part.rs                   | 28 +++++++++++--------
 4 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/fw/core/lib/src/ddi/mbor/get_unwrapping_key.rs b/fw/core/lib/src/ddi/mbor/get_unwrapping_key.rs
index 57eb693e4..19f520bdc 100644
--- a/fw/core/lib/src/ddi/mbor/get_unwrapping_key.rs
+++ b/fw/core/lib/src/ddi/mbor/get_unwrapping_key.rs
@@ -81,7 +81,10 @@ pub(crate) async fn get_unwrapping_key<'p, P: HsmPal>(
     // reserved slot — the PAL converts its vault representation into the
     // wire form (incl. any big-endian↔little-endian flip).
     let frame = DdiGetUnwrappingKeyResp::from_layout(resp, &layout);
-    pal.rsa_priv_pub_key(io, priv_key, Some(frame.pub_key.raw))?;
+    let actual_pub_len = pal.rsa_priv_pub_key(io, priv_key, Some(frame.pub_key.raw))?;
+    if actual_pub_len != pub_len {
+        return Err(HsmError::InvalidArg);
+    }
 
     Ok(resp)
 }
diff --git a/fw/pal/traits/src/crypto/rsa.rs b/fw/pal/traits/src/crypto/rsa.rs
index e57862347..84e0d625a 100644
--- a/fw/pal/traits/src/crypto/rsa.rs
+++ b/fw/pal/traits/src/crypto/rsa.rs
@@ -282,7 +282,12 @@ pub trait HsmRsa {
     /// - `Ok(len)` — wire byte length (`modulus_len + 4`); in `Some`
     ///   mode those bytes are written into `pub_out`.
     /// - `Err(HsmError::InvalidArg)` — `priv_key` is not valid vault
-    ///   format or `pub_out` is too small.
+    ///   format.
+    /// - `Err(HsmError::RsaInvalidKeyLength)` — `pub_out` is too small
+    ///   or the derived public-key components do not fit the wire
+    ///   format.
+    /// - `Err(HsmError::RsaGenerateError)` — deriving or reading the
+    ///   public key failed for another reason.
     fn rsa_priv_pub_key(
         &self,
         io: &impl HsmIo,
diff --git a/fw/plat/std/pal/src/drivers/rsa.rs b/fw/plat/std/pal/src/drivers/rsa.rs
index 58a4944b8..73cfb468a 100644
--- a/fw/plat/std/pal/src/drivers/rsa.rs
+++ b/fw/plat/std/pal/src/drivers/rsa.rs
@@ -103,7 +103,7 @@ pub fn rsa_pub_wire(pubk: &RsaPublicKey, out: Option<&mut [u8]>) -> HsmResult<us
     // reverse the whole field to little-endian.
     let e_len = pubk.e(None).map_err(|_| HsmError::RsaGenerateError)?;
     if e_len > RSA_PUB_EXP_WIRE_LEN {
-        return Err(HsmError::RsaGenerateError);
+        return Err(HsmError::RsaInvalidKeyLength);
     }
     let mut e_be = [0u8; RSA_PUB_EXP_WIRE_LEN];
     pubk.e(Some(&mut e_be[RSA_PUB_EXP_WIRE_LEN - e_len..]))
diff --git a/fw/plat/std/pal/src/part.rs b/fw/plat/std/pal/src/part.rs
index 259d6e944..12045fa49 100644
--- a/fw/plat/std/pal/src/part.rs
+++ b/fw/plat/std/pal/src/part.rs
@@ -558,22 +558,26 @@ impl StdHsmPal {
         let pk = RsaPrivateKey::generate(256).map_err(|_| HsmError::RsaGenerateError)?;
         let priv_len = pk.to_bytes(None).map_err(|_| HsmError::RsaToDerError)?;
         let mut priv_buf = vec![0u8; priv_len];
-        pk.to_bytes(Some(&mut priv_buf[..priv_len]))
-            .map_err(|_| HsmError::RsaToDerError)?;
         let attrs = HsmVaultKeyAttrs::new()
             .with_internal(true)
             .with_local(true)
             .with_unwrap(true);
-
-        let entry = self.active_part_mut(pid)?;
-        let kid = entry.vault.create(
-            &priv_buf[..priv_len],
-            HsmVaultKeyKind::Rsa2kPrivate,
-            None,
-            attrs,
-        )?;
-        entry.unwrapping_key_id = Some(kid);
-        Ok(())
+        let result = (|| {
+            pk.to_bytes(Some(&mut priv_buf[..priv_len]))
+                .map_err(|_| HsmError::RsaToDerError)?;
+
+            let entry = self.active_part_mut(pid)?;
+            let kid = entry.vault.create(
+                &priv_buf[..priv_len],
+                HsmVaultKeyKind::Rsa2kPrivate,
+                None,
+                attrs,
+            )?;
+            entry.unwrapping_key_id = Some(kid);
+            Ok(())
+        })();
+        priv_buf.fill(0);
+        result
     }
 
     /// Borrow a partition that is actively serving host traffic.