fix(dgw): end a recording its producer closes as ended

devolutions-gateway/src/api/fwd.rs

@@ -139,7 +139,9 @@ async fn handle_fwd(
     with_tls: bool,
     agent_tunnel_handle: Option<Arc<agent_tunnel::AgentTunnelHandle>>,
 ) {
-    let (stream, close_handle) = crate::ws::handle(
+    let crate::ws::HandledWebSocket {
+        stream, close_handle, ..
+    } = crate::ws::handle(
         ws,
         crate::ws::KeepAliveShutdownSignal(shutdown_signal),
         Duration::from_secs(conf.debug.ws_keep_alive_interval),
@@ -452,7 +454,11 @@ async fn fwd_http(
         let subscriber_tx = state.subscriber_tx;
 
         client_ws.on_upgrade(move |client_ws| {
-            let (client_stream, client_close_handle) = crate::ws::handle(
+            let crate::ws::HandledWebSocket {
+                stream: client_stream,
+                close_handle: client_close_handle,
+                ..
+            } = crate::ws::handle(
                 client_ws,
                 crate::ws::KeepAliveShutdownSignal(shutdown_signal),
                 Duration::from_secs(conf.debug.ws_keep_alive_interval),

devolutions-gateway/src/api/jmux.rs

@@ -61,7 +61,9 @@ async fn handle_socket(
     source_addr: SocketAddr,
     keep_alive_interval: Duration,
 ) {
-    let (stream, close_handle) = crate::ws::handle(
+    let crate::ws::HandledWebSocket {
+        stream, close_handle, ..
+    } = crate::ws::handle(
         ws,
         crate::ws::KeepAliveShutdownSignal(shutdown_signal),
         keep_alive_interval,

devolutions-gateway/src/api/jrec/mod.rs

@@ -33,7 +33,7 @@ use crate::api::heartbeat::recording_storage_health;
 use crate::artifacts::JrecArtifacts;
 use crate::extract::{JrecToken, RecordingDeleteScope, RecordingsReadScope};
 use crate::http::{HttpError, HttpErrorBuilder};
-use crate::recording::{PushOutcome, RecordingMessageSender};
+use crate::recording::{PushEnd, PushOutcome, RecordingMessageSender};
 use crate::token::{JrecTokenClaims, RecordingFileType, RecordingOperation};
 
 /// Read chunk size when streaming a finished session ZIP from the temp file.
@@ -162,14 +162,22 @@ async fn handle_jrec_push(
     source_addr: SocketAddr,
     keep_alive_interval: Duration,
 ) {
-    let (stream, close_handle) = crate::ws::handle(
+    let crate::ws::HandledWebSocket {
+        stream,
+        close_handle,
+        received_close,
+    } = crate::ws::handle(
         ws,
         crate::ws::KeepAliveShutdownSignal(shutdown_signal.clone()),
         keep_alive_interval,
     );
 
     let result = crate::recording::ClientPush::builder()
         .client_stream(stream)
+        .push_end(move || match received_close.code() {
+            Some(RECORDING_ENDED_CLOSE_CODE) => PushEnd::RecordingEnded,
+            _ => PushEnd::MayReconnect,
+        })
         .recordings(recordings)
         .claims(claims)
         .file_type(file_type)
@@ -200,6 +208,13 @@ async fn handle_jrec_push(
 /// <https://developer.mozilla.org/en-US/docs/Web/API/CloseEvent/code>.
 const STORAGE_FULL_CLOSE_CODE: u16 = 4010;
 
+/// WebSocket close code a producer sends on `/jrec/push/{id}` when its recording is over, not interrupted.
+///
+/// Any other close, or no close frame, keeps the recording open until its reconnection window passes.
+/// Codes in 4000-4999 are reserved for private application use per
+/// <https://developer.mozilla.org/en-US/docs/Web/API/CloseEvent/code>.
+const RECORDING_ENDED_CLOSE_CODE: u16 = 4011;
+
 /// Deletes a recording stored on this instance
 #[cfg_attr(feature = "openapi", utoipa::path(
     delete,

devolutions-gateway/src/api/rdp.rs

@@ -72,7 +72,9 @@ async fn handle_socket(
     synthetic_kdc_registry: crate::credential_injection::SyntheticKdcRegistry,
     agent_tunnel_handle: Option<Arc<agent_tunnel::AgentTunnelHandle>>,
 ) {
-    let (stream, close_handle) = crate::ws::handle(
+    let crate::ws::HandledWebSocket {
+        stream, close_handle, ..
+    } = crate::ws::handle(
         ws,
         crate::ws::KeepAliveShutdownSignal(shutdown_signal),
         Duration::from_secs(conf.debug.ws_keep_alive_interval),

devolutions-gateway/src/recording.intent.md

@@ -1,10 +1,23 @@
 # Recording
-Gateway recording is a continuous stream of bytes pushed from the client who pocesses the valid recording token.
-The url is `/jet/jrec/push/{sessionId}?fileType={fileType}`.
-We expect the `fileType` to be one of the following file types: `webm`, `cast`, `trp` and `slog`, which must be specified.
-When connection is established with request of recordings for the session, if recording is not enabled within a short period of time, the connection will be closed with indication of violation of the recording policy.
+A recording is a continuous stream of bytes pushed by a client that holds a valid recording token.
+A recording contains one or more clips.
+The URL is `/jet/jrec/push/{sessionId}?fileType={fileType}`.
+The `fileType` is required and is one of `webm`, `asciicast`, `trp` and `slog`.
+When a session whose token requires recording starts and no recording starts for it within a short time, Gateway closes the session as a recording policy violation.
 
 
 # Artifacts
-Artifacts are files that are not recordings, currently only have `ai-analysis` with combination to `slog` file type.
-Currently, we do not have an api endpoint to push artifacts, but we will have one in the future.
+Artifacts are files that are not recordings; the only kind today is `ai-analysis`, stored with the `slog` file type.
+An endpoint to push artifacts is wanted later.
+
+
+# Ending a recording
+We use WebSocket close code 4011 to indicate that the producer actively terminates the recording.
+As soon as 4011 is received:
+- the recording is no longer active;
+- the shadow viewers are notified;
+- the associated session is closed subject to the recording policy.
+
+Every other close, including a WebSocket close without a code or with code 1000, and a connection lost without any close, is treated as a network interruption or a clip switch.
+In that case, the recording stays active until its reconnection timeout passes.
+Gateway does not refuse a later push for the same session with a new token after 4011.

devolutions-gateway/src/recording.rs

@@ -26,6 +26,7 @@ use crate::streaming::{RecordingStreamState, StreamLifecycle};
 use crate::token::{JrecTokenClaims, RecordingFileType};
 
 const DISCONNECTED_TTL_EXTRA_LEEWAY: Duration = Duration::from_secs(10);
+
 const BUFFER_WRITER_SIZE: usize = 64 * 1024;
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -74,19 +75,31 @@ pub enum PushOutcome {
     StorageFull,
 }
 
+/// How a producer ended its push stream.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum PushEnd {
+    /// The recording is over: it ends at once, so live viewers see the end right after the last bytes.
+    RecordingEnded,
+    /// The producer may come back, for example with a new clip after a resize: the recording waits for its reconnection window.
+    MayReconnect,
+}
+
 #[derive(TypedBuilder)]
-pub struct ClientPush<S> {
+pub struct ClientPush<S, E> {
     recordings: RecordingMessageSender,
     claims: JrecTokenClaims,
     client_stream: S,
     file_type: RecordingFileType,
     session_id: Uuid,
     shutdown_signal: ShutdownSignal,
+    /// Tells how the producer ended the push, once `client_stream` is over.
+    push_end: E,
 }
 
-impl<S> ClientPush<S>
+impl<S, E> ClientPush<S, E>
 where
     S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
+    E: FnOnce() -> PushEnd,
 {
     pub async fn run(self) -> anyhow::Result<PushOutcome> {
         let Self {
@@ -96,6 +109,7 @@ where
             file_type,
             session_id,
             mut shutdown_signal,
+            push_end,
         } = self;
 
         if session_id != claims.jet_aid {
@@ -196,9 +210,14 @@ where
             Err(e) => Err(anyhow::Error::new(e).context(format!("failed to open file at {recording_file}"))),
         };
 
-        info!(?res, "Recording finished");
+        // UNTESTED: a recording its producer ended ends at once, even when writing its last bytes failed.
+        // The producer is not coming back, and no test makes writing the recording file fail.
+        // When changing this, keep calling `push_end` whatever `res` holds.
+        let end = push_end();
+
+        info!(?res, ?end, "Recording finished");
 
-        recordings.disconnect(session_id).await.context("disconnect")?;
+        recordings.disconnect(session_id, end).await.context("disconnect")?;
 
         res
     }
@@ -359,6 +378,7 @@ enum RecordingManagerMessage {
     },
     Disconnect {
         id: Uuid,
+        end: PushEnd,
     },
     GetState {
         id: Uuid,
@@ -396,7 +416,9 @@ impl fmt::Debug for RecordingManagerMessage {
                 .field("id", id)
                 .field("kind", kind)
                 .finish_non_exhaustive(),
-            RecordingManagerMessage::Disconnect { id } => f.debug_struct("Disconnect").field("id", id).finish(),
+            RecordingManagerMessage::Disconnect { id, end } => {
+                f.debug_struct("Disconnect").field("id", id).field("end", end).finish()
+            }
             RecordingManagerMessage::GetState { id, channel: _ } => {
                 f.debug_struct("GetState").field("id", id).finish_non_exhaustive()
             }
@@ -459,9 +481,9 @@ impl RecordingMessageSender {
         rx.await.context("couldn't receive AddArtifact result")?
     }
 
-    async fn disconnect(&self, id: Uuid) -> anyhow::Result<()> {
+    async fn disconnect(&self, id: Uuid, end: PushEnd) -> anyhow::Result<()> {
         self.channel
-            .send(RecordingManagerMessage::Disconnect { id })
+            .send(RecordingManagerMessage::Disconnect { id, end })
             .await
             .ok()
             .context("couldn't send Remove message")
@@ -728,7 +750,7 @@ impl RecordingManagerTask {
         Ok((recording_file, ClipProgress(stream_state)))
     }
 
-    async fn handle_disconnect(&mut self, id: Uuid) -> anyhow::Result<()> {
+    async fn handle_disconnect(&mut self, id: Uuid, end: PushEnd) -> anyhow::Result<()> {
         let Some(ongoing) = self.ongoing_recordings.get_mut(&id) else {
             return Err(anyhow::anyhow!("unknown recording for ID {id}"));
         };
@@ -740,7 +762,12 @@ impl RecordingManagerTask {
         let end_time = time::OffsetDateTime::now_utc().unix_timestamp();
 
         // Viewers and the recording share this deadline, so a stream never ends while its recording continues.
-        let reconnect_deadline = tokio::time::Instant::now() + ongoing.disconnected_ttl;
+        // An ended recording waits for no reconnection, as if its token disallowed reuse.
+        let disconnected_ttl = match end {
+            PushEnd::RecordingEnded => Duration::ZERO,
+            PushEnd::MayReconnect => ongoing.disconnected_ttl,
+        };
+        let reconnect_deadline = tokio::time::Instant::now() + disconnected_ttl;
         ongoing.state = OnGoingRecordingState::LastSeen { reconnect_deadline };
 
         // Re-read from disk: an artifact may have been added since this recording connected.
@@ -937,8 +964,8 @@ async fn recording_manager_task(
                     RecordingManagerMessage::AddArtifact { id, kind, channel } => {
                         let _ = channel.send(manager.handle_add_artifact(id, kind).await);
                     },
-                    RecordingManagerMessage::Disconnect { id } => {
-                        if let Err(e) = manager.handle_disconnect(id).await {
+                    RecordingManagerMessage::Disconnect { id, end } => {
+                        if let Err(e) = manager.handle_disconnect(id, end).await {
                             error!(error = format!("{e:#}"), "handle_disconnect");
                         }
 
@@ -1008,8 +1035,8 @@ async fn recording_manager_task(
         };
 
         debug!(?msg, "Received message");
-        if let RecordingManagerMessage::Disconnect { id } = msg {
-            if let Err(e) = manager.handle_disconnect(id).await {
+        if let RecordingManagerMessage::Disconnect { id, end } = msg {
+            if let Err(e) = manager.handle_disconnect(id, end).await {
                 error!(error = format!("{e:#}"), "handle_disconnect");
             }
             manager.ongoing_recordings.remove(&id);
@@ -1159,7 +1186,10 @@ mod tests {
         }
 
         async fn disconnect(&self, id: Uuid) {
-            self.sender.disconnect(id).await.expect("disconnect");
+            self.sender
+                .disconnect(id, PushEnd::MayReconnect)
+                .await
+                .expect("disconnect");
             // Messages are processed in order, so this waits for the disconnect to be handled.
             self.sender.get_count().await.expect("sync with manager");
         }
@@ -1346,6 +1376,42 @@ mod tests {
         assert_eq!(state.borrow().lifecycle, StreamLifecycle::Ended);
     }
 
+    // oracle: intent devolutions-gateway/src/recording.intent.md §Ending a recording
+    #[tokio::test]
+    async fn a_recording_its_producer_ended_ends_without_waiting_for_a_reconnection() {
+        let harness = Harness::start();
+        let id = Uuid::new_v4();
+        harness
+            .sender
+            .connect(id, WEBM, Duration::from_secs(3600))
+            .await
+            .expect("connect");
+        let mut stream = harness
+            .sender
+            .subscribe_to_stream(id)
+            .await
+            .expect("subscribe")
+            .expect("ongoing recording");
+
+        harness
+            .sender
+            .disconnect(id, PushEnd::RecordingEnded)
+            .await
+            .expect("disconnect");
+
+        // Far less than the one-hour reconnection window the token allows.
+        tokio::time::timeout(
+            Duration::from_secs(5),
+            stream.wait_for(|state| state.lifecycle == StreamLifecycle::Ended),
+        )
+        .await
+        .expect("stream ended without waiting for a reconnection")
+        .expect("stream state alive");
+        // oracle: mutant the manager holding an ended recording for the token's reconnection window
+        assert!(harness.sender.get_state(id).await.expect("get state").is_none());
+        assert!(!harness.sender.active_recordings.cloned().contains(&id));
+    }
+
     #[tokio::test]
     async fn viewers_and_the_recording_share_the_reconnect_deadline() {
         let harness = Harness::start();

devolutions-gateway/src/ws.rs

@@ -1,4 +1,5 @@
 use core::{future, time};
+use std::sync::{Arc, OnceLock};
 
 use axum::extract::ws::{self, CloseFrame, WebSocket};
 use bytes::Bytes;
@@ -15,18 +16,40 @@ impl transport::KeepAliveShutdown for KeepAliveShutdownSignal {
     }
 }
 
+/// A WebSocket wrapped by [`handle`].
+pub struct HandledWebSocket<S> {
+    /// The WebSocket as a byte stream.
+    pub stream: S,
+    pub close_handle: transport::CloseWebSocketHandle,
+    /// The close code the client sends, known once `stream` has returned end of file for its close frame.
+    pub received_close: ReceivedClose,
+}
+
 /// Spawns a keep-alive task and wraps the WebSocket into a type implementing AsyncRead and AsyncWrite.
 pub fn handle(
     ws: WebSocket,
     shutdown_signal: impl transport::KeepAliveShutdown,
     keep_alive_interval: time::Duration,
-) -> (
-    impl AsyncRead + AsyncWrite + Unpin + Send + 'static,
-    transport::CloseWebSocketHandle,
-) {
+) -> HandledWebSocket<impl AsyncRead + AsyncWrite + Unpin + Send + 'static> {
     let (ws, close_handle) = with_sentinel(ws, shutdown_signal, keep_alive_interval);
+    let received_close = ReceivedClose::default();
+
+    HandledWebSocket {
+        stream: websocket_compat(ws, received_close.clone()),
+        close_handle,
+        received_close,
+    }
+}
+
+/// The close code a client sent when it closed the WebSocket.
+#[derive(Clone, Default)]
+pub struct ReceivedClose(Arc<OnceLock<u16>>);
 
-    (websocket_compat(ws), close_handle)
+impl ReceivedClose {
+    /// The code of the close frame the client sent, if it sent one with a code.
+    pub fn code(&self) -> Option<u16> {
+        self.0.get().copied()
+    }
 }
 
 pub fn handle_messages(
@@ -121,14 +144,22 @@ fn with_sentinel(
     (ws, close_handle)
 }
 
-fn websocket_compat(ws: transport::Shared<WebSocket>) -> impl AsyncRead + AsyncWrite + Unpin + Send + 'static {
+fn websocket_compat(
+    ws: transport::Shared<WebSocket>,
+    received_close: ReceivedClose,
+) -> impl AsyncRead + AsyncWrite + Unpin + Send + 'static {
     let ws_compat = ws
-        .filter_map(|item| {
+        .filter_map(move |item| {
             item.map(|msg| match msg {
                 ws::Message::Text(s) => Some(transport::WsReadMsg::Payload(Bytes::from(s))),
                 ws::Message::Binary(data) => Some(transport::WsReadMsg::Payload(data)),
                 ws::Message::Ping(_) | ws::Message::Pong(_) => None,
-                ws::Message::Close(_) => Some(transport::WsReadMsg::Close),
+                ws::Message::Close(frame) => {
+                    if let Some(frame) = frame {
+                        let _ = received_close.0.set(frame.code);
+                    }
+                    Some(transport::WsReadMsg::Close)
+                }
             })
             .transpose()
             .pipe(future::ready)