vsock: fixed rxq logic

This patch has been cherry-picked from the Firecracker tree. The
reference commit is d2475773557c82d2abad2fc8bdf69e7d01444109.

Fixed a vsock muxer issue that would cause a connection to be removed
from the RX queue, even though it still had pending RX data.

Signed-off-by: Dan Horobeanu <dhr@amazon.com>
Signed-off-by: Gabriel Ionescu <gbi@amazon.com>
Signed-off-by: Stefano Garzarella <sgarzare@redhat.com>

vm-virtio/src/vsock/unix/muxer.rs

@@ -60,7 +60,7 @@ pub struct ConnMapKey {
 
 /// A muxer RX queue item.
 ///
-#[derive(Debug)]
+#[derive(Clone, Copy, Debug)]
 pub enum MuxerRx {
     /// The packet must be fetched from the connection identified by `ConnMapKey`.
     ConnRx(ConnMapKey),
@@ -134,7 +134,7 @@ impl VsockChannel for VsockMuxer {
             self.rxq = MuxerRxQ::from_conn_map(&self.conn_map);
         }
 
-        while let Some(rx) = self.rxq.pop() {
+        while let Some(rx) = self.rxq.peek() {
             let res = match rx {
                 // We need to build an RST packet, going from `local_port` to `peer_port`.
                 MuxerRx::RstPkt {
@@ -151,6 +151,7 @@ impl VsockChannel for VsockMuxer {
                         .set_flags(0)
                         .set_buf_alloc(0)
                         .set_fwd_cnt(0);
+                    self.rxq.pop().unwrap();
                     return Ok(());
                 }
 
@@ -158,9 +159,14 @@ impl VsockChannel for VsockMuxer {
                 // to say.
                 MuxerRx::ConnRx(key) => {
                     let mut conn_res = Err(VsockError::NoData);
+                    let mut do_pop = true;
                     self.apply_conn_mutation(key, |conn| {
                         conn_res = conn.recv_pkt(pkt);
+                        do_pop = !conn.has_pending_rx();
                     });
+                    if do_pop {
+                        self.rxq.pop().unwrap();
+                    }
                     conn_res
                 }
             };

vm-virtio/src/vsock/unix/muxer_rxq.rs

@@ -107,6 +107,12 @@ impl MuxerRxQ {
         false
     }
 
+    /// Peek into the front of the queue.
+    ///
+    pub fn peek(&self) -> Option<MuxerRx> {
+        self.q.front().copied()
+    }
+
     /// Pop an RX item from the front of the queue.
     ///
     pub fn pop(&mut self) -> Option<MuxerRx> {