From 38b748cc4d10402f371a1cd322d92cc97c4c3f6a Mon Sep 17 00:00:00 2001 From: Hui Zhu Date: Wed, 4 Dec 2019 14:26:02 +0800 Subject: [PATCH] hvsock: Update test for command Upgrade 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 --- .../src/virtio/vsock/csm/connection.rs | 47 ++++++- src/devices/src/virtio/vsock/unix/muxer.rs | 120 +++++++++++++++++- 2 files changed, 162 insertions(+), 5 deletions(-) diff --git a/src/devices/src/virtio/vsock/csm/connection.rs b/src/devices/src/virtio/vsock/csm/connection.rs index 8a1023aed669..21d60055dbaf 100644 --- a/src/devices/src/virtio/vsock/csm/connection.rs +++ b/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::::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::::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()); diff --git a/src/devices/src/virtio/vsock/unix/muxer.rs b/src/devices/src/virtio/vsock/unix/muxer.rs index ebaf010f77ec..a429d341a84a 100644 --- a/src/devices/src/virtio/vsock/unix/muxer.rs +++ b/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;