fix(migtd): harden virtqueue, vsock, SPDM and policy against untruste…

deps/td-shim-AzCVMEmu/td-payload-emu/src/mm/shared/mod.rs

@@ -10,6 +10,7 @@ const PAGE_SIZE: usize = 0x1000;
 
 pub struct SharedMemory {
     buf: Vec<u8>,
+    shadow: Vec<u8>,
 }
 
 impl SharedMemory {
@@ -21,6 +22,7 @@ impl SharedMemory {
         let size = pages.checked_mul(4096)?;
         Some(Self {
             buf: Vec::from_iter(core::iter::repeat(0u8).take(size)),
+            shadow: Vec::new(),
         })
     }
 
@@ -34,9 +36,11 @@ impl SharedMemory {
     }
 
     pub fn copy_to_private_shadow(&mut self) -> Option<&[u8]> {
-        // In emulation mode, just return the buffer directly since we're not dealing with
-        // actual shared/private memory conversion like in real TDX
-        Some(&self.buf)
+        // Copy into a separate shadow buffer so the returned slice is an immutable
+        // snapshot, matching the real td-payload shared->private memcpy semantics.
+        self.shadow.clear();
+        self.shadow.extend_from_slice(&self.buf);
+        Some(&self.shadow)
     }
 }
 

src/devices/virtio/src/virtqueue.rs

