#![cfg(feature = "wg")]
use std::net::SocketAddr;
use pktkit::wg::{Config, Handler, PacketType};
#[test]
fn under_load_cookie_reply_then_accepted_retry() {
let responder = Handler::new(Config {
load_threshold: Some(0),
..Default::default()
})
.unwrap();
let initiator = Handler::new(Config::default()).unwrap();
responder.add_peer(initiator.public_key());
initiator.add_peer(responder.public_key());
let peer_addr: SocketAddr = "203.0.113.7:51820".parse().unwrap();
let init1 = initiator
.initiate_handshake(&responder.public_key())
.unwrap();
let res1 = responder.process_packet(&init1, &peer_addr).unwrap();
assert_eq!(
res1.ty,
PacketType::CookieReply,
"expected a cookie reply under load"
);
assert!(!res1.response.is_empty());
let res2 = initiator
.process_packet(&res1.response, &peer_addr)
.unwrap();
assert_eq!(res2.ty, PacketType::CookieReceived);
assert_eq!(res2.peer_key, responder.public_key());
let init2 = initiator
.initiate_handshake(&responder.public_key())
.unwrap();
let res3 = responder.process_packet(&init2, &peer_addr).unwrap();
assert_eq!(
res3.ty,
PacketType::HandshakeResponse,
"retry with valid MAC2 should be accepted"
);
assert_eq!(res3.peer_key, initiator.public_key());
}
#[test]
fn cookie_mac2_is_source_bound() {
let responder = Handler::new(Config {
load_threshold: Some(0),
..Default::default()
})
.unwrap();
let initiator = Handler::new(Config::default()).unwrap();
responder.add_peer(initiator.public_key());
initiator.add_peer(responder.public_key());
let addr_a: SocketAddr = "203.0.113.7:51820".parse().unwrap();
let addr_b: SocketAddr = "203.0.113.8:51820".parse().unwrap();
let init1 = initiator
.initiate_handshake(&responder.public_key())
.unwrap();
let reply = responder.process_packet(&init1, &addr_a).unwrap();
assert_eq!(reply.ty, PacketType::CookieReply);
initiator.process_packet(&reply.response, &addr_a).unwrap();
let init2 = initiator
.initiate_handshake(&responder.public_key())
.unwrap();
let res = responder.process_packet(&init2, &addr_b).unwrap();
assert_eq!(res.ty, PacketType::CookieReply);
}