Wrap the pre-session v2 policy change with timeout

Wrap the pre-session v2 policy change with a timeout to prevent potential
    hang during pre migration.

Author: IntelCaisui

src/migtd/src/migration/session.rs

@@ -13,6 +13,7 @@ use async_io::{AsyncRead, AsyncWrite};
 use core::sync::atomic::AtomicBool;
 #[cfg(any(feature = "vmcall-interrupt", feature = "vmcall-raw"))]
 use core::sync::atomic::Ordering;
+use core::time::Duration;
 use core::{future::poll_fn, mem::size_of, task::Poll};
 #[cfg(any(feature = "vmcall-interrupt", feature = "vmcall-raw"))]
 use event::VMCALL_SERVICE_FLAG;
@@ -847,12 +848,16 @@ pub async fn exchange_msk(info: &MigrationInformation, data: &mut Vec<u8>) -> Re
 
     // Exchange policy firstly because of the message size limitation of TLS protocol
     #[cfg(feature = "policy_v2")]
-    let remote_policy = Box::pin(pre_session_data_exchange(&mut transport)).await?;
+    const PRE_SESSION_TIMEOUT: Duration = Duration::from_secs(60); // 60 seconds
+    #[cfg(feature = "policy_v2")]
+    let remote_policy = Box::pin(with_timeout(
+        PRE_SESSION_TIMEOUT,
+        pre_session_data_exchange(&mut transport),
+    ))
+    .await??;
 
     #[cfg(not(feature = "spdm_attestation"))]
     {
-        use core::time::Duration;
-
         const TLS_TIMEOUT: Duration = Duration::from_secs(60); // 60 seconds
 
         let mut remote_information = ExchangeInformation::default();
@@ -986,12 +991,15 @@ pub async fn exchange_msk(info: &MigrationInformation, data: &mut Vec<u8>) -> Re
         remote_information.key.clear();
     }
 
+    #[cfg(feature = "spdm_attestation")]
+    const SPDM_TIMEOUT: Duration = Duration::from_secs(60); // 60 seconds
+
     #[cfg(feature = "spdm_attestation")]
     if info.is_src() {
         let mut spdm_requester =
             spdm::spdm_requester(transport).map_err(|_| MigrationResult::SecureSessionError)?;
         with_timeout(
-            spdm::SPDM_TIMEOUT,
+            SPDM_TIMEOUT,
             spdm::spdm_requester_transfer_msk(
                 &mut spdm_requester,
                 &info.mig_info,
@@ -1006,7 +1014,7 @@ pub async fn exchange_msk(info: &MigrationInformation, data: &mut Vec<u8>) -> Re
             spdm::spdm_responder(transport).map_err(|_| MigrationResult::SecureSessionError)?;
 
         with_timeout(
-            spdm::SPDM_TIMEOUT,
+            SPDM_TIMEOUT,
             spdm::spdm_responder_transfer_msk(
                 &mut spdm_responder,
                 &info.mig_info,

src/migtd/src/spdm/mod.rs

@@ -16,7 +16,6 @@ use async_trait::async_trait;
 use codec::Codec;
 use codec::Reader;
 use codec::Writer;
-use core::time::Duration;
 use spdmlib::common::SpdmDeviceIo;
 use spdmlib::error::*;
 use spin::Mutex;
@@ -37,8 +36,6 @@ use crate::migration::MigrationResult;
 use crate::migration::MigtdMigrationInformation;
 use crate::spdm::vmcall_msg::VMCALL_SPDM_MESSAGE_HEADER_SIZE;
 
-pub const SPDM_TIMEOUT: Duration = Duration::from_secs(60); // 60 seconds
-
 pub struct MigtdTransport<T: AsyncRead + AsyncWrite + Unpin> {
     pub transport: T,
 }