fw(uno): Live migration support for mbor commands

fw/core/lib/src/ddi/mbor/open_session.rs

@@ -101,23 +101,38 @@ pub(crate) async fn open_session<'p, P: HsmPal>(
     let mk_session = pal.dma_alloc(io, BK_LEN)?;
     pal.rng_fill_bytes(io, mk_session)?;
 
+    // Resolve the SVNs *before* allocating the slot so the only fallible step
+    // left after `session_create` is the response encode, which we roll back.
+    let mfgr_svn = crate::part_state::part_mfgr_svn(pal);
+    let owner_svn =
+        u16::try_from(crate::part_state::part_owner_svn(pal)).map_err(|_| HsmError::InvalidArg)?;
+
     let api_rev_bytes = pack_api_rev(api_rev);
     let sess_id = pal
         .session_create(io, &api_rev_bytes, mk_session, None)
         .await?;
 
     // ── Step 9: Encode response + envelope MK_SESSION under BK_SESSION
-    let resp = encode_response(
-        pal,
-        io,
-        hdr,
-        sess_id,
-        bk_session,
-        mk_session,
-        crate::part_state::part_mfgr_svn(pal),
-        u16::try_from(crate::part_state::part_owner_svn(pal)).map_err(|_| HsmError::InvalidArg)?,
+    //
+    // The slot is now allocated. If the encode (crypto + DMA + MBOR) fails,
+    // roll the slot back — otherwise the logical slot and its session-vault
+    // key stay allocated forever: `sess_id` is never surfaced to the CQE, so
+    // the host never registers the session and can never `CloseSession` it,
+    // leaking a slot for the rest of the boot. This mirrors the reference
+    // firmware's FSM `rollback_open_session` backstop.
+    let resp = match encode_response(
+        pal, io, hdr, sess_id, bk_session, mk_session, mfgr_svn, owner_svn,
     )
-    .await?;
+    .await
+    {
+        Ok(resp) => resp,
+        Err(err) => {
+            // Best-effort teardown of the just-allocated slot; surface the
+            // original failure to the host regardless of teardown result.
+            let _ = pal.session_destroy(io, sess_id).await;
+            return Err(err);
+        }
+    };
 
     // Surface the new session id to the IO layer for the CQE (only on the
     // success path, since the `?` above returns early on failure), letting the

fw/core/lib/src/ddi/mbor/reopen_session.rs

@@ -96,26 +96,40 @@ pub(crate) async fn reopen_session<'p, P: HsmPal>(
         &body.bmk_session[layout.plaintext_offset..layout.plaintext_offset + BK_LEN],
     );
 
+    // Re-envelope MK_SESSION under BK_SESSION for the host to persist
+    // (the SVN etc. recorded in the metadata may have advanced), mirroring
+    // OpenSession's response.  Resolve the SVNs *before* re-keying the slot
+    // so the only fallible step left after `session_create` is the response
+    // encode, which the rollback arm below covers — otherwise a `try_from`
+    // failure here would return early and leak the just-recreated slot.
+    let mfgr_svn = crate::part_state::part_mfgr_svn(pal);
+    let owner_svn =
+        u16::try_from(crate::part_state::part_owner_svn(pal)).map_err(|_| HsmError::InvalidArg)?;
+
     // Recreate the migrated session's own slot with the recovered key.
     let api_rev_bytes = pack_api_rev(api_rev);
     let sess_id = pal
         .session_create(io, &api_rev_bytes, mk_session, Some(reopen_id))
         .await?;
 
