hvsock: Update test for command Upgrade <port_num>

Update test_peer_request and test_local_request in connection.rs.
Add test_local_request_need_reply and
test_local_request_timeout_need_reply to connection.rs.
Add test_local_connection_with_upgrade,
test_local_close_with_upgrade and
test_local_connection_with_upgrade_get_oprst to muxer.rs.

Signed-off-by: Hui Zhu <teawater@antfin.com>

src/devices/src/virtio/vsock/csm/connection.rs

@@ -764,7 +764,7 @@ mod tests {
             Self::new(ConnState::Established)
         }
 
-        fn new(conn_state: ConnState) -> Self {
+        fn new_maybe_need_reply(conn_state: ConnState, need_reply: bool) -> Self {
             let vsock_test_ctx = TestContext::new();
             let mut handler_ctx = vsock_test_ctx.create_epoll_handler_context();
             let stream = TestStream::new();
@@ -782,7 +782,7 @@ mod tests {
                     PEER_BUF_ALLOC,
                 ),
                 ConnState::LocalInit => VsockConnection::<TestStream>::new_local_init(
-                    stream, LOCAL_CID, PEER_CID, LOCAL_PORT, PEER_PORT, false,
+                    stream, LOCAL_CID, PEER_CID, LOCAL_PORT, PEER_PORT, need_reply,
                 ),
                 ConnState::Established => {
                     let mut conn = VsockConnection::<TestStream>::new_peer_init(
@@ -808,6 +808,10 @@ mod tests {
             }
         }
 
+        fn new(conn_state: ConnState) -> Self {
+            Self::new_maybe_need_reply(conn_state, false)
+        }
+
         fn set_stream(&mut self, stream: TestStream) {
             self.conn.stream = stream;
         }
@@ -852,6 +856,7 @@ mod tests {
     fn test_peer_request() {
         let mut ctx = CsmTestContext::new(ConnState::PeerInit);
         assert!(ctx.conn.has_pending_rx());
+        assert_eq!(ctx.conn.need_reply, false);
         ctx.recv();
         // For peer-initiated requests, our connection should always yield a vsock reponse packet,
         // in order to establish the connection.
@@ -870,6 +875,29 @@ mod tests {
     #[test]
     fn test_local_request() {
         let mut ctx = CsmTestContext::new(ConnState::LocalInit);
+        assert_eq!(ctx.conn.need_reply, false);
+        // Host-initiated connections should first yield a connection request packet.
+        assert!(ctx.conn.has_pending_rx());
+        // Before yielding the connection request packet, the timeout kill timer shouldn't be
+        // armed.
+        assert!(!ctx.conn.will_expire());
+        ctx.recv();
+        assert_eq!(ctx.pkt.op(), uapi::VSOCK_OP_REQUEST);
+        // Since the request might time-out, the kill timer should now be armed.
+        assert!(ctx.conn.will_expire());
+        assert!(!ctx.conn.has_expired());
+        ctx.init_pkt(uapi::VSOCK_OP_RESPONSE, 0);
+        ctx.send();
+        // Upon receiving a connection response, the connection should have transitioned to the
+        // established state, and the kill timer should've been disarmed.
+        assert_eq!(ctx.conn.state, ConnState::Established);
+        assert!(!ctx.conn.will_expire());
+    }
+
+    #[test]
+    fn test_local_request_need_reply() {
+        let mut ctx = CsmTestContext::new_maybe_need_reply(ConnState::LocalInit, true);
+        assert_eq!(ctx.conn.need_reply, true);
         // Host-initiated connections should first yield a connection request packet.
         assert!(ctx.conn.has_pending_rx());
         // Before yielding the connection request packet, the timeout kill timer shouldn't be
@@ -891,6 +919,21 @@ mod tests {
     #[test]
     fn test_local_request_timeout() {
         let mut ctx = CsmTestContext::new(ConnState::LocalInit);
+        assert_eq!(ctx.conn.need_reply, false);
+        ctx.recv();
+        assert_eq!(ctx.pkt.op(), uapi::VSOCK_OP_REQUEST);
+        assert!(ctx.conn.will_expire());
+        assert!(!ctx.conn.has_expired());
+        std::thread::sleep(std::time::Duration::from_millis(
+            defs::CONN_REQUEST_TIMEOUT_MS,
+        ));
+        assert!(ctx.conn.has_expired());
+    }
+
+    #[test]
+    fn test_local_request_timeout_need_reply() {
+        let mut ctx = CsmTestContext::new_maybe_need_reply(ConnState::LocalInit, true);
+        assert_eq!(ctx.conn.need_reply, true);
         ctx.recv();
         assert_eq!(ctx.pkt.op(), uapi::VSOCK_OP_REQUEST);
         assert!(ctx.conn.will_expire());

src/devices/src/virtio/vsock/unix/muxer.rs

@@ -847,7 +847,14 @@ mod tests {
             LocalListener::new(format!("{}_{}", self.muxer.host_sock_path, port))
         }
 
-        fn local_connect(&mut self, peer_port: u32) -> (UnixStream, u32) {
+        fn local_connect_maybe_upgrade_oprst(
+            &mut self,
+            peer_port: u32,
+            is_upgrade: bool,
+            is_oprst: bool,
+        ) -> (UnixStream, u32) {
+            assert!(!is_oprst || (is_oprst && is_upgrade));
+
             let (init_local_lsn_count, init_conn_lsn_count) = self.count_epoll_listeners();
 
             let mut stream = UnixStream::connect(self.muxer.host_sock_path.clone()).unwrap();
@@ -861,7 +868,11 @@ mod tests {
             let (local_lsn_count, _) = self.count_epoll_listeners();
             assert_eq!(local_lsn_count, init_local_lsn_count + 1);
 
-            let buf = format!("CONNECT {}\n", peer_port);
+            let buf = if is_upgrade {
+                format!("Upgrade {}\n", peer_port)
+            } else {
+                format!("CONNECT {}\n", peer_port)
+            };
             stream.write_all(buf.as_bytes()).unwrap();
             // The muxer would now get notified that data is available for reading from the locally
             // initiated connection.
@@ -890,11 +901,27 @@ mod tests {
             assert_eq!(self.pkt.dst_port(), peer_port);
             assert_eq!(self.pkt.src_port(), local_port);
 
-            self.init_pkt(local_port, peer_port, uapi::VSOCK_OP_RESPONSE);
+            self.init_pkt(
+                local_port,
+                peer_port,
+                if is_oprst {
+                    uapi::VSOCK_OP_RST
+                } else {
+                    uapi::VSOCK_OP_RESPONSE
+                },
+            );
             self.send();
 
             (stream, local_port)
         }
+
+        fn local_connect(&mut self, peer_port: u32) -> (UnixStream, u32) {
+            self.local_connect_maybe_upgrade_oprst(peer_port, false, false)
+        }
+
+        fn local_connect_with_upgrade(&mut self, peer_port: u32) -> (UnixStream, u32) {
+            self.local_connect_maybe_upgrade_oprst(peer_port, true, false)
+        }
     }
 
     struct LocalListener {
@@ -1071,6 +1098,53 @@ mod tests {
         assert_eq!(ctx.pkt.buf().unwrap()[..data.len()], data);
     }
 
+    #[test]
+    fn test_local_connection_with_upgrade() {
+        let mut ctx = MuxerTestContext::new("local_connection_with_upgrade");
+        let peer_port = 1025;
+        let (mut stream, local_port) = ctx.local_connect_with_upgrade(peer_port);
+
+        // Test the handshake
+        let mut buf = vec![0; 4];
+        stream.read_exact(buf.as_mut_slice()).unwrap();
+        let buf = String::from_utf8(buf).unwrap();
+        assert_eq!(buf, "101\n".to_string());
+
+        // Test guest -> host data flow.
+        let data = [1, 2, 3, 4];
+        ctx.init_data_pkt(local_port, peer_port, &data);
+        ctx.send();
+
+        let mut buf = vec![0u8; data.len()];
+        stream.read_exact(buf.as_mut_slice()).unwrap();
+        assert_eq!(buf.as_slice(), &data);
+
+        // Test host -> guest data flow.
+        let data = [5, 6, 7, 8];
+        stream.write_all(&data).unwrap();
+        ctx.notify_muxer();
+
+        assert!(ctx.muxer.has_pending_rx());
+        ctx.recv();
+        assert_eq!(ctx.pkt.op(), uapi::VSOCK_OP_RW);
+        assert_eq!(ctx.pkt.src_port(), local_port);
+        assert_eq!(ctx.pkt.dst_port(), peer_port);
+        assert_eq!(ctx.pkt.buf().unwrap()[..data.len()], data);
+    }
+
+    #[test]
+    fn test_local_connection_with_upgrade_get_oprst() {
+        let mut ctx = MuxerTestContext::new("local_connection_with_upgrade_get_oprst");
+        let peer_port = 1025;
+        let (mut stream, _local_port) =
+            ctx.local_connect_maybe_upgrade_oprst(peer_port, true, true);
+
+        let mut buf = vec![0; 4];
+        stream.read_exact(buf.as_mut_slice()).unwrap();
+        let buf = String::from_utf8(buf).unwrap();
+        assert_eq!(buf, "503\n".to_string());
+    }
+
     #[test]
     fn test_local_close() {
         let peer_port = 1025;
@@ -1104,6 +1178,46 @@ mod tests {
         assert!(!ctx.muxer.local_port_set.contains(&local_port));
     }
 
+    #[test]
+    fn test_local_close_with_upgrade() {
+        let peer_port = 1025;
+        let mut ctx = MuxerTestContext::new("local_close_with_upgrade");
+        let local_port;
+        {
+            let (mut stream, local_port_) = ctx.local_connect_with_upgrade(peer_port);
+
+            // Test the handshake
+            let mut buf = vec![0; 4];
+            stream.read_exact(buf.as_mut_slice()).unwrap();
+            let buf = String::from_utf8(buf).unwrap();
+            assert_eq!(buf, "101\n".to_string());
+
+            local_port = local_port_;
+        }
+        // Local var `_stream` was now dropped, thus closing the local stream. After the muxer gets
+        // notified via EPOLLIN, it should attempt to gracefully shutdown the connection, issuing a
+        // VSOCK_OP_SHUTDOWN with both no-more-send and no-more-recv indications set.
+        ctx.notify_muxer();
+        assert!(ctx.muxer.has_pending_rx());
+        ctx.recv();
+        assert_eq!(ctx.pkt.op(), uapi::VSOCK_OP_SHUTDOWN);
+        assert_ne!(ctx.pkt.flags() & uapi::VSOCK_FLAGS_SHUTDOWN_SEND, 0);
+        assert_ne!(ctx.pkt.flags() & uapi::VSOCK_FLAGS_SHUTDOWN_RCV, 0);
+        assert_eq!(ctx.pkt.src_port(), local_port);
+        assert_eq!(ctx.pkt.dst_port(), peer_port);
+
+        // The connection should get removed (and its local port freed), after the peer replies
+        // with an RST.
+        ctx.init_pkt(local_port, peer_port, uapi::VSOCK_OP_RST);
+        ctx.send();
+        let key = ConnMapKey {
+            local_port,
+            peer_port,
+        };
+        assert!(!ctx.muxer.conn_map.contains_key(&key));
+        assert!(!ctx.muxer.local_port_set.contains(&local_port));
+    }
+
     #[test]
     fn test_peer_close() {
         let peer_port = 1025;