truncate output

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,

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<Result<Message, rocket_ws::result::Error>>,
     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,
+                mut stdout,
-                stderr,
+                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,