-    // Re-envelope MK_SESSION under BK_SESSION for the host to persist
-    // (the SVN etc. recorded in the metadata may have advanced), mirroring
-    // OpenSession's response.
-    let resp = encode_reopen_response(
-        pal,
-        io,
-        hdr,
-        sess_id,
-        bk_session,
-        mk_session,
-        crate::part_state::part_mfgr_svn(pal),
-        u16::try_from(crate::part_state::part_owner_svn(pal)).map_err(|_| HsmError::InvalidArg)?,
+    let resp = match encode_reopen_response(
+        pal, io, hdr, sess_id, bk_session, mk_session, mfgr_svn, owner_svn,
     )
-    .await?;
+    .await
+    {
+        Ok(resp) => resp,
+        // The slot was already re-keyed by `session_create` above (flipped
+        // from NeedsRenegotiation to Active with a fresh masking key).  If
+        // the response cannot be encoded, the host never learns the reopen
+        // succeeded and will retry it — but a now-Active slot would reject
+        // that retry and the fresh vault key would orphan.  Restore the
+        // slot to NeedsRenegotiation and drop the fresh key, mirroring the
+        // reference firmware's `rollback_open_session` reopen branch.
+        Err(err) => {
+            let _ = pal.session_rollback_reopen(io, sess_id).await;
+            return Err(err);
+        }
+    };
 
     // `ReopenSession` is an in-session command: it reuses the caller's
     // existing session id (echoed in the response header), so unlike

fw/pal/traits/src/session.rs

@@ -301,6 +301,42 @@ pub trait HsmSessionManager {
     ///   session in the caller's partition.
     async fn session_destroy(&self, io: &impl HsmIo, id: HsmSessId) -> HsmResult<()>;
 
+    /// Rolls back an in-place re-key performed by
+    /// [`session_create`](Self::session_create) with `id == Some(_)`
+    /// (the post-migration `ReopenSession` path).
+    ///
+    /// A `ReopenSession` re-keys a
+    /// [`NeedsRenegotiation`](HsmSessionState::NeedsRenegotiation) slot:
+    /// it installs a fresh masking-key vault entry and clears the slot's
+    /// renegotiation flag, flipping it to [`Active`](HsmSessionState::Active).
+    /// If a later step of the same command fails (e.g. the response
+    /// cannot be encoded), the host never learns the session was
+    /// re-keyed and will retry the reopen — but the now-`Active` slot
+    /// would reject that retry.  This call undoes the re-key so the slot
+    /// returns to a freshly-migrated `NeedsRenegotiation` state:
+    ///
+    /// 1. Deletes the fresh masking-key vault entry that
+    ///    [`session_create`](Self::session_create) installed.
+    /// 2. Restores the slot's renegotiation flag and clears its physical
+    ///    vault mapping.
+    ///
+    /// The slot itself is **not** freed (it still represents a migrated
+    /// session awaiting renegotiation), mirroring the reference
+    /// firmware's `rollback_open_session` reopen branch
+    /// (`session_table().rollback_recreation(id)`).
+    ///
+    /// # Parameters
+    ///
+    /// - `io` — caller's I/O context (partition scope).
+    /// - `id` — session that was just re-keyed and must be restored.
+    ///
+    /// # Returns
+    ///
+    /// - `Ok(())` on success.
+    /// - `Err(HsmError::SessionNotFound)` — `id` does not refer to a
+    ///   live slot in the caller's partition.
+    async fn session_rollback_reopen(&self, io: &impl HsmIo, id: HsmSessId) -> HsmResult<()>;
+
     /// Queries the lifecycle state of a session slot.
     ///
     /// This is an infallible probe: an unknown or freed slot is

fw/plat/std/pal/src/drivers/session.rs

@@ -268,6 +268,24 @@ impl SessionTable {
         Ok(id)
     }
 
+    /// Roll back a [`recreate`](Self::recreate): restore the
+    /// renegotiation bit and clear the physical mapping.
+    ///
+    /// The caller is responsible for deleting the orphaned vault entry
+    /// (via [`physical_id`](Self::physical_id)) before invoking this, so
+    /// the physical mapping is read while still intact.
+    ///
+    /// Leaves the slot as a freshly-migrated
+    /// [`NeedsRenegotiation`](HsmSessionState::NeedsRenegotiation) slot
+    /// so the host can retry a failed `ReopenSession` (mirrors the
+    /// reference firmware's `session_table().rollback_recreation(id)`).
+    pub fn rollback_recreation(&mut self, id: HsmSessId) -> HsmResult<()> {
+        let slot = self.active_slot(id)?;
+        self.renego_mask |= 1 << slot;
+        self.phys_ids[slot] = 0;
+        Ok(())
+    }
+
     /// Query the current state of a session slot.
     pub fn state(&self, id: HsmSessId) -> HsmSessionState {
         let Ok(slot) = self.active_slot(id) else {

fw/plat/std/pal/src/session.rs

@@ -75,20 +75,21 @@ impl HsmSessionManager for StdHsmPal {
         let pid = io.pid();
         let entry = self.active_part_mut(pid)?;
 
-        // On re-key: clean up the old session-scoped keys and the old
-        // session key before creating the replacement.  A slot awaiting
-        // renegotiation after a live-migration disable has already had
-        // its vault key material cleared, so there is nothing to tear
-        // down — `recreate` below simply installs the fresh mapping
-        // (mirroring the reference firmware's `recreate_session`).
+        // Validate the reopen target.  The only valid target is a slot
+        // awaiting renegotiation after a live-migration disable; `restore()`
+        // already zeroed its physical vault mapping, so there is nothing to
+        // tear down and `recreate` below simply installs the fresh mapping
+        // (mirroring the reference firmware's `recreate_session`).  Every
+        // other state (Active, Pending, Invalid) is rejected up front with no
+        // side effects: tearing down the old keys here destroyed a live
+        // session and leaked its slot, because `recreate` then rejects the
+        // non-renegotiation slot anyway.
         if let Some(reopen_id) = id {
             if !matches!(
                 entry.session_table.state(reopen_id),
                 HsmSessionState::NeedsRenegotiation
             ) {
-                let old_phys = entry.session_table.physical_id(reopen_id)?;
-                entry.vault.delete_by_session_key(old_phys)?;
-                entry.vault.delete(old_phys)?;
+                return Err(HsmError::InvalidArg);
             }
         }
 
@@ -157,6 +158,26 @@ impl HsmSessionManager for StdHsmPal {
         Ok(())
     }
 
+    /// Roll back a failed post-migration re-key (`ReopenSession`).
+    ///
+    /// Restores the slot to [`NeedsRenegotiation`](HsmSessionState::NeedsRenegotiation)
+    /// and deletes the fresh masking-key vault entry (plus any
+    /// session-scoped keys bound to it) that
+    /// [`session_create`](Self::session_create) installed, so neither the
+    /// slot state nor the vault leaks and the host can retry the reopen.
+    async fn session_rollback_reopen(&self, io: &impl HsmIo, id: HsmSessId) -> HsmResult<()> {
+        let entry = self.active_part_mut(io.pid())?;
+
+        // Delete the fresh masking-key vault entry that the failed reopen
+        // installed, then restore the slot to NeedsRenegotiation (clearing
+        // its physical mapping) so neither the vault nor the slot leaks and
+        // the host can retry the reopen.
+        let physical_id = entry.session_table.physical_id(id)?;
+        entry.vault.delete(physical_id)?;
+        entry.session_table.rollback_recreation(id)?;
+        Ok(())
+    }
+
     /// Query the lifecycle state of a session.
     fn session_state(&self, io: &impl HsmIo, id: HsmSessId) -> HsmSessionState {
         let Ok(entry) = self.active_part(io.pid()) else {

fw/plat/uno/fw/drivers/iic/src/api.rs

@@ -26,8 +26,22 @@ struct ChannelState {
     /// ICQ consumer index — FW advances after reading completions.
     icq_head: u16,
 
+    /// Free-running count of ICQ entries consumed (both completed and
+    /// failed-recycled). Unlike [`icq_head`](Self::icq_head) this never
+    /// wraps at the ring boundary, so a consumer can latch a target value
+    /// and wait for the head to advance past a fixed point. Used by the
+    /// live-migration SQ-delete drain to account for entries the hardware
+    /// has already accepted into the ICQ but that `recv` has not yet pulled
+    /// into the in-flight IO accounting.
+    icq_consumed: u32,
+
     /// Waker for async recv — woken by IRQ or polling.
     waker: WakerRegistration,
+
+    /// Waker for a task tracking ICQ-consumption progress (the LM SQ-delete
+    /// drain). Woken on every consumed entry so the drain re-polls as the IO
+    /// task empties the shared ICQ, independent of the `recv` waker.
+    drain_waker: WakerRegistration,
 }
 
 /// Async IIC (Inbound IO Controller) driver.
@@ -128,7 +142,9 @@ impl<const DEPTH: usize> IicDriver<DEPTH> {
             io_meta: config.io_meta_base as *mut IoMetaEntry,
             state: SingleCell::new(ChannelState {
                 icq_head: 0,
+                icq_consumed: 0,
                 waker: WakerRegistration::new(),
+                drain_waker: WakerRegistration::new(),
             }),
             regs,
             config,
@@ -284,6 +300,13 @@ impl<const DEPTH: usize> IicDriver<DEPTH> {
                     .head
                     .write(ICQ_CHANNEL_HEAD::HEAD.val(s.icq_head.into()));
 
+                // Advance the free-running consumer count and wake any SQ-delete
+                // drain that is waiting for the ICQ to quiesce. This covers both
+                // the success path below and the failed-recycle path, so a drain
+                // makes progress even when the only remaining entries are errors.
+                s.icq_consumed = s.icq_consumed.wrapping_add(1);
+                s.drain_waker.wake();
+
                 // Skip failed entries — recycle buffer, return credit, and keep polling
                 if !success {
                     let _cause = self.regs.interrupt_cause.get();
@@ -315,6 +338,31 @@ impl<const DEPTH: usize> IicDriver<DEPTH> {
         })
     }
 
+    /// Free-running count of ICQ entries consumed since boot.
+    ///
+    /// Monotonic (modulo `u32` wrap) across the ring boundary, so a waiter
+    /// can latch a value and detect when the consumer has advanced past it.
+    pub fn icq_consumed(&self) -> u32 {
+        self.state.with(|s| s.icq_consumed)
+    }
+
+    /// Number of ICQ entries the hardware has accepted but `recv` has not yet
+    /// consumed (ring distance between the shadow tail and the FW head).
+    pub fn icq_pending(&self) -> u16 {
+        self.state.with(|s| {
+            let tail = unsafe { self.icq_tail_shadow.read_volatile() } as u16;
+            tail.wrapping_sub(s.icq_head) & Self::MASK
+        })
+    }
+
+    /// Register `waker` to be notified on every ICQ consumption.
+    ///
+    /// Used by the live-migration SQ-delete drain to re-poll as the IO task
+    /// empties the shared ICQ, independent of the `recv` waker.
+    pub fn register_drain_waker(&self, waker: &core::task::Waker) {
+        self.state.with(|s| s.drain_waker.register(waker));
+    }
+
     /// Wake the driver if the ICQ has pending entries.
     pub fn wake(&self, irq: u16) {
         self.state.with(|s| {

fw/plat/uno/fw/drivers/part_store/src/part_store.rs

@@ -679,19 +679,15 @@ impl Partition {
         unsafe { DmaBuf::from_raw_mut(&mut self.slot_mut().vm_launch_guid) }
     }
 
-    /// Sets the VM-launch GUID.
+    /// Captures the host-provided VM-launch GUID.
     ///
-    /// # Errors
-    ///
-    /// - [`HsmError::InvalidArg`] — `v` is not exactly `GUID_LEN` bytes.
+    /// The GUID arrives as plain bytes in the `SetResource` IPC payload
+    /// (not DMA-sourced), so this takes a fixed-size array rather than a
+    /// [`DmaBuf`]. Mirrors the reference firmware's capture in
+    /// `part_init::handle_set_res_cnt`.
     #[inline(never)]
-    pub fn set_vm_launch_guid(mut self, v: &DmaBuf) -> HsmResult<()> {
-        let src: &[u8] = v;
-        if src.len() != GUID_LEN {
-            return Err(HsmError::InvalidArg);
-        }
-        self.slot_mut().vm_launch_guid.copy_from_slice(src);
-        Ok(())
+    pub fn set_vm_launch_guid(mut self, v: &[u8; GUID_LEN]) {
+        self.slot_mut().vm_launch_guid.copy_from_slice(v);
     }
 
     // ── masked boot key (variable length) ────────────────────────────────
@@ -752,9 +748,6 @@ impl Partition {
             return Err(HsmError::InvalidArg);
         }
         let slot = self.slot_mut();
-        if slot.sealed_bk3.len != 0 {
-            return Err(HsmError::SealedBk3AlreadySet);
-        }
         slot.sealed_bk3.len = src.len() as u32;
         slot.sealed_bk3.data[..src.len()].copy_from_slice(src);
         // Clear the unused tail of the storage slot.