From f756174b9ff8e4476671f2269593b91c79313c15 Mon Sep 17 00:00:00 2001
From: Stefano Garzarella <sgarzare@redhat.com>
Date: Tue, 16 Jun 2020 12:21:17 +0200
Subject: [PATCH] vsock: add muxer rxq regression test

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

Signed-off-by: Gabriel Ionescu <gbi@amazon.com>
Signed-off-by: Stefano Garzarella <sgarzare@redhat.com>
---
 vm-virtio/src/vsock/csm/connection.rs |  5 ++++
 vm-virtio/src/vsock/unix/muxer.rs     | 37 +++++++++++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/vm-virtio/src/vsock/csm/connection.rs b/vm-virtio/src/vsock/csm/connection.rs
index 228842da9..e26392052 100644
--- a/vm-virtio/src/vsock/csm/connection.rs
+++ b/vm-virtio/src/vsock/csm/connection.rs
@@ -772,6 +772,11 @@ mod tests {
         pub(crate) fn fwd_cnt(&self) -> Wrapping<u32> {
             self.fwd_cnt
         }
+
+        /// Forcefully insert a credit update flag.
+        pub(crate) fn insert_credit_update(&mut self) {
+            self.pending_rx.insert(PendingRx::CreditUpdate);
+        }
     }
 
     fn init_pkt(pkt: &mut VsockPacket, op: u16, len: u32) -> &mut VsockPacket {
diff --git a/vm-virtio/src/vsock/unix/muxer.rs b/vm-virtio/src/vsock/unix/muxer.rs
index 956493d97..7625a2674 100644
--- a/vm-virtio/src/vsock/unix/muxer.rs
+++ b/vm-virtio/src/vsock/unix/muxer.rs
@@ -1378,4 +1378,41 @@ mod tests {
         // Check that fwd_cnt is 0 - "OK ..." was not accounted for.
         assert_eq!(conn.fwd_cnt().0, 0);
     }
+
+    #[test]
+    fn test_regression_rxq_pop() {
+        // Address one of the issues found while fixing the following issue:
+        // https://github.com/firecracker-microvm/firecracker/issues/1751
+        // This test checks that a connection is not popped out of the muxer
+        // rxq when multiple flags are set
+        let mut ctx = MuxerTestContext::new("regression_rxq_pop");
+        let peer_port = 1025;
+        let (mut stream, local_port) = ctx.local_connect(peer_port);
+
+        // Send some data.
+        let data = [5u8, 6, 7, 8];
+        stream.write_all(&data).unwrap();
+        ctx.notify_muxer();
+
+        // Get the connection from the connection map.
+        let key = ConnMapKey {
+            local_port,
+            peer_port,
+        };
+        let conn = ctx.muxer.conn_map.get_mut(&key).unwrap();
+
+        // Forcefully insert another flag.
+        conn.insert_credit_update();
+
+        // Call recv twice in order to check that the connection is still
+        // in the rxq.
+        assert!(ctx.muxer.has_pending_rx());
+        ctx.recv();
+        assert!(ctx.muxer.has_pending_rx());
+        ctx.recv();
+
+        // Since initially the connection had two flags set, now there should
+        // not be any pending RX in the muxer.
+        assert!(!ctx.muxer.has_pending_rx());
+    }
 }