max / makenotwork

5.5 KB · 154 lines History Blame Raw

1	//! Axum SSE adapter. Behind the `axum` feature; the rest of the crate has no
2	//! framework dependency.
3	//!
4	//! SSE rather than a WebSocket, on purpose. Upstream traffic is one `POST` per
5	//! message typed, so bidirectionality buys nothing, and a `POST` passes through
6	//! the CSRF, rate-limit, and session middleware both consumers already run. It
7	//! also avoids adding a `connect-src` CSP directive on the Multithreaded side,
8	//! which currently has none and falls back to `default-src 'self'`.
9
10	use std::convert::Infallible;
11	use std::time::Duration;
12
13	use axum::response::sse::{Event, KeepAlive, Sse};
14	use tokio_stream::Stream;
15
16	use crate::event::ChatEvent;
17	use crate::stream::ChatStream;
18
19	/// Keepalive interval.
20	///
21	/// Cloudflare fronts makenot.work and idles a connection out at 100 seconds. A
22	/// 30 second comment clears that with room for a missed tick, and matches what
23	/// the SyncKit SSE endpoint already runs in production.
24	pub const KEEPALIVE_INTERVAL: Duration = Duration::from_secs(30);
25
26	/// Serialize one event into an SSE frame.
27	///
28	/// The event name is carried in the SSE `event:` field so the client can dispatch
29	/// without parsing the payload first, and the payload is the same tagged JSON, so
30	/// a client that prefers to read `type` can.
31	fn frame(event: &ChatEvent) -> Event {
32	let data = serde_json::to_string(event).unwrap_or_else(\|_\| {
33	// A ChatEvent is plain data and cannot fail to serialize. If that ever
34	// changes, a dead connection is worse than a frame the client ignores.
35	tracing::error!("chat event failed to serialize");
36	r#"{"type":"gap"}"#.to_owned()
37	});
38	Event::default().event(event.name()).data(data)
39	}
40
41	/// Wrap a [`ChatStream`] as an axum SSE response.
42	///
43	/// The stream owns the hub subscription, so dropping the response (which axum
44	/// does when the client disconnects) releases the connection slot and prunes the
45	/// room if it was the last listener. There is nothing for a caller to remember to
46	/// clean up.
47	pub fn sse_response(
48	stream: ChatStream,
49	) -> Sse<impl Stream<Item = Result<Event, Infallible>> + Send> {
50	let events = async_stream::stream! {
51	let mut stream = stream;
52	while let Some(event) = stream.next().await {
53	yield Ok(frame(&event));
54	}
55	};
56
57	Sse::new(events).keep_alive(
58	KeepAlive::new()
59	.interval(KEEPALIVE_INTERVAL)
60	.text("keepalive"),
61	)
62	}
63
64	#[cfg(test)]
65	mod tests {
66	use super::*;
67	use crate::ids::{MessageId, RoomId, UserId};
68	use crate::message::Message;
69	use uuid::Uuid;
70
71	fn message(id: i64) -> Message {
72	Message {
73	id: MessageId(id),
74	room_id: RoomId(Uuid::nil()),
75	author_id: UserId(Uuid::nil()),
76	body_html: "hi".into(),
77	created_at: 0,
78	nonce: None,
79	}
80	}
81
82	/// The SSE wire format is a contract with the client island, so assert on the
83	/// bytes rather than on the builder.
84	///
85	/// `Event` exposes no accessor for its encoded form, so this reads the `Debug`
86	/// output, which embeds the real buffer. If an axum upgrade ever makes that
87	/// opaque these assertions go quiet rather than failing, so
88	/// [`debug_output_still_exposes_the_buffer`] guards the technique itself.
89	fn wire(event: &ChatEvent) -> String {
90	format!("{:?}", frame(event))
91	}
92
93	#[test]
94	fn debug_output_still_exposes_the_buffer() {
95	assert!(
96	wire(&ChatEvent::Gap).contains("event: gap\\n"),
97	"Event's Debug no longer embeds the encoded frame, so the wire-format \
98	assertions in this module have stopped checking anything"
99	);
100	}
101
102	#[test]
103	fn every_variant_frames_with_its_own_event_name() {
104	for (event, name) in [
105	(ChatEvent::Message(message(1)), "message"),
106	(ChatEvent::Delete { id: MessageId(1) }, "delete"),
107	(
108	ChatEvent::Purge {
109	author_id: UserId(Uuid::nil()),
110	},
111	"purge",
112	),
113	(ChatEvent::Gap, "gap"),
114	] {
115	// Match the SSE `event:` field specifically. A bare `contains(name)`
116	// would also be satisfied by the tagged JSON in the data field, which
117	// means it would pass even if the event name were never set.
118	assert!(
119	wire(&event).contains(&format!("event: {name}\\n")),
120	"expected an SSE event field named {name}, got: {}",
121	wire(&event)
122	);
123	}
124	}
125
126	#[test]
127	fn body_html_is_carried_verbatim() {
128	// The host renders through docengine before publishing; the transport
129	// must not re-encode or re-escape what it was given.
130	let mut m = message(1);
131	m.body_html = "<em>hi</em>".into();
132	let json = serde_json::to_string(&ChatEvent::Message(m)).unwrap();
133	assert!(json.contains(r#"<em>hi</em>"#), "{json}");
134	}
135
136	#[test]
137	fn a_frame_never_contains_a_bare_newline() {
138	// SSE is newline-delimited. A payload with a raw newline would split one
139	// event into two malformed ones, so the JSON encoding has to be the thing
140	// that protects the framing.
141	let mut m = message(1);
142	m.body_html = "line one\nline two".into();
143	let json = serde_json::to_string(&ChatEvent::Message(m)).unwrap();
144	assert!(
145	!json.contains('\n'),
146	"raw newline reached the frame: {json}"
147	);
148	assert!(
149	json.contains(r"\n"),
150	"newline should survive escaped: {json}"
151	);
152	}
153	}
154