From 5fc52a50568dc6be21ae547b6f385afa843d1707 Mon Sep 17 00:00:00 2001
From: Stefano Garzarella <sgarzare@redhat.com>
Date: Tue, 16 Jun 2020 12:00:07 +0200
Subject: [PATCH] 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     | 10 ++++++++--
 vm-virtio/src/vsock/unix/muxer_rxq.rs |  6 ++++++
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/vm-virtio/src/vsock/unix/muxer.rs b/vm-virtio/src/vsock/unix/muxer.rs
index 9cf1af9bf..b2f57c285 100644
--- a/vm-virtio/src/vsock/unix/muxer.rs
+++ b/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
                 }
             };
diff --git a/vm-virtio/src/vsock/unix/muxer_rxq.rs b/vm-virtio/src/vsock/unix/muxer_rxq.rs
index 79f07494a..3b02b54b1 100644
--- a/vm-virtio/src/vsock/unix/muxer_rxq.rs
+++ b/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> {