max / makenotwork

5.6 KB · 151 lines History Blame Raw

1	//! SSE push notification endpoint for SyncKit subscribers.
2	//!
3	//! Clients connect to `GET /api/sync/subscribe?app_id={uuid}` and receive
4	//! zero-data "changed" events whenever another device pushes changes.
5	//! No payload is included; the client should pull to get actual data.
6	//! This preserves E2E encryption (server never sends plaintext).
7
8	use std::convert::Infallible;
9	use std::sync::atomic::Ordering;
10	use std::time::Duration;
11
12	use axum::{
13	extract::{Query, State},
14	response::sse::{Event, KeepAlive, Sse},
15	};
16	use serde::Deserialize;
17	use tokio_stream::StreamExt;
18	use tokio_stream::wrappers::BroadcastStream;
19
20	use crate::{
21	constants,
22	db::{SyncAppId, UserId},
23	error::{AppError, Result},
24	synckit_auth::SyncUser,
25	};
26
27	/// Drop guard that decrements the per-user SSE connection counter.
28	struct SseConnectionGuard {
29	sse_connections: std::sync::Arc<dashmap::DashMap<UserId, std::sync::atomic::AtomicUsize>>,
30	sync_notify:
31	std::sync::Arc<dashmap::DashMap<(SyncAppId, UserId), tokio::sync::broadcast::Sender<i64>>>,
32	user_id: UserId,
33	app_id: SyncAppId,
34	}
35
36	impl Drop for SseConnectionGuard {
37	fn drop(&mut self) {
38	// Decrement the connection counter, then drop the entry if it hit zero.
39	// `remove_if` runs the predicate under the shard lock, so a connection
40	// arriving concurrently cannot be counted into an entry we then remove.
41	// It also avoids the same-shard deadlock a get() + remove() pair would
42	// hit, which is what the earlier read/drop/remove split was working
43	// around.
44	if let Some(counter) = self.sse_connections.get(&self.user_id) {
45	counter.value().fetch_sub(1, Ordering::AcqRel);
46	}
47	self.sse_connections.remove_if(&self.user_id, \|_, counter\| {
48	counter.load(Ordering::Acquire) == 0
49	});
50
51	// Prune the sync_notify channel once the last receiver drops. This must
52	// be atomic with respect to `Sync::subscribe_channel`: reading
53	// receiver_count(), dropping the guard, then removing lets a subscriber
54	// land in between and hold a receiver on a sender we then discard. The
55	// next push lazily creates a different channel, and that client
56	// silently receives no further SSE events until it reconnects.
57	let key = (self.app_id, self.user_id);
58	self.sync_notify
59	.remove_if(&key, \|_, sender\| sender.receiver_count() == 0);
60	}
61	}
62
63	/// Stream wrapper that holds an SSE connection guard. When the stream is
64	/// dropped (client disconnects), the guard decrements the connection counter.
65	struct GuardedStream<S> {
66	inner: S,
67	_guard: SseConnectionGuard,
68	}
69
70	impl<S: tokio_stream::Stream + Unpin> tokio_stream::Stream for GuardedStream<S> {
71	type Item = S::Item;
72
73	fn poll_next(
74	mut self: std::pin::Pin<&mut Self>,
75	cx: &mut std::task::Context<'_>,
76	) -> std::task::Poll<Option<Self::Item>> {
77	std::pin::Pin::new(&mut self.inner).poll_next(cx)
78	}
79	}
80
81	#[derive(Deserialize)]
82	pub(super) struct SubscribeQuery {
83	pub app_id: SyncAppId,
84	}
85
86	/// SSE endpoint for real-time sync push notifications.
87	///
88	/// `GET /api/sync/subscribe?app_id={uuid}`: JWT auth required.
89	///
90	/// Returns an SSE stream that emits `event: changed` (with empty data) whenever
91	/// a push is made to the same app+user. The client should pull on each event.
92	/// A keepalive comment is sent every 30 seconds to prevent proxy timeouts.
93	#[tracing::instrument(skip_all, name = "synckit::subscribe")]
94	pub(super) async fn sync_subscribe(
95	State(sync): State<crate::Sync>,
96	sync_user: SyncUser,
97	Query(query): Query<SubscribeQuery>,
98	) -> Result<Sse<impl tokio_stream::Stream<Item = std::result::Result<Event, Infallible>>>> {
99	// Validate that the requested app_id matches the JWT's app_id
100	if query.app_id != sync_user.app_id {
101	return Err(AppError::BadRequest(
102	"app_id does not match authenticated session".to_string(),
103	));
104	}
105
106	// Enforce per-user SSE connection limit
107	let counter = sync
108	.sse_connections
109	.entry(sync_user.user_id)
110	.or_insert_with(\|\| std::sync::atomic::AtomicUsize::new(0));
111	let current = counter.value().fetch_add(1, Ordering::AcqRel);
112	if current >= constants::SYNCKIT_MAX_SSE_CONNECTIONS_PER_USER {
113	counter.value().fetch_sub(1, Ordering::AcqRel);
114	return Err(AppError::BadRequest(
115	"Too many concurrent SSE connections".to_string(),
116	));
117	}
118
119	let guard = SseConnectionGuard {
120	sse_connections: sync.sse_connections.clone(),
121	sync_notify: sync.sync_notify.clone(),
122	user_id: sync_user.user_id,
123	app_id: sync_user.app_id,
124	};
125
126	// Get or create the broadcast channel for this app+user
127	let rx = sync.subscribe_channel(sync_user.app_id, sync_user.user_id);
128
129	let stream = BroadcastStream::new(rx).filter_map(\|result\| match result {
130	// Carry the new max seq so a client already at this cursor can skip the
131	// pull. Clients that don't read `seq` just pull as before.
132	Ok(seq) => Some(Ok(Event::default()
133	.event("changed")
134	.data(format!("{{\"seq\":{seq}}}")))),
135	Err(_) => None, // Lagged, skip missed events, client will pull anyway
136	});
137
138	// Wrap stream with the connection guard, when the client disconnects and
139	// the stream is dropped, the guard's Drop impl decrements the counter.
140	let stream = GuardedStream {
141	inner: stream,
142	_guard: guard,
143	};
144
145	Ok(Sse::new(stream).keep_alive(
146	KeepAlive::new()
147	.interval(Duration::from_secs(30))
148	.text("keepalive"),
149	))
150	}
151