api: a stream listens for its token's end from before it is answered

Review 4, F9: a stream answered to a token subscribed to the document's changes only
when its body was first polled, and looked at the token only after a change came: one
revoking it that came between the request's check and the subscription (a node
applying what it was behind on) went unseen until the next change of protocol data, or
the token's expiry.

Now `streams_end` subscribes before the answer is made, and the token is looked at once
as the stream begins (`refused_now`, the document and this node's ended sessions), then
with every change of protocol data as before. A request that isn't answered with a
stream drops its subscription with the answer.

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>,