diff --git a/src/core/config.rs b/src/core/config.rs index 7447234..5a09502 100644 --- a/src/core/config.rs +++ b/src/core/config.rs @@ -4,6 +4,7 @@ const DEFAULT_MAX_CONCURRENT_EXECUTIONS: usize = 16; const DEFAULT_MAX_CONNECTED_DEVICES: usize = 10; const DEFAULT_MAX_EXECUTION_SECS: u64 = 3600; const DEFAULT_MAX_COMMAND_LENGTH: usize = 65_536; +const DEFAULT_MAX_RESULT_BYTES: usize = 10 * 1024 * 1024; const DEFAULT_PING_INTERVAL_SECS: u64 = 30; const DEFAULT_PING_TIMEOUT_SECS: u64 = 10; @@ -12,6 +13,7 @@ pub struct Config { pub max_connected_devices: usize, pub max_execution_time: Duration, pub max_command_length: usize, + pub max_result_bytes: usize, pub ping_interval: Duration, pub ping_timeout: Duration, } @@ -32,6 +34,7 @@ impl Config { DEFAULT_MAX_EXECUTION_SECS, )), max_command_length: parse_env("MAX_COMMAND_LENGTH", DEFAULT_MAX_COMMAND_LENGTH), + max_result_bytes: parse_env("MAX_RESULT_BYTES", DEFAULT_MAX_RESULT_BYTES), ping_interval: Duration::from_secs(parse_env( "PING_INTERVAL_SECS", DEFAULT_PING_INTERVAL_SECS, diff --git a/src/transport/socket/loop_state.rs b/src/transport/socket/loop_state.rs index 6f47192..68403d2 100644 --- a/src/transport/socket/loop_state.rs +++ b/src/transport/socket/loop_state.rs @@ -43,7 +43,7 @@ pub async fn run_loop( } } message = socket.next() => { - match handle_socket_message(socket, message, &pending).await { + match handle_socket_message(socket, message, &pending, config.max_result_bytes).await { SocketAction::Break => break, SocketAction::ClearPongDeadline => awaiting_pong = false, SocketAction::Continue => {} @@ -65,6 +65,7 @@ async fn handle_socket_message( socket: &mut DuplexStream, message: Option>, pending: &PendingReplies, + max_result_bytes: usize, ) -> SocketAction { match message { Some(Ok(Message::Pong(_))) => SocketAction::ClearPongDeadline, @@ -79,11 +80,13 @@ async fn handle_socket_message( if let Ok(ClientMsg::ExecResult { exec_id, exit_code, - stdout, - stderr, + mut stdout, + mut stderr, }) = serde_json::from_str(&text) && let Some((_, reply)) = pending.remove(&exec_id) { + truncate_utf8(&mut stdout, max_result_bytes); + truncate_utf8(&mut stderr, max_result_bytes); reply .send(ExecResult { exit_code, @@ -100,6 +103,19 @@ async fn handle_socket_message( } } +fn truncate_utf8(text: &mut String, max_bytes: usize) { + if text.len() <= max_bytes { + return; + } + + let mut end = max_bytes; + while end > 0 && !text.is_char_boundary(end) { + end -= 1; + } + + text.truncate(end); +} + async fn dispatch( socket: &mut DuplexStream, pending: &PendingReplies,