max / makenotwork

5.6 KB · 155 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	author: None,
80	}
81	}
82
83	/// The SSE wire format is a contract with the client island, so assert on the
84	/// bytes rather than on the builder.
85	///
86	/// `Event` exposes no accessor for its encoded form, so this reads the `Debug`
87	/// output, which embeds the real buffer. If an axum upgrade ever makes that
88	/// opaque these assertions go quiet rather than failing, so
89	/// [`debug_output_still_exposes_the_buffer`] guards the technique itself.
90	fn wire(event: &ChatEvent) -> String {
91	format!("{:?}", frame(event))
92	}
93
94	#[test]
95	fn debug_output_still_exposes_the_buffer() {
96	assert!(
97	wire(&ChatEvent::Gap).contains("event: gap\\n"),
98	"Event's Debug no longer embeds the encoded frame, so the wire-format \
99	assertions in this module have stopped checking anything"
100	);
101	}
102
103	#[test]
104	fn every_variant_frames_with_its_own_event_name() {
105	for (event, name) in [
106	(ChatEvent::Message(message(1)), "message"),
107	(ChatEvent::Delete { id: MessageId(1) }, "delete"),
108	(
109	ChatEvent::Purge {
110	author_id: UserId(Uuid::nil()),
111	},
112	"purge",
113	),
114	(ChatEvent::Gap, "gap"),
115	] {
116	// Match the SSE `event:` field specifically. A bare `contains(name)`
117	// would also be satisfied by the tagged JSON in the data field, which
118	// means it would pass even if the event name were never set.
119	assert!(
120	wire(&event).contains(&format!("event: {name}\\n")),
121	"expected an SSE event field named {name}, got: {}",
122	wire(&event)
123	);
124	}
125	}
126
127	#[test]
128	fn body_html_is_carried_verbatim() {
129	// The host renders through docengine before publishing; the transport
130	// must not re-encode or re-escape what it was given.
131	let mut m = message(1);
132	m.body_html = "<em>hi</em>".into();
133	let json = serde_json::to_string(&ChatEvent::Message(m)).unwrap();
134	assert!(json.contains(r"<em>hi</em>"), "{json}");
135	}
136
137	#[test]
138	fn a_frame_never_contains_a_bare_newline() {
139	// SSE is newline-delimited. A payload with a raw newline would split one
140	// event into two malformed ones, so the JSON encoding has to be the thing
141	// that protects the framing.
142	let mut m = message(1);
143	m.body_html = "line one\nline two".into();
144	let json = serde_json::to_string(&ChatEvent::Message(m)).unwrap();
145	assert!(
146	!json.contains('\n'),
147	"raw newline reached the frame: {json}"
148	);
149	assert!(
150	json.contains(r"\n"),
151	"newline should survive escaped: {json}"
152	);
153	}
154	}
155