Skip to content

Commit

Permalink
hvsock: Update test for command Upgrade <port_num>
Browse files Browse the repository at this point in the history
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>
  • Loading branch information
teawater committed Dec 11, 2019
1 parent e874a35 commit 38b748c
Show file tree
Hide file tree
Showing 2 changed files with 162 additions and 5 deletions.
47 changes: 45 additions & 2 deletions src/devices/src/virtio/vsock/csm/connection.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand All @@ -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(
Expand All @@ -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;
}
Expand Down Expand Up @@ -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.
Expand All @@ -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
Expand All @@ -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());
Expand Down
120 changes: 117 additions & 3 deletions src/devices/src/virtio/vsock/unix/muxer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand All @@ -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.
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down

0 comments on commit 38b748c

Please sign in to comment.