From 0ef5821f7dae04138a0679ffc9d802aa66ee25ad Mon Sep 17 00:00:00 2001
From: siweicai <siwei.cai@intel.com>
Date: Wed, 4 Feb 2026 17:08:31 +0800
Subject: [PATCH] Fix Migtd attestation doesn't close transport

---
 src/migtd/src/migration/session.rs | 16 ++++++++++++++--
 src/migtd/src/spdm/mod.rs          |  6 +++---
 src/migtd/src/spdm/spdm_req.rs     |  5 +++--
 src/migtd/src/spdm/spdm_rsp.rs     |  5 +++--
 4 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/src/migtd/src/migration/session.rs b/src/migtd/src/migration/session.rs
index 9f37f458..18103512 100644
--- a/src/migtd/src/migration/session.rs
+++ b/src/migtd/src/migration/session.rs
@@ -921,8 +921,10 @@ async fn migration_src_exchange_msk(
     info: &MigrationInformation,
     #[cfg(feature = "policy_v2")] remote_policy: Vec<u8>,
 ) -> Result<()> {
+    use core::ops::DerefMut;
+
     const SPDM_TIMEOUT: Duration = Duration::from_secs(60); // 60 seconds
-    let mut spdm_requester = spdm::spdm_requester(transport).map_err(|_e| {
+    let (mut spdm_requester, device_io_ref) = spdm::spdm_requester(transport).map_err(|_e| {
         log::error!(
             "exchange_msk(): Failed in spdm_requester transport. Migration ID: {}\n",
             info.mig_info.mig_request_id
@@ -951,6 +953,10 @@ async fn migration_src_exchange_msk(
         e
     })?;
     log::info!("MSK exchange completed\n");
+
+    let mut transport_lock = device_io_ref.lock();
+    let transport = transport_lock.deref_mut();
+    shutdown_transport(&mut transport.transport, info.mig_info.mig_request_id).await?;
     Ok(())
 }
 
@@ -960,8 +966,10 @@ async fn migration_dst_exchange_msk(
     info: &MigrationInformation,
     #[cfg(feature = "policy_v2")] remote_policy: Vec<u8>,
 ) -> Result<()> {
+    use core::ops::DerefMut;
+
     const SPDM_TIMEOUT: Duration = Duration::from_secs(60); // 60 seconds
-    let mut spdm_responder = spdm::spdm_responder(transport).map_err(|_e| {
+    let (mut spdm_responder, device_io_ref) = spdm::spdm_responder(transport).map_err(|_e| {
         log::error!(
             "exchange_msk(): Failed in spdm_responder transport. Migration ID: {}\n",
             info.mig_info.mig_request_id
@@ -991,6 +999,10 @@ async fn migration_dst_exchange_msk(
         e
     })?;
     log::info!("MSK exchange completed\n");
+
+    let mut transport_lock = device_io_ref.lock();
+    let transport = transport_lock.deref_mut();
+    shutdown_transport(&mut transport.transport, info.mig_info.mig_request_id).await?;
     Ok(())
 }
 
diff --git a/src/migtd/src/spdm/mod.rs b/src/migtd/src/spdm/mod.rs
index 5c2f1924..291fb07d 100644
--- a/src/migtd/src/spdm/mod.rs
+++ b/src/migtd/src/spdm/mod.rs
@@ -36,13 +36,13 @@ use crate::migration::MigrationResult;
 use crate::migration::MigtdMigrationInformation;
 use crate::spdm::vmcall_msg::VMCALL_SPDM_MESSAGE_HEADER_SIZE;
 
-pub struct MigtdTransport<T: AsyncRead + AsyncWrite + Unpin> {
+pub(crate) type SpdmDeviceIoArc<T> = Arc<Mutex<MigtdTransport<T>>>;
+pub struct MigtdTransport<T: AsyncRead + AsyncWrite + Unpin + Send> {
     pub transport: T,
 }
-unsafe impl<T: AsyncRead + AsyncWrite + Unpin> Send for MigtdTransport<T> {}
 
 #[async_trait]
-impl<T: AsyncRead + AsyncWrite + Unpin> SpdmDeviceIo for MigtdTransport<T> {
+impl<T: AsyncRead + AsyncWrite + Unpin + Send> SpdmDeviceIo for MigtdTransport<T> {
     async fn send(&mut self, buffer: Arc<&[u8]>) -> SpdmResult {
         let mut sent = 0;
         while sent < buffer.len() {
diff --git a/src/migtd/src/spdm/spdm_req.rs b/src/migtd/src/spdm/spdm_req.rs
index 09b94cb1..c1c6c5e5 100644
--- a/src/migtd/src/spdm/spdm_req.rs
+++ b/src/migtd/src/spdm/spdm_req.rs
@@ -32,9 +32,10 @@ use crate::{
 
 pub fn spdm_requester<T: AsyncRead + AsyncWrite + Unpin + Send + Sync + 'static>(
     stream: T,
-) -> Result<RequesterContext, SpdmStatus> {
+) -> Result<(RequesterContext, SpdmDeviceIoArc<T>), SpdmStatus> {
     let transport = MigtdTransport { transport: stream };
     let device_io = Arc::new(Mutex::new(transport));
+    let device_io_ref = device_io.clone();
 
     let req_capabilities = SpdmRequestCapabilityFlags::ENCRYPT_CAP
         | SpdmRequestCapabilityFlags::MAC_CAP
@@ -79,7 +80,7 @@ pub fn spdm_requester<T: AsyncRead + AsyncWrite + Unpin + Send + Sync + 'static>
 
     spdmlib::secret::asym_sign::register(SECRET_ASYM_IMPL_INSTANCE.clone());
 
-    Ok(requester_context)
+    Ok((requester_context, device_io_ref))
 }
 
 pub async fn spdm_requester_transfer_msk(
diff --git a/src/migtd/src/spdm/spdm_rsp.rs b/src/migtd/src/spdm/spdm_rsp.rs
index 21d7a320..9eaae0ca 100644
--- a/src/migtd/src/spdm/spdm_rsp.rs
+++ b/src/migtd/src/spdm/spdm_rsp.rs
@@ -61,9 +61,10 @@ pub unsafe fn upcast_mut(inner: &mut ResponderContext) -> &mut ResponderContextE
 
 pub fn spdm_responder<T: AsyncRead + AsyncWrite + Unpin + Send + Sync + 'static>(
     stream: T,
-) -> Result<ResponderContextEx, SpdmStatus> {
+) -> Result<(ResponderContextEx, SpdmDeviceIoArc<T>), SpdmStatus> {
     let transport = MigtdTransport { transport: stream };
     let device_io = Arc::new(Mutex::new(transport));
+    let device_io_ref = device_io.clone();
 
     let rsp_capabilities = SpdmResponseCapabilityFlags::ENCRYPT_CAP
         | SpdmResponseCapabilityFlags::MAC_CAP
@@ -120,7 +121,7 @@ pub fn spdm_responder<T: AsyncRead + AsyncWrite + Unpin + Send + Sync + 'static>
         remote_policy: Vec::new(),
     };
 
-    Ok(responder_context_ex)
+    Ok((responder_context_ex, device_io_ref))
 }
 
 pub async fn spdm_responder_transfer_msk(