1use anyhow::{anyhow, bail};
3use base64::Engine;
4use boringtun::noise::handshake::parse_handshake_anon;
5use boringtun::noise::rate_limiter::RateLimiter;
6use boringtun::noise::{Packet, Tunn, TunnResult};
7use boringtun::x25519::{PublicKey, StaticSecret};
8use rand_core::RngCore;
9use std::collections::HashMap;
10use std::net::{IpAddr, SocketAddr};
11use std::sync::Arc;
12use std::time::{Duration, Instant};
13
14pub mod ws;
15
16pub const MAX_PACKET:usize=65535;
18const HANDSHAKE_RATE_LIMIT:u64=100;
20
21pub fn decode_key(s:&str)->anyhow::Result<[u8;32]>{
22 let v=base64::engine::general_purpose::STANDARD.decode(s.trim())?;
23 v.try_into().map_err(|_| anyhow!("key must decode to 32 bytes"))
24}
25
26pub fn encode_key(k:&[u8;32])->String{base64::engine::general_purpose::STANDARD.encode(k)}
27
28pub fn public_key(private:&[u8;32])->[u8;32]{PublicKey::from(&StaticSecret::from(*private)).to_bytes()}
29
30pub fn generate_key()->[u8;32]{
32 let mut k=[0u8;32];
33 rand_core::OsRng.fill_bytes(&mut k);
34 k[0]&=248;
35 k[31]&=127;
36 k[31]|=64;
37 k
38}
39
40pub fn src_address(packet:&[u8])->Option<IpAddr>{
42 match packet.first()?>>4{
43 4 if packet.len()>=20=>Some(IpAddr::from(<[u8;4]>::try_from(&packet[12..16]).ok()?)),
44 6 if packet.len()>=40=>Some(IpAddr::from(<[u8;16]>::try_from(&packet[8..24]).ok()?)),
45 _=>None,
46 }
47}
48
49pub fn dst_address(packet:&[u8])->Option<IpAddr>{Tunn::dst_address(packet)}
51
52pub struct Peer{
53 pub public_key:[u8;32],
54 pub endpoint:Option<SocketAddr>,
56 preshared_key:Option<[u8;32]>,
57 keepalive:Option<u16>,
58 tunn:Tunn,
59}
60
61const MAX_PEERS:usize=1<<24;
63
64pub struct Wg{
66 private:StaticSecret,
67 public:PublicKey,
68 peers:Vec<Option<Peer>>,
70 by_key:HashMap<[u8;32], usize>,
71 free:Vec<usize>,
73 scratch:Vec<u8>,
74 limiter:Arc<RateLimiter>,
75 quiet:Arc<RateLimiter>,
78 last_reset:Instant,
79}
80
81#[derive(Clone, Copy, Debug, PartialEq, Eq)]
83pub struct Bound{
84 pub peer:usize,
85 pub proven:bool,
87 pub keepalive:bool,
89}
90
91impl Wg{
92 pub fn new(private_key:[u8;32])->Wg{
93 let private=StaticSecret::from(private_key);
94 let public=PublicKey::from(&private);
95 let (limiter, quiet)=(Arc::new(RateLimiter::new(&public, HANDSHAKE_RATE_LIMIT)), Arc::new(RateLimiter::new(&public, u64::MAX)));
96 Wg{private, public, peers:Vec::new(), by_key:HashMap::new(), free:Vec::new(), scratch:vec![0u8;MAX_PACKET], limiter, quiet, last_reset:Instant::now()}
97 }
98
99 pub fn public_key(&self)->[u8;32]{self.public.to_bytes()}
100
101 pub fn add_peer(&mut self, public_key:[u8;32], preshared_key:Option<[u8;32]>, keepalive:Option<u16>)->anyhow::Result<usize>{
104 if self.by_key.contains_key(&public_key){bail!("duplicate peer public key")}
105 let idx=self.free.last().copied().unwrap_or(self.peers.len());
106 if idx>=MAX_PEERS{bail!("too many peers")}
107 let tunn=Tunn::new(self.private.clone(), PublicKey::from(public_key), preshared_key, keepalive, idx as u32, Some(self.quiet.clone()));
108 let peer=Some(Peer{public_key, endpoint:None, preshared_key, keepalive, tunn});
109 if self.free.pop().is_some(){self.peers[idx]=peer}else{self.peers.push(peer)}
110 self.by_key.insert(public_key, idx);
111 Ok(idx)
112 }
113
114 pub fn remove_peer(&mut self, public_key:&[u8;32])->bool{
116 let Some(idx)=self.by_key.remove(public_key) else{return false};
117 self.peers[idx]=None;
118 self.free.push(idx);
119 true
120 }
121
122 pub fn drop_sessions(&mut self, public_key:&[u8;32])->bool{
124 let Some((idx, peer))=self.by_key.get(public_key).and_then(|idx| Some((*idx, self.peers.get_mut(*idx)?.as_mut()?))) else{return false};
125 peer.tunn=Tunn::new(self.private.clone(), PublicKey::from(peer.public_key), peer.preshared_key, peer.keepalive, idx as u32, Some(self.quiet.clone()));
126 peer.endpoint=None;
127 true
128 }
129
130 pub fn index(&self, public_key:&[u8;32])->Option<usize>{self.by_key.get(public_key).copied()}
132
133 pub fn key(&self, idx:usize)->Option<[u8;32]>{Some(self.peers.get(idx)?.as_ref()?.public_key)}
135
136 pub fn peers(&self)->impl Iterator<Item=&Peer>{self.peers.iter().flatten()}
137
138 pub fn endpoint(&self, idx:usize)->Option<SocketAddr>{self.peers.get(idx)?.as_ref()?.endpoint}
139
140 pub fn recv_from_network(&mut self, from:SocketAddr, datagram:&[u8], net_out:&mut Vec<Vec<u8>>, ip_out:&mut Vec<Vec<u8>>)->Option<usize>{
143 self.recv_authenticated(from, datagram, net_out, ip_out).map(|(idx, _)| idx)
144 }
145
146 pub fn recv_authenticated(&mut self, from:SocketAddr, datagram:&[u8], net_out:&mut Vec<Vec<u8>>, ip_out:&mut Vec<Vec<u8>>)->Option<(usize, bool)>{
149 let limiter=self.limiter.clone();
152 let (idx, kind)=match limiter.verify_packet(Some(from.ip()), datagram, &mut self.scratch){
153 Ok(Packet::HandshakeInit(p))=>{
154 let hs=parse_handshake_anon(&self.private, &self.public, &p).ok()?;
155 let idx=*self.by_key.get(&hs.peer_static_public)?;
156 tracing::debug!("handshake init from {from} for peer {idx}");
157 (idx, 1u8)
158 }
159 Ok(Packet::HandshakeResponse(p))=>((p.receiver_idx>>8) as usize, 2),
160 Ok(Packet::PacketCookieReply(p))=>((p.receiver_idx>>8) as usize, 3),
161 Ok(Packet::PacketData(p))=>((p.receiver_idx>>8) as usize, 4),
162 Err(TunnResult::WriteToNetwork(cookie))=>{tracing::debug!("under load, cookie reply to {from}"); net_out.push(cookie.to_vec()); return None}
163 Err(e)=>{tracing::debug!("datagram from {from} failed the mac check: {e:?}"); return None}
164 };
165 let before=net_out.len();
166 let Wg{peers, scratch, ..}=self;
167 let peer=peers.get_mut(idx)?.as_mut()?;
168 let mut input:&[u8]=datagram;
169 let mut decrypted=false;
170 loop{
171 match peer.tunn.decapsulate(Some(from.ip()), input, scratch){
172 TunnResult::Done=>break,
173 TunnResult::Err(e)=>{tracing::debug!("peer {idx} decapsulate: {e:?}"); return None}
174 TunnResult::WriteToNetwork(b)=>{net_out.push(b.to_vec()); input=&[]}
175 TunnResult::WriteToTunnelV4(b, _)=>{ip_out.push(b.to_vec()); decrypted=true; break}
176 TunnResult::WriteToTunnelV6(b, _)=>{ip_out.push(b.to_vec()); decrypted=true; break}
177 }
178 }
179 let answered=net_out[before..].iter().any(|d| d.first()!=Some(&3));
182 let proven=decrypted|| kind==4|| (kind!=3 && answered);
183 if proven{peer.endpoint=Some(from)}
184 Some((idx, proven))
185 }
186
187 pub fn recv_bound(&mut self, public_key:&[u8;32], from:SocketAddr, datagram:&[u8], net_out:&mut Vec<Vec<u8>>, ip_out:&mut Vec<Vec<u8>>)->Option<Bound>{
191 let idx=*self.by_key.get(public_key)?;
192 let kind=match Tunn::parse_incoming_packet(datagram).ok()?{
193 Packet::HandshakeInit(p)=>{
194 if parse_handshake_anon(&self.private, &self.public, &p).ok()?.peer_static_public!=*public_key{return None}
195 1u8
196 }
197 Packet::HandshakeResponse(p) if (p.receiver_idx>>8) as usize==idx=>2,
198 Packet::PacketData(p) if (p.receiver_idx>>8) as usize==idx=>4,
199 _=>return None,
200 };
201 let (before, carried)=(net_out.len(), ip_out.len());
202 let Wg{peers, scratch, ..}=self;
203 let peer=peers.get_mut(idx)?.as_mut()?;
204 let mut input:&[u8]=datagram;
205 loop{
206 match peer.tunn.decapsulate(None, input, scratch){
207 TunnResult::Done=>break,
208 TunnResult::Err(e)=>{tracing::debug!("peer {idx} decapsulate on a bound session: {e:?}"); return None}
209 TunnResult::WriteToNetwork(b)=>{net_out.push(b.to_vec()); input=&[]}
210 TunnResult::WriteToTunnelV4(b, _)|TunnResult::WriteToTunnelV6(b, _)=>{ip_out.push(b.to_vec()); break}
211 }
212 }
213 let mut at=before;
215 while at<net_out.len(){if net_out[at].first()==Some(&3){net_out.remove(at);}else{at+=1}}
216 let proven=kind==4|| net_out.len()>before;
217 if proven{peer.endpoint=Some(from)}
218 Some(Bound{peer:idx, proven, keepalive:kind==4&& ip_out.len()==carried})
219 }
220
221 pub fn encapsulate(&mut self, idx:usize, packet:&[u8], net_out:&mut Vec<Vec<u8>>){
223 let Wg{peers, scratch, ..}=self;
224 let Some(peer)=peers.get_mut(idx).and_then(|peer| peer.as_mut()) else{return};
225 match peer.tunn.encapsulate(packet, scratch){
226 TunnResult::WriteToNetwork(b)=>net_out.push(b.to_vec()),
227 TunnResult::Err(e)=>tracing::debug!("peer {idx} encapsulate: {e:?}"),
228 _=>{}
229 }
230 }
231
232 pub fn initiate(&mut self, idx:usize, net_out:&mut Vec<Vec<u8>>){
234 let Wg{peers, scratch, ..}=self;
235 let Some(peer)=peers.get_mut(idx).and_then(|peer| peer.as_mut()) else{return};
236 if let TunnResult::WriteToNetwork(b)=peer.tunn.format_handshake_initiation(scratch, true){net_out.push(b.to_vec())}
237 }
238
239 pub fn time_since_last_handshake(&self, idx:usize)->Option<std::time::Duration>{self.peers.get(idx)?.as_ref()?.tunn.time_since_last_handshake()}
241
242 pub fn update_timers(&mut self, out:&mut Vec<(usize, Vec<u8>)>){
244 if self.last_reset.elapsed()>=Duration::from_secs(1){
246 self.limiter.reset_count();
247 self.last_reset=Instant::now();
248 }
249 let Wg{peers, scratch, ..}=self;
250 for (i, peer) in peers.iter_mut().enumerate(){
251 let Some(peer)=peer else{continue};
252 match peer.tunn.update_timers(scratch){
253 TunnResult::WriteToNetwork(b)=>out.push((i, b.to_vec())),
254 TunnResult::Err(e)=>tracing::debug!("peer {i} timers: {e:?}"),
255 _=>{}
256 }
257 }
258 }
259}
260
261#[cfg(test)]
262mod tests{
263 use super::*;
264
265 #[test]
266 fn keys_roundtrip(){
267 let k=generate_key();
268 assert_eq!(decode_key(&encode_key(&k)).unwrap(), k);
269 assert_eq!(public_key(&k), PublicKey::from(&StaticSecret::from(k)).to_bytes());
270 assert!(decode_key("bm90IDMyIGJ5dGVz").is_err());
271 }
272
273 #[test]
274 fn addresses_parse(){
275 let mut p=vec![0x45, 0, 0, 20, 0, 0, 0, 0, 64, 6, 0, 0];
276 p.extend_from_slice(&[10, 7, 0, 2]);
277 p.extend_from_slice(&[1, 1, 1, 1]);
278 assert_eq!(src_address(&p).unwrap(), "10.7.0.2".parse::<IpAddr>().unwrap());
279 assert_eq!(Tunn::dst_address(&p).unwrap(), "1.1.1.1".parse::<IpAddr>().unwrap());
280 assert!(src_address(&[]).is_none());
281 }
282
283 #[test]
284 fn a_forged_flood_cannot_steal_an_endpoint(){
285 let (srv, a)=(generate_key(), generate_key());
286 let mut wg=Wg::new(srv);
287 let ia=wg.add_peer(public_key(&a), None, None).unwrap();
288 let peer:SocketAddr="127.0.0.1:1111".parse().unwrap();
289 let attacker:SocketAddr="127.0.0.1:2222".parse().unwrap();
290 let mut client=Tunn::new(StaticSecret::from(a), PublicKey::from(public_key(&srv)), None, None, 0, None);
291 let mut buf=vec![0u8;MAX_PACKET];
292 let TunnResult::WriteToNetwork(init)=client.format_handshake_initiation(&mut buf, false) else{panic!("no initiation")};
293 let init=init.to_vec();
294 let (mut net, mut ip)=(Vec::new(), Vec::new());
295 assert_eq!(wg.recv_from_network(peer, &init, &mut net, &mut ip), Some(ia));
296 assert_eq!(wg.endpoint(ia), Some(peer));
297 let mut cookies=0;
300 for _ in 0..(HANDSHAKE_RATE_LIMIT as usize)*3{
301 net.clear();
302 assert_eq!(wg.recv_from_network(attacker, &init, &mut net, &mut ip), None);
303 cookies+=net.iter().filter(|d| d.first()==Some(&3)).count();
304 }
305 assert!(cookies>0, "the flood never reached the rate limit, so nothing was proven");
306 assert_eq!(wg.endpoint(ia), Some(peer), "a flood of forged handshakes moved the endpoint");
307 }
308
309 #[test]
310 fn the_limiter_recovers_after_a_reset(){
311 let (srv, a)=(generate_key(), generate_key());
312 let mut wg=Wg::new(srv);
313 wg.add_peer(public_key(&a), None, None).unwrap();
314 let from:SocketAddr="127.0.0.1:1111".parse().unwrap();
315 let mut client=Tunn::new(StaticSecret::from(a), PublicKey::from(public_key(&srv)), None, None, 0, None);
316 let mut buf=vec![0u8;MAX_PACKET];
317 let TunnResult::WriteToNetwork(init)=client.format_handshake_initiation(&mut buf, false) else{panic!("no initiation")};
318 let init=init.to_vec();
319 let (mut net, mut ip)=(Vec::new(), Vec::new());
320 for _ in 0..(HANDSHAKE_RATE_LIMIT as usize)*2{wg.recv_from_network(from, &init, &mut net, &mut ip);}
321 assert!(net.iter().any(|d| d.first()==Some(&3)), "the limiter never engaged");
322 std::thread::sleep(Duration::from_millis(1100));
324 wg.update_timers(&mut Vec::new());
325 net.clear();
326 wg.recv_from_network(from, &init, &mut net, &mut ip);
327 assert!(net.iter().all(|d| d.first()!=Some(&3)), "still handing out cookies after a reset");
328 }
329
330 fn connected(wg:&mut Wg, srv:&[u8;32], key:[u8;32], from:SocketAddr)->(Tunn, usize){
332 let mut client=Tunn::new(StaticSecret::from(key), PublicKey::from(public_key(srv)), None, None, 0, None);
333 let (mut buf, mut out, mut net, mut ip)=(vec![0u8;MAX_PACKET], vec![0u8;MAX_PACKET], Vec::new(), Vec::new());
334 let TunnResult::WriteToNetwork(init)=client.format_handshake_initiation(&mut buf, false) else{panic!("no initiation")};
335 let idx=wg.recv_from_network(from, init, &mut net, &mut ip).expect("the handshake was refused");
336 let TunnResult::WriteToNetwork(keepalive)=client.decapsulate(None, &net[0], &mut out) else{panic!("no keepalive after the response")};
337 net.clear();
338 assert_eq!(wg.recv_from_network(from, keepalive, &mut net, &mut ip), Some(idx));
339 assert!(wg.time_since_last_handshake(idx).is_some());
340 (client, idx)
341 }
342
343 fn data(client:&mut Tunn)->Vec<u8>{
344 let mut buf=vec![0u8;MAX_PACKET];
345 let packet=[0x45, 0, 0, 20, 0, 0, 0, 0, 64, 17, 0, 0, 10, 7, 0, 2, 1, 1, 1, 1];
346 let TunnResult::WriteToNetwork(datagram)=client.encapsulate(&packet, &mut buf) else{panic!("nothing to send")};
347 datagram.to_vec()
348 }
349
350 #[test]
351 fn a_removed_peer_is_refused_and_the_others_keep_their_place(){
352 let (srv, a, b, c)=(generate_key(), generate_key(), generate_key(), generate_key());
353 let mut wg=Wg::new(srv);
354 for key in [a, b]{wg.add_peer(public_key(&key), None, None).unwrap();}
355 let (from_a, from_b):(SocketAddr, SocketAddr)=("127.0.0.1:1111".parse().unwrap(), "127.0.0.1:2222".parse().unwrap());
356 let (mut client_a, ia)=connected(&mut wg, &srv, a, from_a);
357 let (mut client_b, ib)=connected(&mut wg, &srv, b, from_b);
358 assert!(wg.remove_peer(&public_key(&a)) && !wg.remove_peer(&public_key(&a)));
359 assert_eq!((wg.index(&public_key(&a)), wg.index(&public_key(&b)), wg.key(ia), wg.key(ib), wg.peers().count()), (None, Some(ib), None, Some(public_key(&b)), 1));
360 let (mut net, mut ip, mut buf)=(Vec::new(), Vec::new(), vec![0u8;MAX_PACKET]);
361 assert_eq!(wg.recv_from_network(from_a, &data(&mut client_a), &mut net, &mut ip), None);
363 wg.encapsulate(ia, &[0x45; 20], &mut net);
364 assert!(net.is_empty() && ip.is_empty() && wg.endpoint(ia).is_none() && wg.time_since_last_handshake(ia).is_none());
365 let mut again=Tunn::new(StaticSecret::from(a), PublicKey::from(public_key(&srv)), None, None, 0, None);
367 let TunnResult::WriteToNetwork(init)=again.format_handshake_initiation(&mut buf, false) else{panic!("no initiation")};
368 assert_eq!(wg.recv_from_network(from_a, init, &mut net, &mut ip), None);
369 assert!(net.is_empty(), "a removed peer's handshake was answered");
370 assert_eq!(wg.recv_from_network(from_b, &data(&mut client_b), &mut net, &mut ip), Some(ib));
372 assert_eq!(ip.len(), 1);
373 assert_eq!(wg.add_peer(public_key(&c), None, None).unwrap(), ia);
375 let (_client_c, ic)=connected(&mut wg, &srv, c, "127.0.0.1:3333".parse().unwrap());
376 ip.clear();
377 assert_eq!((ic, wg.recv_from_network(from_a, &data(&mut client_a), &mut net, &mut ip)), (ia, None));
378 assert!(ip.is_empty() && wg.endpoint(ic)==Some("127.0.0.1:3333".parse().unwrap()), "a removed peer reached the peer that took its index");
379 }
380
381 #[test]
382 fn dropped_sessions_need_a_new_handshake(){
383 let (srv, a)=(generate_key(), generate_key());
384 let mut wg=Wg::new(srv);
385 wg.add_peer(public_key(&a), None, None).unwrap();
386 let from:SocketAddr="127.0.0.1:1111".parse().unwrap();
387 let (mut client, idx)=connected(&mut wg, &srv, a, from);
388 assert!(wg.drop_sessions(&public_key(&a)) && !wg.drop_sessions(&[9;32]));
389 let (mut net, mut ip)=(Vec::new(), Vec::new());
390 assert_eq!(wg.recv_from_network(from, &data(&mut client), &mut net, &mut ip), None);
391 assert!(ip.is_empty() && wg.endpoint(idx).is_none() && wg.time_since_last_handshake(idx).is_none());
392 assert_eq!(connected(&mut wg, &srv, a, from).1, idx, "the peer could not shake hands again");
393 }
394
395 #[test]
397 fn a_bound_session_takes_only_its_own_peer(){
398 let (srv, a, b)=(generate_key(), generate_key(), generate_key());
399 let mut wg=Wg::new(srv);
400 let (ia, ib)=(wg.add_peer(public_key(&a), None, None).unwrap(), wg.add_peer(public_key(&b), None, None).unwrap());
401 let from:SocketAddr="127.0.0.1:1234".parse().unwrap();
402 let initiation=|key:[u8;32]|{
403 let mut client=Tunn::new(StaticSecret::from(key), PublicKey::from(public_key(&srv)), None, None, 0, None);
404 let mut buf=vec![0u8;MAX_PACKET];
405 let TunnResult::WriteToNetwork(init)=client.format_handshake_initiation(&mut buf, false) else{panic!("no initiation")};
406 let init=init.to_vec();
407 (client, init)
408 };
409 let (mut net, mut ip)=(Vec::new(), Vec::new());
410 let (_, from_b)=initiation(b);
411 assert_eq!(wg.recv_bound(&public_key(&a), from, &from_b, &mut net, &mut ip), None);
412 assert!(net.is_empty()&& wg.endpoint(ia).is_none()&& wg.endpoint(ib).is_none(), "another peer's initiation was answered on a bound session");
413 let (mut client, from_a)=initiation(a);
414 assert_eq!(wg.recv_bound(&public_key(&a), from, &from_a, &mut net, &mut ip), Some(Bound{peer:ia, proven:true, keepalive:false}));
415 assert_eq!((net.len(), net[0].first(), wg.endpoint(ia)), (1, Some(&2), Some(from)));
416 assert_eq!(wg.recv_bound(&public_key(&a), from, &from_a, &mut net, &mut ip), None);
418 assert_eq!(net.len(), 1);
419 let mut out=vec![0u8;MAX_PACKET];
421 let TunnResult::WriteToNetwork(keepalive)=client.decapsulate(None, &net[0], &mut out) else{panic!("no keepalive after the response")};
422 let keepalive=keepalive.to_vec();
423 assert_eq!(wg.recv_bound(&public_key(&b), from, &keepalive, &mut net, &mut ip), None);
424 assert_eq!(wg.recv_bound(&public_key(&a), from, &keepalive, &mut net, &mut ip), Some(Bound{peer:ia, proven:true, keepalive:true}));
425 let sent=data(&mut client);
426 assert_eq!(wg.recv_bound(&public_key(&b), from, &sent, &mut net, &mut ip), None);
427 assert_eq!(wg.recv_bound(&public_key(&a), from, &sent, &mut net, &mut ip), Some(Bound{peer:ia, proven:true, keepalive:false}));
428 assert_eq!((ip.len(), wg.recv_bound(&public_key(&a), from, &sent, &mut net, &mut ip)), (1, None), "a replayed data packet authenticated");
429 }
430
431 #[test]
433 fn a_bound_session_is_not_held_to_the_load_limit_and_gets_no_cookie(){
434 let (srv, a, b)=(generate_key(), generate_key(), generate_key());
435 let mut wg=Wg::new(srv);
436 wg.add_peer(public_key(&a), None, None).unwrap();
437 wg.add_peer(public_key(&b), None, None).unwrap();
438 let from:SocketAddr="127.0.0.1:1234".parse().unwrap();
439 let mut flood=Tunn::new(StaticSecret::from(b), PublicKey::from(public_key(&srv)), None, None, 0, None);
440 let mut buf=vec![0u8;MAX_PACKET];
441 let TunnResult::WriteToNetwork(init)=flood.format_handshake_initiation(&mut buf, false) else{panic!("no initiation")};
442 let init=init.to_vec();
443 let (mut net, mut ip)=(Vec::new(), Vec::new());
444 for _ in 0..(HANDSHAKE_RATE_LIMIT as usize)*2{wg.recv_from_network(from, &init, &mut net, &mut ip);}
445 assert!(net.iter().any(|d| d.first()==Some(&3)), "the limiter never engaged");
446 net.clear();
447 let mut client=Tunn::new(StaticSecret::from(a), PublicKey::from(public_key(&srv)), None, None, 0, None);
448 let mut buf=vec![0u8;MAX_PACKET];
449 let TunnResult::WriteToNetwork(init)=client.format_handshake_initiation(&mut buf, false) else{panic!("no initiation")};
450 assert!(wg.recv_bound(&public_key(&a), from, init, &mut net, &mut ip).is_some_and(|bound| bound.proven));
451 assert_eq!((net.len(), net[0].first()), (1, Some(&2)));
452 }
453
454 #[test]
456 fn cookie_replies_begin_above_the_handshake_limit_and_not_at_half_of_it(){
457 let (srv, a)=(generate_key(), generate_key());
458 let mut wg=Wg::new(srv);
459 wg.add_peer(public_key(&a), None, None).unwrap();
460 let from:SocketAddr="127.0.0.1:1234".parse().unwrap();
461 let mut client=Tunn::new(StaticSecret::from(a), PublicKey::from(public_key(&srv)), None, None, 0, None);
462 let mut buf=vec![0u8;MAX_PACKET];
463 let TunnResult::WriteToNetwork(init)=client.format_handshake_initiation(&mut buf, false) else{panic!("no initiation")};
464 let init=init.to_vec();
465 let (mut net, mut ip, mut first)=(Vec::new(), Vec::new(), None);
466 for n in 0..300{
467 net.clear();
468 wg.recv_from_network(from, &init, &mut net, &mut ip);
469 if first.is_none()&& net.iter().any(|d| d.first()==Some(&3)){first=Some(n)}
470 }
471 assert_eq!(first, Some(HANDSHAKE_RATE_LIMIT as usize));
472 }
473
474 #[test]
475 fn handshake_routes_to_the_right_peer(){
476 let (srv, a, b)=(generate_key(), generate_key(), generate_key());
477 let mut wg=Wg::new(srv);
478 let ia=wg.add_peer(public_key(&a), None, None).unwrap();
479 let ib=wg.add_peer(public_key(&b), None, None).unwrap();
480 assert!(wg.add_peer(public_key(&b), None, None).is_err());
481 let from:SocketAddr="127.0.0.1:1234".parse().unwrap();
482 for (key, want) in [(b, ib), (a, ia)]{
483 let mut client=Tunn::new(StaticSecret::from(key), PublicKey::from(public_key(&srv)), None, None, 0, None);
484 let mut buf=vec![0u8;MAX_PACKET];
485 let TunnResult::WriteToNetwork(init)=client.format_handshake_initiation(&mut buf, false) else{panic!("no initiation")};
486 let init=init.to_vec();
487 let (mut net, mut ip)=(Vec::new(), Vec::new());
488 assert_eq!(wg.recv_from_network(from, &init, &mut net, &mut ip), Some(want));
489 assert_eq!(net.len(), 1, "expected a handshake response");
490 assert_eq!(wg.endpoint(want), Some(from));
491 }
492 }
493}