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