clover's git

snowbound

1	//! The relay: rooms of WebSocket peers, each message passed on to the peer it names, with
2	//! limits on everything a stranger can make it hold. A thread reads each connection and
3	//! another writes it, from a queue capped in bytes.
4
5	use crate::{BROADCAST, GROUP, Notice, SLOT, Verdict, ws};
6	use std::{
7	collections::{BTreeMap, HashMap, VecDeque},
8	io::{BufReader, Write},
9	net::{IpAddr, Shutdown, TcpListener, TcpStream},
10	ops::RangeInclusive,
11	sync::{
12	Arc, Condvar, Mutex,
13	atomic::{AtomicU64, AtomicUsize, Ordering},
14	},
15	thread,
16	time::{Duration, Instant},
17	};
18
19	/// What a relay allows; `snowbound-relay --help` explains each.
20	#[derive(Clone, Debug)]
21	pub struct Config {
22	/// Counts a peer by the last `X-Forwarded-For` entry, the one its own proxy added.
23	pub trust_forwarded: bool,
24	pub max_connections: usize,
25	pub max_connections_per_address: usize,
26	pub max_rooms: usize,
27	pub max_room_peers: usize,
28	/// The largest message, in bytes.
29	pub max_message: usize,
30	/// The most bytes waiting to go to one peer before the relay hangs up on it.
31	pub queue: usize,
32	/// How long a connection may send nothing.
33	pub idle: Duration,
34	pub room_bytes_per_second: u64,
35	pub joins_per_minute: u32,
36	pub room_joins_per_minute: u32,
37	/// Wrong codes an address may try in a minute before it is locked out, a minute the first
38	/// time and twice as long each time after, up to an hour.
39	pub failures_per_minute: u32,
40	/// Wrong codes after which a code admits no one new.
41	pub burn_after: u32,
42	/// How long a peer joining a code's room has to meet its owner.
43	pub pending: Duration,
44	}
45
46	impl Default for Config {
47	fn default() -> Self {
48	Self {
49	trust_forwarded: false,
50	max_connections: 256,
51	max_connections_per_address: 128,
52	max_rooms: 128,
53	max_room_peers: 64,
54	max_message: 256 << 10,
55	queue: 1 << 20,
56	idle: Duration::from_secs(600),
57	room_bytes_per_second: 4 << 20,
58	joins_per_minute: 240,
59	room_joins_per_minute: 120,
60	failures_per_minute: 10,
61	burn_after: 5,
62	pending: Duration::from_secs(20),
63	}
64	}
65	}
66
67	const HANDSHAKE: Duration = Duration::from_secs(10);
68	/// How long a polled session's `GET` waits for something to bring, how long one of its
69	/// connections may idle between requests, and how long a session may ask nothing.
70	const WAIT: Duration = Duration::from_secs(25);
71	const KEPT: Duration = Duration::from_secs(60);
72	const IDLE_POLL: Duration = Duration::from_secs(60);
73	const WRITE: Duration = Duration::from_secs(30);
74	const MINUTE: Duration = Duration::from_secs(60);
75	const HOUR: Duration = Duration::from_secs(3600);
76	/// Addresses remembered at once; past it, new ones wait.
77	const ADDRESSES: usize = 1 << 16;
78	const NAMEPLATES: RangeInclusive<u32> = 1..=999;
79	const STACK: usize = 256 << 10;
80
81	/// Serves `listener` until it fails.
82	pub fn serve(listener: TcpListener, config: Config) -> std::io::Result<()> {
83	let relay = Arc::new(Relay {
84	config,
85	state: Mutex::default(),
86	connections: AtomicUsize::new(0),
87	relayed: Default::default(),
88	started: Instant::now(),
89	});
90	let sweeping = Arc::clone(&relay);
91	thread::Builder::new().name("sweep".into()).spawn(move || {
92	loop {
93	thread::sleep(Duration::from_secs(1));
94	sweeping.sweep(Instant::now());
95	}
96	})?;
97	for stream in listener.incoming() {
98	let Ok(stream) = stream else {
99	// Out of descriptors, most likely: let some close.
100	thread::sleep(Duration::from_millis(50));
101	continue;
102	};
103	if relay.connections.fetch_add(1, Ordering::AcqRel) >= relay.config.max_connections {
104	relay.connections.fetch_sub(1, Ordering::AcqRel);
105	continue;
106	}
107	let serving = Arc::clone(&relay);
108	let spawned = thread::Builder::new().stack_size(STACK).spawn(move || {
109	serving.connection(stream);
110	serving.connections.fetch_sub(1, Ordering::AcqRel);
111	});
112	if spawned.is_err() {
113	relay.connections.fetch_sub(1, Ordering::AcqRel);
114	}
115	}
116	Ok(())
117	}
118
119	struct Relay {
120	config: Config,
121	state: Mutex<State>,
122	/// Connections open, joined or not.
123	connections: AtomicUsize,
124	/// Bytes of peers' messages taken in, and given out.
125	relayed: [AtomicU64; 2],
126	started: Instant,
127	}
128
129	#[derive(Default)]
130	struct State {
131	rooms: HashMap<String, Room>,
132	addresses: HashMap<IpAddr, Address>,
133	/// Peers that reach the relay by requests rather than a WebSocket, by session.
134	polls: HashMap<String, Poll>,
135	}
136
137	/// A peer in a room by requests: where its messages wait for its next `GET`.
138	struct Poll {
139	tag: String,
140	slot: u32,
141	outbox: Arc<Outbox>,
142	/// When it last asked anything; one quiet past `IDLE_POLL` has gone.
143	last: Instant,
144	}
145
146	struct Room {
147	/// The slot of the peer that claimed a code's room, while it is there.
148	owner: Option<u32>,
149	next: u32,
150	members: BTreeMap<u32, Member>,
151	joins: Bucket,
152	bytes: Bucket,
153	failures: u32,
154	/// When wrong codes burned it, then when it emptied: its number stays out of use for a
155	/// minute after, so an old code never leads to its owner's next one.
156	burned: Option<Instant>,
157	}
158
159	struct Member {
160	outbox: Arc<Outbox>,
161	address: IpAddr,
162	/// When a peer joining a code's room joined, until its owner says it met it.
163	pending: Option<Instant>,
164	}
165
166	/// What one address (or IPv6 /64) is doing.
167	struct Address {
168	connections: usize,
169	pending: usize,
170	joins: Bucket,
171	/// Wrong codes in the last minute.
172	failures: VecDeque<Instant>,
173	/// Lockouts in a row, each twice as long.
174	strikes: u32,
175	struck: Option<Instant>,
176	locked: Option<Instant>,
177	}
178
179	enum Ask {
180	/// A code's room, numbered by the relay or, coming back, as it was.
181	Claim(Option<u32>),
182	Room(String),
183	}
184
185	enum Refusal {
186	NotFound,
187	Gone,
188	Wait(Duration),
189	Full,
190	}
191
192	impl Relay {
193	/// Answers a connection's requests one after another, as a polled session sends them,
194	/// until one takes the connection over as a WebSocket or it ends.
195	fn connection(&self, stream: TcpStream) {
196	let _ = stream.set_nodelay(true);
197	let _ = stream.set_read_timeout(Some(HANDSHAKE));
198	let _ = stream.set_write_timeout(Some(WRITE));
199	let Ok(reading) = stream.try_clone() else {
200	return;
201	};
202	let mut reader = BufReader::new(reading);
203	while let Ok(head) = ws::head(&mut reader) {
204	if !self.request(&stream, &mut reader, &head) {
205	return;
206	}
207	// Between a session's requests, a connection may idle as long as one waits.
208	let _ = stream.set_read_timeout(Some(KEPT));
209	}
210	}
211
212	/// Answers the request `head`: whether the connection serves another.
213	fn request(&self, stream: &TcpStream, reader: &mut BufReader<TcpStream>, head: &str) -> bool {
214	let address = crate::peer(stream, head, self.config.trust_forwarded);
215	let target = head.split(' ').nth(1).unwrap_or_default();
216	let (path, query) = target.split_once('?').unwrap_or((target, ""));
217	if let Some(id) = path.strip_prefix("/v1/poll/") {
218	return self.polled(stream, reader, head, id);
219	}
220	let ask = match path {
221	"/health" => {
222	respond(stream, "200 OK", "application/json", "", &self.health());
223	return false;
224	}
225	"/v1/claim" => Ask::Claim(
226	query
227	.split('&')
228	.find_map(|pair| pair.strip_prefix("nameplate="))
229	.and_then(|number| number.parse().ok()),
230	),
231	_ => match path.strip_prefix("/v1/room/").filter(|tag| valid(tag)) {
232	Some(tag) => Ask::Room(tag.into()),
233	None => {
234	respond(stream, "404 Not Found", "text/plain", "", "No such page\n");
235	return false;
236	}
237	},
238	};
239	if query.split('&').any(|pair| pair == "poll=1") {
240	let outbox = Arc::new(Outbox::mailbox(self.config.queue));
241	return match self.join(address, ask, &outbox, Instant::now()) {
242	Ok((tag, slot)) => {
243	let mut id = [0; 16];
244	if getrandom::fill(&mut id).is_err() {
245	return false;
246	}
247	let id: String = id.iter().map(|byte| format!("{byte:02x}")).collect();
248	let poll = Poll {
249	tag,
250	slot,
251	outbox,
252	last: Instant::now(),
253	};
254	self.state.lock().unwrap().polls.insert(id.clone(), poll);
255	reply(stream, "200 OK", format!("session {id}").as_bytes())
256	}
257	Err(refusal) => {
258	refuse(stream, refusal);
259	false
260	}
261	};
262	}
263	let upgrade = ws::header(head, "Upgrade")
264	.is_some_and(|value| value.eq_ignore_ascii_case("websocket"));
265	let (true, Some(key)) = (
266	upgrade && head.starts_with("GET "),
267	ws::header(head, "Sec-WebSocket-Key"),
268	) else {
269	respond(
270	stream,
271	"400 Bad Request",
272	"text/plain",
273	"",
274	"A WebSocket, or ?poll=1\n",
275	);
276	return false;
277	};
278	let Ok(writing) = stream.try_clone() else {
279	return false;
280	};
281	let outbox = Arc::new(Outbox::new(writing, self.config.queue));
282	let (tag, slot) = match self.join(address, ask, &outbox, Instant::now()) {
283	Ok(joined) => joined,
284	Err(refusal) => {
285	refuse(stream, refusal);
286	return false;
287	}
288	};
289	let switching = format!(
290	"HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\
291	Sec-WebSocket-Accept: {}\r\n\r\n",
292	ws::accept(key)
293	);
294	let writer = Arc::clone(&outbox);
295	if (&*stream).write_all(switching.as_bytes()).is_ok()
296	&& thread::Builder::new()
297	.stack_size(STACK)
298	.spawn(move || writer.drain())
299	.is_ok()
300	{
301	let _ = stream.set_read_timeout(Some(self.config.idle));
302	let mut reader = ws::Reader::new(reader, self.config.max_message, true);
303	while let Ok(message) = reader.read() {
304	if !self.heard(&tag, slot, &outbox, message) {
305	break;
306	}
307	}
308	}
309	let mut state = self.state.lock().unwrap();
310	depart(&mut state, &self.config, &tag, slot, Instant::now());
311	false
312	}
313
314	/// A polled session's request: `GET` waits for what is to go to it, `POST` brings what
315	/// it sends, each a run of WebSocket frames. Whether the connection serves another.
316	fn polled(
317	&self,
318	stream: &TcpStream,
319	reader: &mut BufReader<TcpStream>,
320	head: &str,
321	id: &str,
322	) -> bool {
323	let length: usize = ws::header(head, "Content-Length")
324	.and_then(|length| length.parse().ok())
325	.unwrap_or(0);
326	if length > self.config.max_message * 4 {
327	return false;
328	}
329	let mut body = vec![0; length];
330	if std::io::Read::read_exact(reader, &mut body).is_err() {
331	return false;
332	}
333	let session = {
334	let mut state = self.state.lock().unwrap();
335	state.polls.get_mut(id).map(|poll| {
336	poll.last = Instant::now();
337	(poll.tag.clone(), poll.slot, Arc::clone(&poll.outbox))
338	})
339	};
340	let Some((tag, slot, outbox)) = session else {
341	return reply(stream, "410 Gone", b"No such session\n");
342	};
343	if head.starts_with("POST ") {
344	let mut frames = ws::Reader::new(&body[..], self.config.max_message, true);
345	while let Ok(message) = frames.read() {
346	if !self.heard(&tag, slot, &outbox, message) {
347	self.end_poll(id);
348	return reply(stream, "410 Gone", b"Closed\n");
349	}
350	}
351	return reply(stream, "200 OK", b"");
352	}
353	let _ = stream.set_write_timeout(Some(WRITE));
354	match outbox.take(WAIT) {
355	Some(bytes) => {
356	if let Some(poll) = self.state.lock().unwrap().polls.get_mut(id) {
357	poll.last = Instant::now();
358	}
359	reply(stream, "200 OK", &bytes)
360	}
361	None => {
362	self.end_poll(id);
363	reply(stream, "410 Gone", b"Closed\n")
364	}
365	}
366	}
367
368	/// Ends the polled session `id`: it leaves its room.
369	fn end_poll(&self, id: &str) {
370	let mut state = self.state.lock().unwrap();
371	if let Some(poll) = state.polls.remove(id) {
372	depart(
373	&mut state,
374	&self.config,
375	&poll.tag,
376	poll.slot,
377	Instant::now(),
378	);
379	}
380	}
381
382	fn health(&self) -> String {
383	let state = self.state.lock().unwrap();
384	let peers: usize = state.rooms.values().map(|room| room.members.len()).sum();
385	let [taken, given] = &self.relayed;
386	format!(
387	"{{\"rooms\":{},\"peers\":{peers},\"connections\":{},\"seconds\":{},\
388	\"bytes_in\":{},\"bytes_out\":{}}}\n",
389	state.rooms.len(),
390	self.connections.load(Ordering::Acquire),
391	self.started.elapsed().as_secs(),
392	taken.load(Ordering::Relaxed),
393	given.load(Ordering::Relaxed),
394	)
395	}
396
397	/// Puts a peer from `address` in the room it asks for, telling it and those it may talk
398	/// to: its room's tag and its slot.
399	fn join(
400	&self,
401	address: IpAddr,
402	ask: Ask,
403	outbox: &Arc<Outbox>,
404	now: Instant,
405	) -> Result<(String, u32), Refusal> {
406	let config = &self.config;
407	let mut state = self.state.lock().unwrap();
408	let State {
409	rooms, addresses, ..
410	} = &mut *state;
411	if !addresses.contains_key(&address) && addresses.len() >= ADDRESSES {
412	return Err(Refusal::Full);
413	}
414	let client = addresses
415	.entry(address)
416	.or_insert_with(|| Address::new(config, now));
417	client.refresh(now);
418	if let Some(until) = client.locked {
419	return Err(Refusal::Wait(until - now));
420	}
421	if client.connections >= config.max_connections_per_address {
422	return Err(Refusal::Wait(Duration::from_secs(10)));
423	}
424	let per_minute = f64::from(config.joins_per_minute);
425	client
426	.joins
427	.take(per_minute, per_minute / 60.0, now)
428	.map_err(Refusal::Wait)?;
429	let (tag, nameplate) = match ask {
430	Ask::Claim(back) => {
431	let free = |number: u32| {
432	rooms
433	.get(&code(number))
434	.is_none_or(|room| room.owner.is_none() && room.burned.is_none())
435	};
436	let number = back
437	.filter(|number| NAMEPLATES.contains(number) && free(*number))
438	.or_else(|| {
439	(0..32)
440	.filter_map(|_| {
441	let mut bytes = [0; 4];
442	getrandom::fill(&mut bytes).ok()?;
443	Some(1 + u32::from_le_bytes(bytes) % NAMEPLATES.end())
444	})
445	.chain(NAMEPLATES)
446	.find(|number| free(*number))
447	})
448	.ok_or(Refusal::Full)?;
449	(code(number), Some(number))
450	}
451	Ask::Room(tag) if tag.starts_with("code-") => {
452	let room = rooms.get(&tag).ok_or(Refusal::NotFound)?;
453	if room.burned.is_some() {
454	return Err(Refusal::Gone);
455	}
456	if room.owner.is_none() {
457	return Err(Refusal::NotFound);
458	}
459	if client.failures.len() + client.pending >= config.failures_per_minute as usize {
460	let wait = client.failures.front().map_or(config.pending, |first| {
461	(*first + MINUTE).saturating_duration_since(now)
462	});
463	return Err(Refusal::Wait(wait));
464	}
465	(tag, None)
466	}
467	Ask::Room(tag) => (tag, None),
468	};
469	if !rooms.contains_key(&tag) && rooms.len() >= config.max_rooms {
470	return Err(Refusal::Full);
471	}
472	let room = rooms.entry(tag.clone()).or_insert_with(|| Room {
473	owner: None,
474	next: 1,
475	members: BTreeMap::new(),
476	joins: Bucket::full(f64::from(config.room_joins_per_minute), now),
477	bytes: Bucket::full(config.room_bytes_per_second as f64, now),
478	failures: 0,
479	burned: None,
480	});
481	let per_minute = f64::from(config.room_joins_per_minute);
482	let admitted = room.members.len() < config.max_room_peers;
483	let joined = admitted.then(|| room.joins.take(per_minute, per_minute / 60.0, now));
484	let refusal = match joined {
485	None => Some(Refusal::Full),
486	Some(Err(wait)) => Some(Refusal::Wait(wait)),
487	Some(Ok(())) => None,
488	};
489	if let Some(refusal) = refusal {
490	if room.members.is_empty() {
491	rooms.remove(&tag);
492	}
493	return Err(refusal);
494	}
495	let slot = room.next;
496	room.next += 1;
497	let pending = room.owner.is_some() && nameplate.is_none();
498	let visible: Vec<u32> = match room.owner {
499	Some(owner) if pending => vec![owner],
500	_ => (room.members.iter())
501	.filter(|(_, member)| member.pending.is_none())
502	.map(|(slot, _)| *slot)
503	.collect(),
504	};
505	if let Some(number) = nameplate {
506	room.owner = Some(slot);
507	outbox.push(notice(Notice::Nameplate(number)));
508	}
509	outbox.push(notice(Notice::Welcome {
510	you: slot,
511	members: visible.clone(),
512	}));
513	for other in visible {
514	room.members[&other]
515	.outbox
516	.push(notice(Notice::Joined(slot)));
517	}
518	room.members.insert(
519	slot,
520	Member {
521	outbox: Arc::clone(outbox),
522	address,
523	pending: pending.then_some(now),
524	},
525	);
526	client.connections += 1;
527	client.pending += usize::from(pending);
528	Ok((tag, slot))
529	}
530
531	/// Acts on a message from `slot`: false once it is gone or broke the protocol.
532	fn heard(&self, tag: &str, slot: u32, outbox: &Outbox, message: ws::Message) -> bool {
533	match message {
534	ws::Message::Binary(data) => {
535	let Some((to, rest)) = data.split_first_chunk::<SLOT>() else {
536	return false;
537	};
538	let to = u32::from_be_bytes(*to);
539	// The slots a group message names, then what it carries.
540	let (named, payload) = match to {
541	BROADCAST => (None, rest),
542	to if to & GROUP != 0 => {
543	let Some((named, payload)) =
544	rest.split_at_checked((to & !GROUP) as usize * SLOT)
545	else {
546	return false;
547	};
548	let named = named.as_chunks::<SLOT>().0.iter();
549	(Some(named.map(|slot| u32::from_be_bytes(*slot))), payload)
550	}
551	_ => (None, rest),
552	};
553	let state = self.state.lock().unwrap();
554	let Some(room) = state.rooms.get(tag) else {
555	return false;
556	};
557	let Some(from) = room.members.get(&slot) else {
558	return false;
559	};
560	// One waiting for the code's owner talks to the owner alone.
561	let allowed = |to: u32| {
562	let target = room.members.get(&to)?;
563	(to != slot
564	&& (from.pending.is_none() || room.owner == Some(to))
565	&& (target.pending.is_none() || room.owner == Some(slot)))
566	.then_some(&target.outbox)
567	};
568	let (source, targets): (u32, Vec<&Arc<Outbox>>) = match (to, named) {
569	(BROADCAST, _) => (
570	slot | GROUP,
571	room.members.keys().filter_map(|to| allowed(*to)).collect(),
572	),
573	(_, Some(named)) => (slot | GROUP, named.filter_map(allowed).collect()),
574	(to, None) => (slot, allowed(to).into_iter().collect()),
575	};
576	let message = ws::frame(
577	ws::BINARY,
578	&[&source.to_be_bytes()[..], payload].concat(),
579	None,
580	);
581	let targets: Vec<Arc<Outbox>> = targets.into_iter().cloned().collect();
582	drop(state);
583	// Waiting here slows the sender alone, as its socket fills.
584	let given = (SLOT + payload.len()) * targets.len();
585	let rate = self.config.room_bytes_per_second as f64;
586	let wait = match self.state.lock().unwrap().rooms.get_mut(tag) {
587	Some(room) => room.bytes.spend(given as f64, rate, Instant::now()),
588	None => return false,
589	};
590	thread::sleep(wait);
591	for target in targets {
592	target.push(message.clone());
593	}
594	let [taken, out] = &self.relayed;
595	taken.fetch_add(data.len() as u64, Ordering::Relaxed);
596	out.fetch_add(given as u64, Ordering::Relaxed);
597	true
598	}
599	ws::Message::Text(text) => {
600	if let Ok(verdict) = text.parse() {
601	self.judge(tag, slot, verdict);
602	}
603	true
604	}
605	ws::Message::Ping(payload) => {
606	outbox.push(ws::frame(ws::PONG, &payload, None));
607	true
608	}
609	ws::Message::Pong => true,
610	ws::Message::Close => false,
611	}
612	}
613
614	/// Takes a code's owner's word on the peer in a slot waiting to meet it.
615	fn judge(&self, tag: &str, from: u32, verdict: Verdict) {
616	let mut state = self.state.lock().unwrap();
617	let State {
618	rooms, addresses, ..
619	} = &mut *state;
620	let Some(room) = rooms.get_mut(tag).filter(|room| room.owner == Some(from)) else {
621	return;
622	};
623	match verdict {
624	Verdict::Met(slot) => {
625	let Some(member) = room.members.get_mut(&slot) else {
626	return;
627	};
628	if member.pending.take().is_none() {
629	return;
630	}
631	if let Some(address) = addresses.get_mut(&member.address) {
632	address.pending -= 1;
633	}
634	let newcomer = Arc::clone(&member.outbox);
635	for (other, peer) in &room.members {
636	if *other != slot && *other != from && peer.pending.is_none() {
637	newcomer.push(notice(Notice::Joined(*other)));
638	peer.outbox.push(notice(Notice::Joined(slot)));
639	}
640	}
641	}
642	Verdict::Failed(slot) => {
643	if room
644	.members
645	.get(&slot)
646	.is_some_and(|member| member.pending.is_some())
647	{
648	depart(&mut state, &self.config, tag, slot, Instant::now());
649	}
650	}
651	}
652	}
653
654	/// Fails the peers that waited too long to meet a code's owner, and forgets burned rooms
655	/// empty a minute and quiet addresses.
656	fn sweep(&self, now: Instant) {
657	let config = &self.config;
658	let mut state = self.state.lock().unwrap();
659	let late: Vec<(String, u32)> = (state.rooms.iter())
660	.flat_map(|(tag, room)| {
661	(room.members.iter())
662	.filter(|(_, member)| {
663	member.pending.is_some_and(|since| {
664	now.saturating_duration_since(since) >= config.pending
665	})
666	})
667	.map(|(slot, _)| (tag.clone(), *slot))
668	})
669	.collect();
670	for (tag, slot) in late {
671	depart(&mut state, config, &tag, slot, now);
672	}
673	let quiet: Vec<String> = (state.polls.iter())
674	.filter(|(_, poll)| now.saturating_duration_since(poll.last) >= IDLE_POLL)
675	.map(|(id, _)| id.clone())
676	.collect();
677	for id in quiet {
678	if let Some(poll) = state.polls.remove(&id) {
679	depart(&mut state, config, &poll.tag, poll.slot, now);
680	}
681	}
682	state.rooms.retain(|_, room| {
683	!room.members.is_empty()
684	|| room
685	.burned
686	.is_some_and(|since| now.saturating_duration_since(since) < MINUTE)
687	});
688	state.addresses.retain(|_, address| {
689	address.refresh(now);
690	!address.quiet(config, now)
691	});
692	}
693	}
694
695	/// Takes `slot` out of room `tag` and hangs up on it. One still waiting to meet a code's
696	/// owner tried a wrong code; the owner leaving excuses those waiting for it.
697	fn depart(state: &mut State, config: &Config, tag: &str, slot: u32, now: Instant) {
698	let State {
699	rooms, addresses, ..
700	} = state;
701	let Some(room) = rooms.get_mut(tag) else {
702	return;
703	};
704	let Some(member) = room.members.remove(&slot) else {
705	return;
706	};
707	member.outbox.close();
708	if let Some(address) = addresses.get_mut(&member.address) {
709	address.connections -= 1;
710	if member.pending.is_some() {
711	address.pending -= 1;
712	if let Some(lockout) = address.fail(now, config.failures_per_minute) {
713	eprintln!(
714	"{}: locked out for {} s after too many wrong codes",
715	member.address,
716	lockout.as_secs()
717	);
718	}
719	}
720	}
721	for peer in room.members.values() {
722	peer.outbox.push(notice(Notice::Left(slot)));
723	}
724	if member.pending.is_some() {
725	room.failures += 1;
726	if room.failures >= config.burn_after && room.burned.is_none() {
727	room.burned = Some(now);
728	eprintln!("{tag}: burned after {} wrong codes", room.failures);
729	if let Some(owner) = room.owner.and_then(|owner| room.members.get(&owner)) {
730	owner.outbox.push(notice(Notice::Burned));
731	}
732	}
733	}
734	if room.owner == Some(slot) {
735	room.owner = None;
736	for member in room.members.values_mut() {
737	if member.pending.take().is_some() {
738	if let Some(address) = addresses.get_mut(&member.address) {
739	address.pending -= 1;
740	}
741	member.outbox.close();
742	}
743	}
744	}
745	if room.members.is_empty() {
746	match &mut room.burned {
747	Some(since) => *since = now,
748	None => {
749	rooms.remove(tag);
750	}
751	}
752	}
753	}
754
755	fn notice(notice: Notice) -> Vec<u8> {
756	ws::frame(ws::TEXT, notice.to_string().as_bytes(), None)
757	}
758
759	fn code(number: u32) -> String {
760	format!("code-{number}")
761	}
762
763	/// A room tag as clients make them: a hash in hex, or `code-` and a number.
764	fn valid(tag: &str) -> bool {
765	(1..=64).contains(&tag.len())
766	&& tag
767	.bytes()
768	.all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'-')
769	}
770
771	/// Answers a polled session's request, keeping the connection: whether that worked.
772	fn reply(mut stream: &TcpStream, status: &str, body: &[u8]) -> bool {
773	let head = format!(
774	"HTTP/1.1 {status}\r\nContent-Type: application/octet-stream\r\nContent-Length: {}\r\n\
775	Cache-Control: no-store\r\n\r\n",
776	body.len()
777	);
778	stream.write_all(head.as_bytes()).is_ok() && stream.write_all(body).is_ok()
779	}
780
781	fn refuse(stream: &TcpStream, refusal: Refusal) {
782	let (status, headers, body) = match refusal {
783	Refusal::NotFound => ("404 Not Found", String::new(), "No such code\n"),
784	Refusal::Gone => ("410 Gone", String::new(), "The code has expired\n"),
785	Refusal::Wait(wait) => (
786	"429 Too Many Requests",
787	// Rounded up, so that a client waiting so long finds it over.
788	format!(
789	"Retry-After: {}\r\n",
790	wait.as_secs() + u64::from(wait.subsec_nanos() > 0)
791	),
792	"Too many tries\n",
793	),
794	Refusal::Full => (
795	"503 Service Unavailable",
796	"Retry-After: 30\r\n".into(),
797	"The relay is full\n",
798	),
799	};
800	respond(stream, status, "text/plain", &headers, body);
801	}
802
803	fn respond(mut stream: &TcpStream, status: &str, kind: &str, headers: &str, body: &str) {
804	let _ = write!(
805	stream,
806	"HTTP/1.1 {status}\r\nContent-Type: {kind}\r\nContent-Length: {}\r\nConnection: close\r\n\
807	{headers}\r\n{body}",
808	body.len()
809	);
810	}
811
812	impl Address {
813	fn new(config: &Config, now: Instant) -> Self {
814	Self {
815	connections: 0,
816	pending: 0,
817	joins: Bucket::full(f64::from(config.joins_per_minute), now),
818	failures: VecDeque::new(),
819	strikes: 0,
820	struck: None,
821	locked: None,
822	}
823	}
824
825	/// Forgets failures over a minute old, a lockout that has ended, and strikes after a
826	/// quiet hour.
827	fn refresh(&mut self, now: Instant) {
828	while self
829	.failures
830	.front()
831	.is_some_and(|failed| now.saturating_duration_since(*failed) >= MINUTE)
832	{
833	self.failures.pop_front();
834	}
835	self.locked = self.locked.filter(|until| *until > now);
836	if self
837	.struck
838	.is_some_and(|struck| now.saturating_duration_since(struck) >= HOUR)
839	{
840	self.strikes = 0;
841	self.struck = None;
842	}
843	}
844
845	/// Counts a wrong code: at `limit` in a minute, the lockout it starts.
846	fn fail(&mut self, now: Instant, limit: u32) -> Option<Duration> {
847	self.refresh(now);
848	self.failures.push_back(now);
849	if self.failures.len() < limit as usize {
850	return None;
851	}
852	self.failures.clear();
853	self.strikes += 1;
854	self.struck = Some(now);
855	let lockout = MINUTE
856	.saturating_mul(1 << (self.strikes - 1).min(6))
857	.min(HOUR);
858	self.locked = Some(now + lockout);
859	Some(lockout)
860	}
861
862	/// Whether nothing about this address needs remembering.
863	fn quiet(&self, config: &Config, now: Instant) -> bool {
864	let mut joins = self.joins;
865	let per_minute = f64::from(config.joins_per_minute);
866	self.connections == 0
867	&& self.pending == 0
868	&& self.failures.is_empty()
869	&& self.locked.is_none()
870	&& self.strikes == 0
871	&& joins.take(per_minute, per_minute / 60.0, now).is_ok()
872	}
873	}
874
875	/// A token bucket.
876	#[derive(Clone, Copy)]
877	struct Bucket {
878	tokens: f64,
879	at: Instant,
880	}
881
882	impl Bucket {
883	fn full(capacity: f64, now: Instant) -> Self {
884	Self {
885	tokens: capacity,
886	at: now,
887	}
888	}
889
890	fn refill(&mut self, capacity: f64, rate: f64, now: Instant) {
891	let elapsed = now.saturating_duration_since(self.at).as_secs_f64();
892	self.tokens = (self.tokens + rate * elapsed).min(capacity);
893	self.at = self.at.max(now);
894	}
895
896	/// Takes one token from a bucket of `capacity` refilling at `rate` a second; else how
897	/// long until one is there.
898	fn take(&mut self, capacity: f64, rate: f64, now: Instant) -> Result<(), Duration> {
899	self.refill(capacity, rate, now);
900	if self.tokens >= 1.0 {
901	self.tokens -= 1.0;
902	Ok(())
903	} else {
904	Err(Duration::from_secs_f64((1.0 - self.tokens) / rate))
905	}
906	}
907
908	/// Spends `amount` from a bucket holding a second's worth at `rate`, into debt if need
909	/// be: how long until the debt is repaid.
910	fn spend(&mut self, amount: f64, rate: f64, now: Instant) -> Duration {
911	self.refill(rate, rate, now);
912	self.tokens -= amount;
913	Duration::from_secs_f64((-self.tokens / rate).max(0.0))
914	}
915	}
916
917	/// What waits to be written to one peer, capped in bytes.
918	struct Outbox {
919	/// Where its frames go; none for a polled session's, which its `GET`s take.
920	stream: Option<TcpStream>,
921	queue: Mutex<Queue>,
922	ready: Condvar,
923	most: usize,
924	}
925
926	#[derive(Default)]
927	struct Queue {
928	frames: VecDeque<Vec<u8>>,
929	bytes: usize,
930	closed: bool,
931	}
932
933	impl Outbox {
934	fn new(stream: TcpStream, most: usize) -> Self {
935	Self {
936	stream: Some(stream),
937	queue: Mutex::default(),
938	ready: Condvar::new(),
939	most,
940	}
941	}
942
943	fn mailbox(most: usize) -> Self {
944	Self {
945	stream: None,
946	queue: Mutex::default(),
947	ready: Condvar::new(),
948	most,
949	}
950	}
951
952	/// Everything queued, waiting up to `wait` for something; none once closed and empty.
953	fn take(&self, wait: Duration) -> Option<Vec<u8>> {
954	let mut queue = self.queue.lock().unwrap();
955	if queue.frames.is_empty() && !queue.closed {
956	queue = self.ready.wait_timeout(queue, wait).unwrap().0;
957	}
958	if queue.closed && queue.frames.is_empty() {
959	return None;
960	}
961	queue.bytes = 0;
962	Some(queue.frames.drain(..).flatten().collect())
963	}
964
965	/// Queues `frame`, or hangs up on a peer that reads too slowly to take it.
966	fn push(&self, frame: Vec<u8>) {
967	let mut queue = self.queue.lock().unwrap();
968	if queue.closed {
969	return;
970	}
971	if queue.bytes + frame.len() > self.most {
972	drop(queue);
973	self.abort();
974	return;
975	}
976	queue.bytes += frame.len();
977	queue.frames.push_back(frame);
978	self.ready.notify_one();
979	}
980
981	/// Hangs up after writing what is queued, which tells a refused guest its code was wrong.
982	fn close(&self) {
983	self.queue.lock().unwrap().closed = true;
984	self.ready.notify_one();
985	}
986
987	/// Hangs up now, dropping what is queued.
988	fn abort(&self) {
989	*self.queue.lock().unwrap() = Queue {
990	closed: true,
991	..Queue::default()
992	};
993	self.ready.notify_one();
994	if let Some(stream) = &self.stream {
995	let _ = stream.shutdown(Shutdown::Both);
996	}
997	}
998
999	/// Writes what is queued until closed.
1000	fn drain(&self) {
1001	loop {
1002	let frame = {
1003	let mut queue = self.queue.lock().unwrap();
1004	loop {
1005	if let Some(frame) = queue.frames.pop_front() {
1006	queue.bytes -= frame.len();
1007	break frame;
1008	}
1009	if queue.closed {
1010	drop(queue);
1011	return self.abort();
1012	}
1013	queue = self.ready.wait(queue).unwrap();
1014	}
1015	};
1016	let Some(mut stream) = self.stream.as_ref() else {
1017	return;
1018	};
1019	if stream.write_all(&frame).is_err() {
1020	return self.abort();
1021	}
1022	}
1023	}
1024	}
1025
1026	#[cfg(test)]
1027	mod tests {
1028	use super::*;
1029
1030	/// Ten wrong codes in a minute lock an address out for a minute, then two, then four,
1031	/// and a quiet hour forgives it.
1032	#[test]
1033	fn wrong_codes_lock_out_for_longer_each_time() {
1034	let config = Config::default();
1035	let start = Instant::now();
1036	let mut address = Address::new(&config, start);
1037	let mut now = start;
1038	for expected in [1, 2, 4] {
1039	for _ in 0..9 {
1040	assert_eq!(address.fail(now, 10), None);
1041	}
1042	assert_eq!(address.fail(now, 10), Some(MINUTE * expected));
1043	now += MINUTE * expected;
1044	}
1045	now += HOUR;
1046	address.refresh(now);
1047	assert_eq!(address.strikes, 0);
1048	assert!(address.quiet(&config, now));
1049	}
1050
1051	/// Failures spread out over more than a minute never add up to a lockout.
1052	#[test]
1053	fn slow_wrong_codes_never_lock_out() {
1054	let start = Instant::now();
1055	let mut address = Address::new(&Config::default(), start);
1056	for minute in 0..30 {
1057	for second in [0, 30] {
1058	let now = start + MINUTE * minute + Duration::from_secs(second);
1059	assert_eq!(address.fail(now, 10), None);
1060	}
1061	}
1062	}
1063
1064	#[test]
1065	fn buckets_refill_and_debts_are_waited_out() {
1066	let now = Instant::now();
1067	let mut joins = Bucket::full(2.0, now);
1068	assert!(joins.take(2.0, 1.0, now).is_ok());
1069	assert!(joins.take(2.0, 1.0, now).is_ok());
1070	assert_eq!(joins.take(2.0, 1.0, now), Err(Duration::from_secs(1)));
1071	assert!(joins.take(2.0, 1.0, now + Duration::from_secs(1)).is_ok());
1072
1073	let mut bytes = Bucket::full(100.0, now);
1074	assert_eq!(bytes.spend(100.0, 100.0, now), Duration::ZERO);
1075	assert_eq!(bytes.spend(50.0, 100.0, now), Duration::from_millis(500));
1076	}
1077
1078	#[test]
1079	fn ipv6_addresses_count_by_their_64() {
1080	let listener = TcpListener::bind("127.0.0.1:0").unwrap();
1081	let stream = TcpStream::connect(listener.local_addr().unwrap()).unwrap();
1082	let head = "GET / HTTP/1.1\r\nX-Forwarded-For: 1.2.3.4, 2001:db8:1:2:3:4:5:6\r\n\r\n";
1083	assert_eq!(
1084	crate::peer(&stream, head, true),
1085	"2001:db8:1:2::".parse::<IpAddr>().unwrap()
1086	);
1087	assert_eq!(
1088	crate::peer(&stream, head, false),
1089	"127.0.0.1".parse::<IpAddr>().unwrap()
1090	);
1091	}
1092	}