diff --git a/devolutions-gateway/src/api/fwd.rs b/devolutions-gateway/src/api/fwd.rs index b2099519a5..b89506267c 100644 --- a/devolutions-gateway/src/api/fwd.rs +++ b/devolutions-gateway/src/api/fwd.rs @@ -139,7 +139,9 @@ async fn handle_fwd( with_tls: bool, agent_tunnel_handle: Option>, ) { - 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), diff --git a/devolutions-gateway/src/api/jmux.rs b/devolutions-gateway/src/api/jmux.rs index e40bcf3e66..aa57327a18 100644 --- a/devolutions-gateway/src/api/jmux.rs +++ b/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, diff --git a/devolutions-gateway/src/api/jrec/mod.rs b/devolutions-gateway/src/api/jrec/mod.rs index 32f4b7d143..4dfa03784c 100644 --- a/devolutions-gateway/src/api/jrec/mod.rs +++ b/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,7 +162,11 @@ 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, @@ -170,6 +174,10 @@ async fn handle_jrec_push( 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( /// . 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 +/// . +const RECORDING_ENDED_CLOSE_CODE: u16 = 4011; + /// Deletes a recording stored on this instance #[cfg_attr(feature = "openapi", utoipa::path( delete, diff --git a/devolutions-gateway/src/api/rdp.rs b/devolutions-gateway/src/api/rdp.rs index 29e5d161e2..18050de532 100644 --- a/devolutions-gateway/src/api/rdp.rs +++ b/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>, ) { - 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), diff --git a/devolutions-gateway/src/recording.intent.md b/devolutions-gateway/src/recording.intent.md index 14bb0371e6..0a6b6ca5e5 100644 --- a/devolutions-gateway/src/recording.intent.md +++ b/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. \ No newline at end of file +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. diff --git a/devolutions-gateway/src/recording.rs b/devolutions-gateway/src/recording.rs index 8ece9b9b68..0ebc4c730d 100644 --- a/devolutions-gateway/src/recording.rs +++ b/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 { +pub struct ClientPush { 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 ClientPush +impl ClientPush where S: AsyncRead + AsyncWrite + Unpin + Send + 'static, + E: FnOnce() -> PushEnd, { pub async fn run(self) -> anyhow::Result { 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(); diff --git a/devolutions-gateway/src/ws.rs b/devolutions-gateway/src/ws.rs index 144dbc4c41..07ec38d94a 100644 --- a/devolutions-gateway/src/ws.rs +++ b/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 { + /// 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 { 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>); - (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 { + self.0.get().copied() + } } pub fn handle_messages( @@ -121,14 +144,22 @@ fn with_sentinel( (ws, close_handle) } -fn websocket_compat(ws: transport::Shared) -> impl AsyncRead + AsyncWrite + Unpin + Send + 'static { +fn websocket_compat( + ws: transport::Shared, + 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) diff --git a/testsuite/tests/cli/dgw/jrec_push.rs b/testsuite/tests/cli/dgw/jrec_push.rs new file mode 100644 index 0000000000..1707f886e4 --- /dev/null +++ b/testsuite/tests/cli/dgw/jrec_push.rs @@ -0,0 +1,186 @@ +//! `/jet/jrec/push` tests: how a producer closes its push decides whether the recording waits for a reconnection. + +use std::time::{Duration, Instant}; + +use base64::Engine as _; +use futures_util::{SinkExt as _, StreamExt as _}; +use testsuite::cli::start_dgw; +use testsuite::dgw_config::DgwConfig; +use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::tungstenite::protocol::CloseFrame; +use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode; +use uuid::Uuid; + +/// The close code a producer sends when its recording is over (`RECORDING_ENDED_CLOSE_CODE` in Gateway). +const RECORDING_ENDED_CLOSE_CODE: u16 = 4011; + +/// Far longer than any test runs, so only an ended recording can leave the active list. +const REUSE_WINDOW_SECONDS: u32 = 3600; + +/// An unsigned token; the Gateway under test runs with token validation disabled. +fn unsigned_token(header: serde_json::Value, claims: serde_json::Value) -> String { + let engine = base64::engine::general_purpose::URL_SAFE_NO_PAD; + let header = engine.encode(serde_json::to_vec(&header).expect("serialize header")); + let payload = engine.encode(serde_json::to_vec(&claims).expect("serialize claims")); + let signature = engine.encode(b"signature"); + format!("{header}.{payload}.{signature}") +} + +fn push_token(session_id: Uuid) -> String { + unsigned_token( + serde_json::json!({"alg": "RS256", "cty": "JREC"}), + serde_json::json!({ + "jet_aid": session_id, + "jet_rop": "push", + "jet_reuse": REUSE_WINDOW_SECONDS, + "jti": Uuid::new_v4(), + "nbf": 1_733_669_999i64, + "exp": 3_331_553_599i64, + }), + ) +} + +fn recordings_read_token() -> String { + unsigned_token( + serde_json::json!({"typ": "JWT", "alg": "RS256"}), + serde_json::json!({ + "type": "scope", + "scope": "gateway.recordings.read", + "jti": Uuid::new_v4(), + "iat": 1_733_669_999i64, + "nbf": 1_733_669_999i64, + "exp": 3_331_553_599i64, + }), + ) +} + +async fn active_recordings(http_port: u16) -> Vec { + reqwest::Client::new() + .get(format!("http://127.0.0.1:{http_port}/jet/jrec/list?active=true")) + .bearer_auth(recordings_read_token()) + .timeout(Duration::from_secs(10)) + .send() + .await + .expect("list active recordings") + .error_for_status() + .expect("list active recordings succeeded") + .json() + .await + .expect("parse active recordings") +} + +/// Pushes a few bytes for `session_id`, closes the push with `code`, and waits for Gateway to answer the close. +async fn push_and_close(http_port: u16, session_id: Uuid, code: CloseCode) { + let url = format!( + "ws://127.0.0.1:{http_port}/jet/jrec/push/{session_id}?fileType=webm&token={}", + push_token(session_id) + ); + let (mut socket, _) = tokio_tungstenite::connect_async(url).await.expect("open push"); + socket + .send(Message::Binary(vec![0x1A, 0x45, 0xDF, 0xA3].into())) + .await + .expect("push bytes"); + + let started = Instant::now(); + while !active_recordings(http_port).await.contains(&session_id) { + assert!( + started.elapsed() < Duration::from_secs(10), + "the push never became an active recording" + ); + tokio::time::sleep(Duration::from_millis(50)).await; + } + + socket + .send(Message::Close(Some(CloseFrame { + code, + reason: "".into(), + }))) + .await + .expect("close push"); + while let Some(message) = tokio::time::timeout(Duration::from_secs(10), socket.next()) + .await + .expect("Gateway answered the close") + { + if message.is_err() { + break; + } + } +} + +// oracle: intent devolutions-gateway/src/recording.intent.md §Ending a recording +#[tokio::test] +async fn a_push_closed_as_ended_leaves_the_active_recordings_without_waiting_for_a_reconnection() { + let recordings = tempfile::tempdir().expect("create recording dir"); + let config = DgwConfig::builder() + .disable_token_validation(true) + .recording_path(recordings.path().to_path_buf()) + .build() + .init() + .expect("init gateway config"); + let _gateway = start_dgw(&config).await.expect("start gateway"); + let session_id = Uuid::new_v4(); + + push_and_close( + config.http_port(), + session_id, + CloseCode::from(RECORDING_ENDED_CLOSE_CODE), + ) + .await; + + let started = Instant::now(); + while active_recordings(config.http_port()).await.contains(&session_id) { + assert!( + started.elapsed() < Duration::from_secs(10), + "the ended recording is still waiting for a reconnection" + ); + tokio::time::sleep(Duration::from_millis(50)).await; + } +} + +// oracle: intent devolutions-gateway/src/recording.intent.md §Ending a recording +#[tokio::test] +async fn a_push_closed_as_ended_leaves_the_active_recordings_while_another_waits_for_a_reconnection() { + let recordings = tempfile::tempdir().expect("create recording dir"); + let config = DgwConfig::builder() + .disable_token_validation(true) + .recording_path(recordings.path().to_path_buf()) + .build() + .init() + .expect("init gateway config"); + let _gateway = start_dgw(&config).await.expect("start gateway"); + let waiting = Uuid::new_v4(); + let ended = Uuid::new_v4(); + + // The waiting recording holds Gateway's removal timer at its one-hour deadline. + push_and_close(config.http_port(), waiting, CloseCode::Normal).await; + push_and_close(config.http_port(), ended, CloseCode::from(RECORDING_ENDED_CLOSE_CODE)).await; + + let started = Instant::now(); + while active_recordings(config.http_port()).await.contains(&ended) { + assert!( + started.elapsed() < Duration::from_secs(10), + "the ended recording waits for the other recording's reconnection" + ); + tokio::time::sleep(Duration::from_millis(50)).await; + } +} + +// oracle: intent devolutions-gateway/src/recording.intent.md §Ending a recording +#[tokio::test] +async fn a_push_closed_normally_keeps_the_recording_open_for_a_reconnection() { + let recordings = tempfile::tempdir().expect("create recording dir"); + let config = DgwConfig::builder() + .disable_token_validation(true) + .recording_path(recordings.path().to_path_buf()) + .build() + .init() + .expect("init gateway config"); + let _gateway = start_dgw(&config).await.expect("start gateway"); + let session_id = Uuid::new_v4(); + + // A producer that resizes its session closes the push like this and reconnects with the next clip. + push_and_close(config.http_port(), session_id, CloseCode::Normal).await; + tokio::time::sleep(Duration::from_secs(2)).await; + + assert!(active_recordings(config.http_port()).await.contains(&session_id)); +} diff --git a/testsuite/tests/cli/dgw/mod.rs b/testsuite/tests/cli/dgw/mod.rs index 5bd4e0c49d..23da8ba6c8 100644 --- a/testsuite/tests/cli/dgw/mod.rs +++ b/testsuite/tests/cli/dgw/mod.rs @@ -1,6 +1,7 @@ mod benign_disconnect; mod cli_args; mod heartbeat; +mod jrec_push; mod kdc_proxy; mod preflight; mod recording_search;