migtd: shutdown vsock stream manually.

Signed-off-by: Jiaqi Gao <[email protected]>
intel · Dec 11, 2024 · 3053a5a · 3053a5a
1 parent 453d539
commit 3053a5a
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/src/crypto/src/rustls_impl/tls.rs b/src/crypto/src/rustls_impl/tls.rs
@@ -44,6 +44,10 @@ where
         SecureChannel { conn }
     }
 
+    pub fn transport_mut(&mut self) -> &mut T {
+        self.conn.transport_mut()
+    }
+
     pub async fn write(&mut self, data: &[u8]) -> Result<usize> {
         self.conn.write(data).await
     }
@@ -89,6 +93,13 @@ impl<T: AsyncRead + AsyncWrite + Unpin> TlsConnection<T> {
             _ => Error::TlsStream,
         }
     }
+
+    fn transport_mut(&mut self) -> &mut T {
+        match self {
+            Self::Server(conn) => &mut conn.transport,
+            Self::Client(conn) => &mut conn.transport,
+        }
+    }
 }
 
 pub struct TlsConfig {

diff --git a/src/migtd/src/migration/session.rs b/src/migtd/src/migration/session.rs
@@ -336,6 +336,8 @@ pub async fn exchange_msk(info: &MigrationInformation) -> Result<()> {
         if size < size_of::<ExchangeInformation>() {
             return Err(MigrationResult::NetworkError);
         }
+        #[cfg(not(feature = "virtio-serial"))]
+        ratls_client.transport_mut().shutdown().await?;
     } else {
         log::info!("Start a TLS server for migration destination\n");
         // TLS server
@@ -355,6 +357,8 @@ pub async fn exchange_msk(info: &MigrationInformation) -> Result<()> {
         if size < size_of::<ExchangeInformation>() {
             return Err(MigrationResult::NetworkError);
         }
+        #[cfg(not(feature = "virtio-serial"))]
+        ratls_server.transport_mut().shutdown().await?;
     }
 
     let mig_ver = cal_mig_version(info.is_src(), &exchange_information, &remote_information)?;