1 files changed, 33 insertions, 17 deletions

diff --git a/src/api/login.rs b/src/api/login.rs
index dc9d895..83e7f25 100644
--- a/src/api/login.rs
+++ b/src/api/login.rs
@@ -379,8 +379,11 @@ pub(crate) async fn streams_end(
         .extensions()
         .get::<Token>()
         .map(|token| Arc::clone(&token.claims));
+    // Listening before the answer is made: a change revoking the token that comes while
+    // it is (the node applying what it was behind on) isn't missed.
+    let changes = token.as_ref().map(|_| state.change_events.subscribe());
     let response = next.run(request).await;
-    let Some(claims) = token else {
+    let (Some(claims), Some(changes)) = (token, changes) else {
         return response;
     };
     let streaming = response
@@ -394,15 +397,23 @@ pub(crate) async fn streams_end(
         return response;
     }
     let (parts, body) = response.into_parts();
-    let ended = token_ends(state, claims);
+    let ended = token_ends(state, claims, changes);
     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();
+/// sessions) no longer takes it, looked at once at first (what changed since the
+/// request was let in) and with every change of protocol data in `changes` after.
+async fn token_ends(
+    state: AppState,
+    claims: Arc<Claims>,
+    mut changes: tokio::sync::broadcast::Receiver<Arc<crate::document::store::ChangeSet>>,
+) {
+    if let Some(refused) = refused_now(&state, &claims).await {
+        debug!(user = %claims.user, reason = %refused, "Ended a stream with its token");
+        return;
+    }
     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);
@@ -420,24 +431,29 @@ async fn token_ends(state: AppState, claims: Arc<Claims>) {
                 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 {
+        if let Some(refused) = refused_now(&state, &claims).await {
             debug!(user = %claims.user, reason = %refused, "Ended a stream with its token");
             return;
         }
     }
 }
 
+/// Why the token of `claims` isn't taken any more, if it isn't: by the document, or
+/// this node's ended sessions.
+async fn refused_now(state: &AppState, claims: &Claims) -> Option<token::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)
+    })
+}
+
 /// `POST /auth/refresh`.
 pub(crate) async fn refresh(
     State(state): State<AppState>,