Skip to main content

max / makenotwork

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