1 files changed, 77 insertions, 0 deletions
diff --git a/src/api/login.rs b/src/api/login.rs
index ec0c9cf..57c1859 100644
--- a/src/api/login.rs
+++ b/src/api/login.rs
@@ -362,6 +362,81 @@ pub(crate) async fn login_key(
     answer(&state, &claims)
 }
 
+/// A stream answered to a request with a token (server-sent events, the store's
+/// streams) ends with the token: when it expires, or as soon as the document says it
+/// isn't taken any more (`token::still_taken`: its user disabled or removed, its tokens
+/// revoked, its session ended, its issuer banned), looked at when a change touches
+/// protocol data, or when a session ends on this node. The page reconnects with the
+/// token it has then, if any; a refreshed one is a new token, so a stream lasts at most
+/// a token's lifetime. Answers that aren't streams pass as they are.
+pub(crate) async fn streams_end(
+    State(state): State<AppState>,
+    request: Request<Body>,
+    next: axum::middleware::Next,
+) -> Response {
+    let token = request
+        .extensions()
+        .get::<Token>()
+        .map(|token| Arc::clone(&token.claims));
+    let response = next.run(request).await;
+    let Some(claims) = token else {
+        return response;
+    };
+    let streaming = response
+        .headers()
+        .get(axum::http::header::CONTENT_TYPE)
+        .and_then(|kind| kind.to_str().ok())
+        .is_some_and(|kind| {
+            kind.starts_with("text/event-stream") || kind.starts_with("application/x-ndjson")
+        });
+    if !streaming {
+        return response;
+    }
+    let (parts, body) = response.into_parts();
+    let ended = token_ends(state, claims);
+    let body = futures::StreamExt::take_until(body.into_data_stream(), Box::pin(ended));
+    Response::from_parts(parts, Body::from_stream(body))
+}
+
+/// When the token of `claims` ends: it expires, or the document (or this node's ended
+/// sessions) no longer takes it.
+async fn token_ends(state: AppState, claims: Arc<Claims>) {
+    let mut changes = state.change_events.subscribe();
+    let left = claims.exp.saturating_sub(state.clock.now_millis().as_u64());
+    let expiry = tokio::time::sleep(std::time::Duration::from_millis(left));
+    tokio::pin!(expiry);
+    loop {
+        let session_ended = state.session_ended.notified();
+        tokio::select! {
+            () = &mut expiry => return,
+            () = session_ended => {}
+            change = changes.recv() => match change {
+                // Protocol data, the root (a switch) among it.
+                Ok(set) if !set.paths.iter().any(|path| {
+                    path.first().is_none_or(|first| first == crate::document::meta::ROOT)
+                }) => continue,
+                Ok(_) | Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {}
+                Err(tokio::sync::broadcast::error::RecvError::Closed) => return,
+            },
+        }
+        let refused = {
+            let replica = state.doc.read().await;
+            token::still_taken(&replica.doc, &claims).err()
+        }
+        .or_else(|| {
+            state
+                .ended_sessions
+                .lock_or_recover()
+                .contains_key(&claims.session)
+                .then_some(token::Refused::Ended)
+        });
+        if let Some(refused) = refused {
+            debug!(user = %claims.user, reason = %refused, "Ended a stream with its token");
+            return;
+        }
+    }
+}
+
 /// `POST /auth/refresh`.
 pub(crate) async fn refresh(
     State(state): State<AppState>,
@@ -477,6 +552,8 @@ async fn end(state: &AppState, claims: &Claims, by: &str) -> Result<(), ApiError
         let mut ended = state.ended_sessions.lock_or_recover();
         ended.retain(|_, until| *until > now);
         ended.insert(claims.session.clone(), until);
+        drop(ended);
+        state.session_ended.notify_waiters();
     }
     info!(user = %claims.user, session = %claims.session, by = %by, recorded, "Ended a session");
     Ok(())