use phantom_protocol::transport::handshake::{
ClientHello, HandshakeClient, HandshakeResponse, HandshakeServer,
};
use rand::Rng;
const FUZZ_ITERATIONS: usize = 1000;
#[test]
fn fuzz_process_client_hello() {
let mut rng = rand::thread_rng();
let server = HandshakeServer::new().unwrap();
let client_ip = "127.0.0.1".parse().unwrap();
for _ in 0..FUZZ_ITERATIONS {
let mut bytes = vec![0u8; 1024];
rng.fill(&mut bytes[..]);
if let Ok(hello) = borsh::from_slice::<ClientHello>(&bytes) {
let _ = server.process_client_hello(&hello, 0, client_ip);
}
}
}
#[test]
fn fuzz_garbage_input_to_server() {
let mut rng = rand::thread_rng();
let server = HandshakeServer::new().unwrap();
let client_ip = "127.0.0.1".parse().unwrap();
for _ in 0..FUZZ_ITERATIONS {
let len = rng.gen_range(0..5000);
let mut bytes = vec![0u8; len];
rng.fill(&mut bytes[..]);
if let Ok(hello) = borsh::from_slice::<ClientHello>(&bytes) {
let _ = server.process_client_hello(&hello, 5, client_ip);
}
}
}
#[test]
fn test_valid_handshake_flow() {
let server = HandshakeServer::new().unwrap();
let client = HandshakeClient::new().expect("HandshakeClient::new");
let client_ip = "127.0.0.1".parse().unwrap();
let hello = client.create_client_hello();
let response = server.process_client_hello(&hello, 0, client_ip);
match response {
HandshakeResponse::Retry(r) => {
assert!(r.cookie.is_some());
let mut hello_retry = hello.clone();
hello_retry.cookie = r.cookie;
let response2 = server.process_client_hello(&hello_retry, 0, client_ip);
match response2 {
HandshakeResponse::Success(server_hello, _, _) => {
client
.process_server_hello(
&hello_retry,
&server_hello,
Some(server.verifying_key()),
)
.expect("Handshake failed");
}
_ => panic!("Expected success on second try"),
}
}
HandshakeResponse::Success(server_hello, _, _) => {
client
.process_server_hello(&hello, &server_hello, Some(server.verifying_key()))
.expect("Handshake failed");
}
HandshakeResponse::Reject(r) => panic!("unexpected reject: {:?}", r),
HandshakeResponse::Fail(e) => panic!("Handshake failed: {:?}", e),
}
}