use rand::Rng;
use sha2::{Digest, Sha256};
use std::io::{Read, Write};
use vsock::{get_local_cid, VsockAddr, VsockListener, VsockStream, VMADDR_CID_HOST};
const TEST_BLOB_SIZE: usize = 1_000_000;
const TEST_BLOCK_SIZE: usize = 5_000;
const SERVER_CID: u32 = 3;
const SERVER_PORT: u32 = 8000;
const LISTEN_PORT: u32 = 9000;
#[test]
fn test_vsock() {
let mut rng = rand::rng();
let mut blob: Vec<u8> = vec![];
let mut rx_blob = vec![];
let mut tx_pos = 0;
blob.resize(TEST_BLOB_SIZE, 0);
rx_blob.resize(TEST_BLOB_SIZE, 0);
rng.fill_bytes(&mut blob);
let mut stream =
VsockStream::connect(&VsockAddr::new(SERVER_CID, SERVER_PORT)).expect("connection failed");
while tx_pos < TEST_BLOB_SIZE {
let written_bytes = stream
.write(&blob[tx_pos..tx_pos + TEST_BLOCK_SIZE])
.expect("write failed");
if written_bytes == 0 {
panic!("stream unexpectedly closed");
}
let mut rx_pos = tx_pos;
while rx_pos < (tx_pos + written_bytes) {
let read_bytes = stream.read(&mut rx_blob[rx_pos..]).expect("read failed");
if read_bytes == 0 {
panic!("stream unexpectedly closed");
}
rx_pos += read_bytes;
}
tx_pos += written_bytes;
}
let expected = Sha256::digest(&blob);
let actual = Sha256::digest(&rx_blob);
assert_eq!(expected, actual);
}
#[test]
fn test_get_local_cid() {
assert_eq!(get_local_cid().unwrap(), VMADDR_CID_HOST);
}
#[test]
fn test_listener_local_addr() {
let listener = VsockListener::bind(&VsockAddr::new(VMADDR_CID_HOST, LISTEN_PORT)).unwrap();
let local_addr = listener.local_addr().unwrap();
assert_eq!(local_addr.cid(), VMADDR_CID_HOST);
assert_eq!(local_addr.port(), LISTEN_PORT);
}
#[test]
fn test_stream_addresses() {
let stream =
VsockStream::connect(&VsockAddr::new(SERVER_CID, SERVER_PORT)).expect("connection failed");
let local_addr = stream.local_addr().unwrap();
assert!([libc::VMADDR_CID_ANY, VMADDR_CID_HOST].contains(&local_addr.cid()));
let peer_addr = stream.peer_addr().unwrap();
assert_eq!(peer_addr.cid(), SERVER_CID);
assert_eq!(peer_addr.port(), SERVER_PORT);
}