@@ -47,6 +47,20 @@ pub struct VirtQueue {
     free_head: u16,
     avail_idx: u16,
     last_used_idx: u16,
+
+    /// Private shadow of the descriptor `next` links. This is the authoritative
+    /// source for free-list traversal and chain recycling; the host-writable
+    /// `desc[].next` field in shared DMA is never trusted to drive guest-side
+    /// accounting.
+    desc_next: [u16; MAX_QUEUE_SIZE],
+    /// Private shadow of each descriptor's direction (true = device-writable,
+    /// i.e. host-to-guest). Recorded at `add` time so recycling does not rely on
+    /// the host-writable `desc[].flags`.
+    desc_write: [bool; MAX_QUEUE_SIZE],
+    /// Number of descriptors in the in-flight chain headed by each index. A
+    /// value of 0 means the index is not the head of an in-flight chain, so a
+    /// host-reported used id that is not a real chain head is rejected.
+    chain_len: [u16; MAX_QUEUE_SIZE],
 }
 
 impl VirtQueue {
@@ -82,8 +96,10 @@ impl VirtQueue {
             unsafe { &mut *((dma_addr as usize + layout.used_offset) as *mut UsedRing) };
 
         // link descriptors together
+        let mut desc_next = [0u16; MAX_QUEUE_SIZE];
         for i in 0..(queue_size - 1) {
             desc[i as usize].next.write(i + 1);
+            desc_next[i as usize] = i + 1;
         }
 
         Ok(VirtQueue {
@@ -96,6 +112,9 @@ impl VirtQueue {
             free_head: 0,
             avail_idx: 0,
             last_used_idx: 0,
+            desc_next,
+            desc_write: [false; MAX_QUEUE_SIZE],
+            chain_len: [0; MAX_QUEUE_SIZE],
         })
     }
 
@@ -129,7 +148,11 @@ impl VirtQueue {
         flags.remove(DescFlags::NEXT);
         desc.flags.write(flags);
 
-        self.num_used += (g2h.len() + h2g.len()) as u16;
+        let chain_len = (g2h.len() + h2g.len()) as u16;
+        self.num_used += chain_len;
+        // Record the submitted chain length privately so recycling validates the
+        // host-reported used id against a chain we actually built.
+        self.chain_len[head as usize] = chain_len;
 
         let avail_slot = self.avail_idx & (self.queue_size - 1);
         self.avail
@@ -160,61 +183,82 @@ impl VirtQueue {
     }
 
     fn add_descriptor(&mut self, buf: &VirtqueueBuf, flag: DescFlags) -> Result<u16> {
+        let index = self.free_head;
         let desc = self
             .desc
-            .get_mut(self.free_head as usize)
+            .get_mut(index as usize)
             .ok_or(VirtioError::InvalidDescriptorIndex)?;
         desc.set_buf(buf);
         desc.flags.write(flag);
 
-        // Update the free head
-        let last = self.free_head;
-        self.free_head = desc.next.read();
+        // Advance the free head using the private shadow link instead of the
+        // host-writable `desc.next`, and record the buffer direction privately.
+        self.free_head = self.desc_next[index as usize];
+        self.desc_write[index as usize] = flag.contains(DescFlags::WRITE);
 
-        Ok(last)
+        Ok(index)
     }
 
     /// Recycle descriptors in the list specified by head.
     ///
     /// This will push all linked descriptors at the front of the free list.
     fn recycle_descriptors(
         &mut self,
-        mut head: u16,
+        head: u16,
         g2h: &mut Vec<VirtqueueBuf>,
         h2g: &mut Vec<VirtqueueBuf>,
     ) -> Result<()> {
+        // Validate, against private state, that `head` (reported by the host via
+        // the used ring) is the head of a chain we actually submitted. The walk
+        // below is driven entirely by the private shadow, never by the
+        // host-writable `desc.next`/`desc.flags`, so the host cannot redirect it
+        // to never-submitted slots or leak in-flight descriptors.
+        let len = self
+            .chain_len
+            .get(head as usize)
+            .copied()
+            .ok_or(VirtioError::InvalidDescriptorIndex)?;
+        if len == 0 || len > self.num_used {
+            return Err(VirtioError::InvalidDescriptor);
+        }
+
         let origin_free_head = self.free_head;
-        self.free_head = head;
+        let mut cur = head;
+        for i in 0..len {
+            let index = cur as usize;
+            let next = self.desc_next[index];
+            let is_write = self.desc_write[index];
+            let is_last = i + 1 == len;
 
-        while self.num_used > 0 {
             let desc = self
                 .desc
-                .get_mut(head as usize)
+                .get_mut(index)
                 .ok_or(VirtioError::InvalidDescriptorIndex)?;
-
-            let flags = desc.flags.read();
-            self.num_used -= 1;
-
             let addr = desc.addr.read();
-            let len = desc.len.read();
+            let buf_len = desc.len.read();
+            if is_last {
+                // Relink the tail of the recycled chain to the previous free
+                // head in the shared table so the device still observes a
+                // consistent free list.
+                desc.next.write(origin_free_head);
+            }
 
-            if flags.contains(DescFlags::WRITE) {
-                h2g.push(VirtqueueBuf::new(addr, len));
+            if is_write {
+                h2g.push(VirtqueueBuf::new(addr, buf_len));
             } else {
-                g2h.push(VirtqueueBuf::new(addr, len));
+                g2h.push(VirtqueueBuf::new(addr, buf_len));
             }
 
-            if flags.contains(DescFlags::NEXT) {
-                if self.num_used == 0 {
-                    return Err(VirtioError::InvalidDescriptor);
-                }
-                head = desc.next.read();
-            } else {
-                desc.next.write(origin_free_head);
-                break;
+            if is_last {
+                self.desc_next[index] = origin_free_head;
             }
+            cur = next;
         }
 
+        self.free_head = head;
+        self.num_used -= len;
+        self.chain_len[head as usize] = 0;
+
         Ok(())
     }
 

src/devices/vsock/src/stream.rs

@@ -109,11 +109,11 @@ impl VsockStream {
     }
 
     pub fn bind(&mut self, addr: &VsockAddr) -> Result {
-        if USED_PORT.lock().contains(&addr.port()) {
+        let mut used_port = USED_PORT.lock();
+        if !used_port.insert(addr.port()) {
             return Err(VsockError::AddressAlreadyUsed);
         }
-
-        USED_PORT.lock().insert(addr.port());
+        drop(used_port);
         self.addr.local.set_port(addr.port);
         Ok(())
     }
@@ -457,14 +457,16 @@ lazy_static! {
 }
 
 pub fn get_unused_port() -> Option<u32> {
-    let mut port = UNUSED_PORT_COUNTER.lock().checked_add(1)?;
+    let mut counter = UNUSED_PORT_COUNTER.lock();
+    let mut used_port = USED_PORT.lock();
+    let mut port = counter.checked_add(1)?;
 
-    while USED_PORT.lock().contains(&port) {
+    while used_port.contains(&port) {
         port = port.checked_add(1)?;
     }
 
-    USED_PORT.lock().insert(port);
-    *UNUSED_PORT_COUNTER.lock() = port;
+    used_port.insert(port);
+    *counter = port;
 
     Some(port)
 }

src/devices/vsock/src/transport/vmcall.rs

@@ -192,12 +192,20 @@ async fn vmcall_service_migtd_send(
             return Poll::Pending;
         }
 
+        // The VMM controls the response length; reject any response whose data
+        // section is too short to hold the fields parsed below, otherwise the
+        // indexing into `reply.data()` would panic.
+        let reply_data = reply.data();
+        if reply_data.len() < 12 {
+            return Poll::Ready(Err(VsockTransportError::InvalidParameter));
+        }
+
         // Do the sanity check
         if reply.guid() != VMCALL_SERVICE_MIGTD_GUID.as_bytes()
             || reply.status() != 0
-            || reply.data()[0] != CURRENT_VERSION
-            || reply.data()[1] != COMMAND_SEND
-            || u64::from_le_bytes(reply.data()[4..12].try_into().unwrap()) != mid
+            || reply_data[0] != CURRENT_VERSION
+            || reply_data[1] != COMMAND_SEND
+            || u64::from_le_bytes(reply_data[4..12].try_into().unwrap()) != mid
         {
             return Poll::Ready(Err(VsockTransportError::InvalidParameter));
         }
@@ -230,17 +238,25 @@ async fn vmcall_service_migtd_receive(
             return Poll::Pending;
         }
 
+        // The VMM controls the response length; reject any response whose data
+        // section is too short to hold the fields parsed below, otherwise the
+        // indexing into `reply.data()` would panic.
+        let reply_data = reply.data();
+        if reply_data.len() < 12 {
+            return Poll::Ready(Err(VsockTransportError::InvalidParameter));
+        }
+
         // Do the sanity check
         if reply.guid() != VMCALL_SERVICE_MIGTD_GUID.as_bytes()
             || reply.status() != 0
-            || reply.data()[0] != CURRENT_VERSION
-            || reply.data()[1] != COMMAND_RECV
-            || u64::from_le_bytes(reply.data()[4..12].try_into().unwrap()) != mid
+            || reply_data[0] != CURRENT_VERSION
+            || reply_data[1] != COMMAND_RECV
+            || u64::from_le_bytes(reply_data[4..12].try_into().unwrap()) != mid
         {
             return Poll::Ready(Err(VsockTransportError::InvalidParameter));
         }
 
-        recv_packet(&reply.data()[12..])?;
+        recv_packet(&reply_data[12..])?;
         Poll::Ready(pop_stream_queues(addrs).ok_or(VsockTransportError::InvalidVsockPacket))
     })
     .await

src/migtd/src/migration/session.rs

@@ -585,11 +585,10 @@ pub async fn wait_for_request() -> Result<MigrationInformation> {
                     .ok_or(MigrationResult::InvalidParameter)?;
             let request_id = mig_info.mig_info.mig_request_id;
 
-            if REQUESTS.lock().contains(&request_id) {
-                Poll::Pending
-            } else {
-                REQUESTS.lock().insert(request_id);
+            if REQUESTS.lock().insert(request_id) {
                 Poll::Ready(Ok(mig_info))
+            } else {
+                Poll::Pending
             }
         } else if wfr.operation == 0 {
             Poll::Pending

src/migtd/src/spdm/mod.rs

@@ -85,7 +85,7 @@ impl<T: AsyncRead + AsyncWrite + Unpin + Send> SpdmDeviceIo for MigtdTransport<T
         while recvd < VMCALL_SPDM_MESSAGE_HEADER_SIZE {
             let n = self
                 .transport
-                .read(&mut buffer[recvd..])
+                .read(&mut buffer[recvd..VMCALL_SPDM_MESSAGE_HEADER_SIZE])
                 .await
                 // ConnectionAborted maps from VmcallRawError::VmmCanceled.
                 // The SpdmDeviceIo trait cannot represent this error, so just
@@ -108,10 +108,11 @@ impl<T: AsyncRead + AsyncWrite + Unpin + Send> SpdmDeviceIo for MigtdTransport<T
             return Err(0_usize);
         }
 
-        while recvd < payload_size + VMCALL_SPDM_MESSAGE_HEADER_SIZE {
+        let total = payload_size + VMCALL_SPDM_MESSAGE_HEADER_SIZE;
+        while recvd < total {
             let n = self
                 .transport
-                .read(&mut buffer[recvd..])
+                .read(&mut buffer[recvd..total])
                 .await
                 .map_err(|e| {
                     if e.kind() == rust_std_stub::io::ErrorKind::ConnectionAborted {
@@ -122,7 +123,7 @@ impl<T: AsyncRead + AsyncWrite + Unpin + Send> SpdmDeviceIo for MigtdTransport<T
             recvd += n;
         }
 
-        Ok(recvd)
+        Ok(total)
     }
 
     async fn flush_all(&mut self) -> SpdmResult {