From 49dd355116822d4798f2be1f52e748db6a4b8387 Mon Sep 17 00:00:00 2001
From: Andrew Walbran <qwandor@google.com>
Date: Thu, 30 May 2024 12:46:31 +0100
Subject: [PATCH 1/2] repr(packed) still needs repr(C) to guarantee the correct
 layout.

---
 src/device/socket/protocol.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/device/socket/protocol.rs b/src/device/socket/protocol.rs
index 00ca7c06..71d1543f 100644
--- a/src/device/socket/protocol.rs
+++ b/src/device/socket/protocol.rs
@@ -45,7 +45,7 @@ pub struct VirtioVsockConfig {
 }
 
 /// The message header for data packets sent on the tx/rx queues
-#[repr(packed)]
+#[repr(C, packed)]
 #[derive(AsBytes, Clone, Copy, Debug, Eq, FromBytes, FromZeroes, PartialEq)]
 pub struct VirtioVsockHdr {
     pub src_cid: U64<LittleEndian>,

From e5f9868e42a1668be4015c9ad6570f35712e5928 Mon Sep 17 00:00:00 2001
From: Andrew Walbran <qwandor@google.com>
Date: Thu, 30 May 2024 12:53:31 +0100
Subject: [PATCH 2/2] Allow vsock stream shutdown hints to be specified.

---
 src/device/socket/connectionmanager.rs |  9 ++++++---
 src/device/socket/protocol.rs          | 17 ++++++++++++++++
 src/device/socket/vsock.rs             | 27 +++++++++++++++++++++++---
 3 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/src/device/socket/connectionmanager.rs b/src/device/socket/connectionmanager.rs
index 8690ca30..518b5547 100644
--- a/src/device/socket/connectionmanager.rs
+++ b/src/device/socket/connectionmanager.rs
@@ -252,7 +252,8 @@ impl<H: Hal, T: Transport> VsockConnectionManager<H, T> {
         }
     }
 
-    /// Requests to shut down the connection cleanly.
+    /// Requests to shut down the connection cleanly, telling the peer that we won't send or receive
+    /// any more data.
     ///
     /// This returns as soon as the request is sent; you should wait until `poll` returns a
     /// `VsockEventType::Disconnected` event if you want to know that the peer has acknowledged the
@@ -389,7 +390,9 @@ mod tests {
     use super::*;
     use crate::{
         device::socket::{
-            protocol::{SocketType, VirtioVsockConfig, VirtioVsockHdr, VirtioVsockOp},
+            protocol::{
+                SocketType, StreamShutdown, VirtioVsockConfig, VirtioVsockHdr, VirtioVsockOp,
+            },
             vsock::{VsockBufferStatus, QUEUE_SIZE, RX_QUEUE_IDX, TX_QUEUE_IDX},
         },
         hal::fake::FakeHal,
@@ -557,7 +560,7 @@ mod tests {
                     dst_port: host_port.into(),
                     len: 0.into(),
                     socket_type: SocketType::Stream.into(),
-                    flags: 0.into(),
+                    flags: (StreamShutdown::SEND | StreamShutdown::RECEIVE).into(),
                     buf_alloc: 1024.into(),
                     fwd_cnt: (hello_from_host.len() as u32).into(),
                 }
diff --git a/src/device/socket/protocol.rs b/src/device/socket/protocol.rs
index 71d1543f..4eac3245 100644
--- a/src/device/socket/protocol.rs
+++ b/src/device/socket/protocol.rs
@@ -214,3 +214,20 @@ bitflags! {
         const NOTIFICATION_DATA     = 1 << 38;
     }
 }
+
+bitflags! {
+    /// Flags sent with a shutdown request to hint that the peer won't send or receive more data.
+    #[derive(Copy, Clone, Debug, Default, Eq, PartialEq)]
+    pub struct StreamShutdown: u32 {
+        /// The peer will not receive any more data.
+        const RECEIVE = 1 << 0;
+        /// The peer will not send any more data.
+        const SEND = 1 << 1;
+    }
+}
+
+impl From<StreamShutdown> for U32<LittleEndian> {
+    fn from(flags: StreamShutdown) -> Self {
+        flags.bits().into()
+    }
+}
diff --git a/src/device/socket/vsock.rs b/src/device/socket/vsock.rs
index 84e5df65..e3ba1107 100644
--- a/src/device/socket/vsock.rs
+++ b/src/device/socket/vsock.rs
@@ -2,7 +2,9 @@
 #![deny(unsafe_op_in_unsafe_fn)]
 
 use super::error::SocketError;
-use super::protocol::{Feature, VirtioVsockConfig, VirtioVsockHdr, VirtioVsockOp, VsockAddr};
+use super::protocol::{
+    Feature, StreamShutdown, VirtioVsockConfig, VirtioVsockHdr, VirtioVsockOp, VsockAddr,
+};
 use crate::hal::Hal;
 use crate::queue::VirtQueue;
 use crate::transport::Transport;
@@ -409,19 +411,38 @@ impl<H: Hal, T: Transport> VirtIOSocket<H, T> {
         result
     }
 
-    /// Requests to shut down the connection cleanly.
+    /// Requests to shut down the connection cleanly, sending hints about whether we will send or
+    /// receive more data.
     ///
     /// This returns as soon as the request is sent; you should wait until `poll` returns a
     /// `VsockEventType::Disconnected` event if you want to know that the peer has acknowledged the
     /// shutdown.
-    pub fn shutdown(&mut self, connection_info: &ConnectionInfo) -> Result {
+    pub fn shutdown_with_hints(
+        &mut self,
+        connection_info: &ConnectionInfo,
+        hints: StreamShutdown,
+    ) -> Result {
         let header = VirtioVsockHdr {
             op: VirtioVsockOp::Shutdown.into(),
+            flags: hints.into(),
             ..connection_info.new_header(self.guest_cid)
         };
         self.send_packet_to_tx_queue(&header, &[])
     }
 
+    /// Requests to shut down the connection cleanly, telling the peer that we won't send or receive
+    /// any more data.
+    ///
+    /// This returns as soon as the request is sent; you should wait until `poll` returns a
+    /// `VsockEventType::Disconnected` event if you want to know that the peer has acknowledged the
+    /// shutdown.
+    pub fn shutdown(&mut self, connection_info: &ConnectionInfo) -> Result {
+        self.shutdown_with_hints(
+            connection_info,
+            StreamShutdown::SEND | StreamShutdown::RECEIVE,
+        )
+    }
+
     /// Forcibly closes the connection without waiting for the peer.
     pub fn force_close(&mut self, connection_info: &ConnectionInfo) -> Result {
         let header = VirtioVsockHdr {