api: a stream opened with a token ends with it

A token was checked when a request came, and a stream answered to it (the dashboard's
events, a subscription's links and messages, the store's streams) then went on for as
long as the client kept it open: past the token's expiry, past its user's tokens being
revoked, its user disabled or removed, its session ended.

Now a middleware inside `login::bearer` (`login::streams_end`) ends a streaming answer
(`text/event-stream`, `application/x-ndjson`) to a request with a token at the token's
deadline, its expiry, or as soon as the document no longer takes it
(`token::still_taken`, the checks `verify` makes of the document, now its own
function): looked at when a change touches protocol data (the root, a switch, among
it), and when a session ends on the node itself in a swarm without a registry
(`AppState::session_ended`). No polling: a timer for the expiry, and the changes as
they come. The client reconnects with the token it has then; a refreshed token is a
new one, so a stream lasts a token's lifetime at most (12 hours by default).

`futures` is a dependency of every build now (for `take_until`), not only of libp2p's
and Tor's.

2 files changed, 82 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(())
diff --git a/src/api/mod.rs b/src/api/mod.rs
index 6631bc9..97724ff 100644
--- a/src/api/mod.rs
+++ b/src/api/mod.rs
@@ -121,6 +121,11 @@ pub(crate) fn router(state: AppState, auth_state: AuthState, debug_api: bool) ->
     ))
     // Who has a password here (`/auth/keys`: a user's last key goes only then).
     .layer(axum::Extension(auth_state))
+    // A stream answered to a token ends with it (expired, revoked, its session ended).
+    .layer(middleware::from_fn_with_state(
+        state.clone(),
+        login::streams_end,
+    ))
     // Outside authentication: a token, if the request brings one, names the user.
     .layer(middleware::from_fn_with_state(state.clone(), login::bearer))
     // Outermost: pages of other origins, and their preflights, which bring no