Merge pull request #136 from rcore-os/shutdown

 Allow vsock stream shutdown hints to be specified

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(),
                 }

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>,
@@ -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()
+    }
+}

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 {