| 1 | use std::collections::VecDeque; |
| 2 | use std::sync::Arc; |
| 3 | use std::time::Duration; |
| 4 | |
| 5 | use anyhow::{Context, Result}; |
| 6 | use tokio::io::{AsyncBufReadExt, AsyncWriteExt}; |
| 7 | use tokio::process::{Child, ChildStdin, ChildStdout}; |
| 8 | use tokio::sync::Mutex as TokioMutex; |
| 9 | |
| 10 | use super::wire::MAX_MCP_RESPONSE_BYTES; |
| 11 | use super::{McpServerConfig, McpTransport}; |
| 12 | use crate::child_env; |
| 13 | |
| 14 | pub(super) struct StdioTransport { |
| 15 | pub(super) child: Arc<TokioMutex<Child>>, |
| 16 | pub(super) stdin: ChildStdin, |
| 17 | pub(super) reader: tokio::io::BufReader<ChildStdout>, |
| 18 | /// Partial frame bytes survive cancellation of the receive future. |
| 19 | pub(super) pending_line: Vec<u8>, |
| 20 | /// Tail of stderr lines from the spawned MCP server. A background task |
| 21 | /// drains the child's stderr into this buffer so a mid-run crash leaves |
| 22 | /// some context behind instead of `Stdio::null` swallowing it. |
| 23 | pub(super) stderr_tail: Arc<StderrTail>, |
| 24 | /// Plugin authority can change in another process while this child is |
| 25 | /// idle. The connection-level cancellation token therefore also owns a |
| 26 | /// process watcher instead of waiting for a later tool call to drop the |
| 27 | /// transport. |
| 28 | pub(super) authority_cancel_watch: Option<tokio::task::JoinHandle<()>>, |
| 29 | /// Holds reviewed executable/script handles for the process lifetime. |
| 30 | pub(super) _reviewed_launch: Option<super::ReviewedStdioLaunch>, |
| 31 | } |
| 32 | |
| 33 | /// How long `StdioTransport::shutdown` waits for the child to exit on SIGTERM |
| 34 | /// before `kill_on_drop` fires SIGKILL. Tuned short so a hung MCP server |
| 35 | /// can't stall TUI exit; well-behaved servers almost always exit within |
| 36 | /// a few hundred ms. |
| 37 | pub(super) const STDIO_SHUTDOWN_GRACE: Duration = Duration::from_millis(2_000); |
| 38 | |
| 39 | /// How many lines of MCP-server stderr to keep around for crash diagnostics. |
| 40 | /// Bounded so a chatty server can't grow this without limit; large enough to |
| 41 | /// catch typical Node/Python startup or panic output. |
| 42 | const STDERR_TAIL_CAPACITY: usize = 64; |
| 43 | |
| 44 | /// Bounded ring buffer for the most recent stderr lines from a spawned MCP |
| 45 | /// server. Used by `StdioTransport` to surface server-side context when the |
| 46 | /// transport read side fails (server crashed, exited early, etc). |
| 47 | #[derive(Default)] |
| 48 | pub(super) struct StderrTail { |
| 49 | lines: TokioMutex<VecDeque<String>>, |
| 50 | } |
| 51 | |
| 52 | impl StderrTail { |
| 53 | pub(super) fn new() -> Arc<Self> { |
| 54 | Arc::new(Self { |
| 55 | lines: TokioMutex::new(VecDeque::with_capacity(STDERR_TAIL_CAPACITY)), |
| 56 | }) |
| 57 | } |
| 58 | |
| 59 | pub(super) async fn push(&self, line: String) { |
| 60 | let mut buf = self.lines.lock().await; |
| 61 | if buf.len() >= STDERR_TAIL_CAPACITY { |
| 62 | buf.pop_front(); |
| 63 | } |
| 64 | buf.push_back(line); |
| 65 | } |
| 66 | |
| 67 | async fn snapshot(&self) -> Vec<String> { |
| 68 | self.lines.lock().await.iter().cloned().collect() |
| 69 | } |
| 70 | } |
| 71 | |
| 72 | impl StdioTransport { |
| 73 | pub(super) fn spawn( |
| 74 | server_name: &str, |
| 75 | command: &str, |
| 76 | config: &McpServerConfig, |
| 77 | cancel_token: tokio_util::sync::CancellationToken, |
| 78 | ) -> Result<Self> { |
| 79 | let reviewed_launch = if let Some(reviewed_plugin) = config.reviewed_plugin.as_ref() { |
| 80 | // This is deliberately the last trust check before constructing |
| 81 | // and spawning the lazy stdio child. It re-reads only the |
| 82 | // Codewhale-owned plugin bundle, never user MCP/provider config or |
| 83 | // credential files, and fails closed on any content/capability |
| 84 | // drift after pool construction. |
| 85 | Some(reviewed_plugin.prepare_stdio_launch( |
| 86 | server_name, |
| 87 | command, |
| 88 | &config.args, |
| 89 | config.cwd.as_deref(), |
| 90 | )?) |
| 91 | } else { |
| 92 | None |
| 93 | }; |
| 94 | let mut cmd = reviewed_launch.as_ref().map_or_else( |
| 95 | || { |
| 96 | let mut command_process = tokio::process::Command::new(command); |
| 97 | command_process.args(&config.args); |
| 98 | command_process |
| 99 | }, |
| 100 | |launch| { |
| 101 | let mut command_process = tokio::process::Command::new(&launch.command); |
| 102 | command_process.args(&launch.args); |
| 103 | command_process |
| 104 | }, |
| 105 | ); |
| 106 | crate::utils::suppress_tokio_console_window(&mut cmd); |
| 107 | cmd.stdin(std::process::Stdio::piped()) |
| 108 | .stdout(std::process::Stdio::piped()) |
| 109 | .stderr(std::process::Stdio::piped()) |
| 110 | .kill_on_drop(true); |
| 111 | let launch_cwd = reviewed_launch |
| 112 | .as_ref() |
| 113 | .and_then(|launch| launch.cwd.as_ref()) |
| 114 | .or(config.cwd.as_ref().filter(|_| reviewed_launch.is_none())); |
| 115 | if let Some(cwd) = launch_cwd { |
| 116 | cmd.current_dir(cwd); |
| 117 | } |
| 118 | #[cfg(unix)] |
| 119 | if let Some(cwd_fd) = reviewed_launch |
| 120 | .as_ref() |
| 121 | .and_then(|launch| launch.cwd_fd.as_ref()) |
| 122 | { |
| 123 | use std::os::fd::AsRawFd as _; |
| 124 | let fd = cwd_fd.as_raw_fd(); |
| 125 | // SAFETY: the closure calls only async-signal-safe `fchdir` on an |
| 126 | // inherited directory descriptor before exec. |
| 127 | unsafe { |
| 128 | cmd.pre_exec(move || { |
| 129 | if libc::fchdir(fd) == 0 { |
| 130 | Ok(()) |
| 131 | } else { |
| 132 | Err(std::io::Error::last_os_error()) |
| 133 | } |
| 134 | }); |
| 135 | } |
| 136 | } |
| 137 | |
| 138 | // Expand `${NAME}` placeholders so secret env values can be sourced |
| 139 | // from the process environment instead of being stored in cleartext |
| 140 | // in the MCP config. The child env is allowlist-sanitized below, so |
| 141 | // these vars would not otherwise be inherited by the child. |
| 142 | let expanded_env = super::expanded_mcp_stdio_env(config) |
| 143 | .with_context(|| format!("MCP server '{server_name}' env expansion failed"))?; |
| 144 | |
| 145 | // User-configured MCP keeps the compatibility-oriented Node/Python |
| 146 | // bootstrap allowlist (#1244). Reviewed plugins receive only the base |
| 147 | // secret-scrubbed child environment plus their explicitly reviewed |
| 148 | // mappings, so namespaces such as NPM_CONFIG_* are never inherited |
| 149 | // ambiently across the consent boundary. |
| 150 | if let Some(reviewed_plugin) = config.reviewed_plugin.as_ref() { |
| 151 | cmd.env_clear(); |
| 152 | for (key, value) in child_env::sanitized_plugin_mcp_env_from( |
| 153 | reviewed_plugin.host_environment.entries().iter().cloned(), |
| 154 | child_env::string_map_env(&expanded_env), |
| 155 | ) { |
| 156 | cmd.env(key, value); |
| 157 | } |
| 158 | } else { |
| 159 | child_env::apply_to_tokio_command_mcp( |
| 160 | &mut cmd, |
| 161 | child_env::string_map_env(&expanded_env), |
| 162 | ); |
| 163 | } |
| 164 | |
| 165 | let mut child = cmd.spawn().map_err(|error| { |
| 166 | let message = if error.kind() == std::io::ErrorKind::NotFound |
| 167 | && super::is_node_command(command) |
| 168 | && launch_cwd.is_none_or(|directory| directory.is_dir()) |
| 169 | { |
| 170 | format!("MCP server {server_name} could not start because Node.js was not found. Install Node.js from https://nodejs.org/ and restart Codewhale with node on PATH. Built-in Computer Use requires Node.js 20 or newer.") |
| 171 | } else if config.reviewed_plugin.is_some() { |
| 172 | format!( |
| 173 | "MCP stdio spawn failed (transport=stdio server={server_name} reviewed-plugin argv_count={} env_count={})", |
| 174 | config.args.len(), |
| 175 | expanded_env.len(), |
| 176 | ) |
| 177 | } else { |
| 178 | let env_keys: Vec<&str> = expanded_env.keys().map(String::as_str).collect(); |
| 179 | format!( |
| 180 | "MCP stdio spawn failed (transport=stdio server={server_name} cmd={command:?} args={:?} env_keys={env_keys:?})", |
| 181 | config.args, |
| 182 | ) |
| 183 | }; |
| 184 | anyhow::Error::new(error).context(message) |
| 185 | })?; |
| 186 | |
| 187 | let stdin = child.stdin.take().context("Failed to get MCP stdin")?; |
| 188 | let stdout = child.stdout.take().context("Failed to get MCP stdout")?; |
| 189 | let stderr = child.stderr.take().context("Failed to get MCP stderr")?; |
| 190 | |
| 191 | // Drain stderr into a bounded ring buffer so a crash mid-run leaves |
| 192 | // diagnostic breadcrumbs instead of disappearing into `Stdio::null`. |
| 193 | // The task exits naturally when the child closes its stderr |
| 194 | // (kill_on_drop / exit / explicit shutdown). |
| 195 | let stderr_tail = StderrTail::new(); |
| 196 | { |
| 197 | let tail = Arc::clone(&stderr_tail); |
| 198 | // A reviewed plugin child receives environment-backed values that |
| 199 | // are intentionally absent from its manifest. Still drain its |
| 200 | // stderr to avoid blocking, but do not retain or surface arbitrary |
| 201 | // child output that could echo those credentials into a chat or |
| 202 | // persisted transcript. |
| 203 | let capture_lines = config.reviewed_plugin.is_none(); |
| 204 | tokio::spawn(async move { |
| 205 | let mut lines = tokio::io::BufReader::new(stderr).lines(); |
| 206 | while let Ok(Some(line)) = lines.next_line().await { |
| 207 | if capture_lines { |
| 208 | tail.push(line).await; |
| 209 | } |
| 210 | } |
| 211 | }); |
| 212 | } |
| 213 | |
| 214 | let child = Arc::new(TokioMutex::new(child)); |
| 215 | let authority_cancel_watch = config.reviewed_plugin.as_ref().map(|_| { |
| 216 | let watched_child = Arc::clone(&child); |
| 217 | tokio::spawn(async move { |
| 218 | cancel_token.cancelled().await; |
| 219 | terminate_child_for_authority_change(&watched_child).await; |
| 220 | }) |
| 221 | }); |
| 222 | |
| 223 | Ok(Self { |
| 224 | child, |
| 225 | stdin, |
| 226 | reader: tokio::io::BufReader::new(stdout), |
| 227 | pending_line: Vec::new(), |
| 228 | stderr_tail, |
| 229 | authority_cancel_watch, |
| 230 | _reviewed_launch: reviewed_launch, |
| 231 | }) |
| 232 | } |
| 233 | } |
| 234 | |
| 235 | /// Format the captured stderr tail for inclusion in an error message. Empty |
| 236 | /// tails return `None` so the caller can fall back to its original message. |
| 237 | async fn format_stderr_context(tail: &StderrTail) -> Option<String> { |
| 238 | let lines = tail.snapshot().await; |
| 239 | if lines.is_empty() { |
| 240 | return None; |
| 241 | } |
| 242 | Some(format!( |
| 243 | "MCP server stderr (last {} line{}):\n{}", |
| 244 | lines.len(), |
| 245 | if lines.len() == 1 { "" } else { "s" }, |
| 246 | lines.join("\n"), |
| 247 | )) |
| 248 | } |
| 249 | |
| 250 | /// Best-effort SIGTERM. On Unix uses `libc::kill`; on Windows there's no |
| 251 | /// equivalent so we let `kill_on_drop` (TerminateProcess) handle it via the |
| 252 | /// subsequent Drop. Returns whether a signal was actually sent. |
| 253 | fn send_sigterm(child: &Child) -> bool { |
| 254 | #[cfg(unix)] |
| 255 | { |
| 256 | if let Some(pid) = child.id() { |
| 257 | // SAFETY: pid was just obtained from `child.id()`. `libc::kill` |
| 258 | // with `SIGTERM` is async-signal-safe and never observes invalid |
| 259 | // memory. Worst case (pid wrap / process already gone) returns |
| 260 | // ESRCH, which we deliberately ignore. |
| 261 | unsafe { |
| 262 | let _ = libc::kill(pid as i32, libc::SIGTERM); |
| 263 | } |
| 264 | return true; |
| 265 | } |
| 266 | false |
| 267 | } |
| 268 | #[cfg(not(unix))] |
| 269 | { |
| 270 | let _ = child; |
| 271 | false |
| 272 | } |
| 273 | } |
| 274 | |
| 275 | async fn terminate_child_for_authority_change(child: &Arc<TokioMutex<Child>>) { |
| 276 | let mut child = child.lock().await; |
| 277 | terminate_child(&mut child).await; |
| 278 | } |
| 279 | |
| 280 | async fn terminate_child(child: &mut Child) { |
| 281 | // Reap an already-exited child before resolving its PID. Until it is |
| 282 | // reaped, the OS cannot recycle that identity; after it is reaped there is |
| 283 | // nothing left to signal. This avoids a PID-only watcher ever targeting an |
| 284 | // unrelated process after rapid PID reuse. |
| 285 | if child.try_wait().is_ok_and(|status| status.is_some()) { |
| 286 | return; |
| 287 | } |
| 288 | |
| 289 | #[cfg(unix)] |
| 290 | send_sigterm(child); |
| 291 | |
| 292 | #[cfg(not(unix))] |
| 293 | let _ = child.start_kill(); |
| 294 | |
| 295 | match tokio::time::timeout(STDIO_SHUTDOWN_GRACE, child.wait()).await { |
| 296 | Ok(Ok(_)) => {} |
| 297 | Ok(Err(_)) | Err(_) => { |
| 298 | // SIGTERM is advisory. Revocation and explicit shutdown must not |
| 299 | // leave the reviewed child alive indefinitely. |
| 300 | let _ = child.start_kill(); |
| 301 | let _ = child.wait().await; |
| 302 | } |
| 303 | } |
| 304 | } |
| 305 | |
| 306 | #[async_trait::async_trait] |
| 307 | impl McpTransport for StdioTransport { |
| 308 | async fn send(&mut self, mut msg: Vec<u8>) -> Result<()> { |
| 309 | msg.push(b'\n'); |
| 310 | self.stdin.write_all(&msg).await?; |
| 311 | self.stdin.flush().await?; |
| 312 | Ok(()) |
| 313 | } |
| 314 | |
| 315 | /// Non-blocking liveness probe: a reaped child means the transport is |
| 316 | /// dead even though the `Ready` flag is still set (#6187). The sync |
| 317 | /// trait contract forbids awaiting the lock, so a contended lock reads |
| 318 | /// as alive — the read side observes the death on the next call. |
| 319 | fn probe_dead(&self) -> bool { |
| 320 | match self.child.try_lock() { |
| 321 | Ok(mut child) => matches!(child.try_wait(), Ok(Some(_))), |
| 322 | Err(_) => false, |
| 323 | } |
| 324 | } |
| 325 | |
| 326 | async fn recv(&mut self) -> Result<Vec<u8>> { |
| 327 | loop { |
| 328 | // Bounded read: a server emitting a newline-free multi-GB "line" |
| 329 | // must not OOM us (read_line is unbounded). |
| 330 | let bytes = match read_line_capped( |
| 331 | &mut self.reader, |
| 332 | &mut self.pending_line, |
| 333 | MAX_MCP_RESPONSE_BYTES, |
| 334 | ) |
| 335 | .await |
| 336 | { |
| 337 | Ok(b) => b, |
| 338 | Err(err) => { |
| 339 | if let Some(stderr) = format_stderr_context(&self.stderr_tail).await { |
| 340 | anyhow::bail!("Stdio transport read error: {err}\n{stderr}"); |
| 341 | } |
| 342 | return Err(err.into()); |
| 343 | } |
| 344 | }; |
| 345 | if bytes == 0 { |
| 346 | // Let the stderr drain task catch up before snapshotting, and |
| 347 | // name the exit status: a reviewed plugin's stderr is never |
| 348 | // retained, so the status is the only reason the operator |
| 349 | // gets when the child dies before the handshake (#5916). |
| 350 | tokio::task::yield_now().await; |
| 351 | let exit = self.child.lock().await.try_wait().ok().flatten(); |
| 352 | let exit = exit.map_or_else(String::new, |status| format!(" ({status})")); |
| 353 | if let Some(stderr) = format_stderr_context(&self.stderr_tail).await { |
| 354 | anyhow::bail!("Stdio transport closed{exit}\n{stderr}"); |
| 355 | } |
| 356 | anyhow::bail!("Stdio transport closed{exit}"); |
| 357 | } |
| 358 | |
| 359 | let line_bytes = std::mem::take(&mut self.pending_line); |
| 360 | let line = String::from_utf8_lossy(&line_bytes); |
| 361 | let trimmed = line.trim(); |
| 362 | if trimmed.is_empty() { |
| 363 | continue; |
| 364 | } |
| 365 | |
| 366 | return Ok(trimmed.as_bytes().to_vec()); |
| 367 | } |
| 368 | } |
| 369 | |
| 370 | /// Send SIGTERM and wait up to `STDIO_SHUTDOWN_GRACE` for graceful exit, |
| 371 | /// then force termination and reap the child as the backstop. |
| 372 | async fn shutdown(&mut self) { |
| 373 | let mut child = self.child.lock().await; |
| 374 | terminate_child(&mut child).await; |
| 375 | } |
| 376 | } |
| 377 | |
| 378 | /// Session changes can drop a pool without explicitly awaiting shutdown. |
| 379 | /// Keep the owned child alive for the same bounded cleanup as explicit |
| 380 | /// shutdown, so servers can release input and recording resources. Runtime |
| 381 | /// teardown still drops the cleanup future and invokes `kill_on_drop`. |
| 382 | impl Drop for StdioTransport { |
| 383 | fn drop(&mut self) { |
| 384 | if let Some(watch) = self.authority_cancel_watch.take() { |
| 385 | watch.abort(); |
| 386 | } |
| 387 | if let Ok(runtime) = tokio::runtime::Handle::try_current() { |
| 388 | let child = Arc::clone(&self.child); |
| 389 | let reviewed_launch = self._reviewed_launch.take(); |
| 390 | runtime.spawn(async move { |
| 391 | let _reviewed_launch = reviewed_launch; |
| 392 | let mut child = child.lock().await; |
| 393 | terminate_child(&mut child).await; |
| 394 | }); |
| 395 | return; |
| 396 | } |
| 397 | if let Ok(mut child) = self.child.try_lock() |
| 398 | && !child.try_wait().is_ok_and(|status| status.is_some()) |
| 399 | { |
| 400 | send_sigterm(&child); |
| 401 | } |
| 402 | } |
| 403 | } |
| 404 | |
| 405 | /// Continue one newline-terminated line in caller-owned `out`, aborting if it |
| 406 | /// exceeds `max` bytes. Cancellation retains consumed bytes; the caller clears |
| 407 | /// the buffer only after receiving a complete frame. Returns the total bytes |
| 408 | /// accumulated; 0 means EOF. |
| 409 | async fn read_line_capped<R>( |
| 410 | reader: &mut R, |
| 411 | out: &mut Vec<u8>, |
| 412 | max: usize, |
| 413 | ) -> std::io::Result<usize> |
| 414 | where |
| 415 | R: tokio::io::AsyncBufRead + Unpin, |
| 416 | { |
| 417 | use tokio::io::AsyncBufReadExt; |
| 418 | loop { |
| 419 | let (chunk, consumed, done) = { |
| 420 | let available = reader.fill_buf().await?; |
| 421 | if available.is_empty() { |
| 422 | (Vec::new(), 0usize, true) |
| 423 | } else if let Some(pos) = available.iter().position(|&b| b == b'\n') { |
| 424 | (available[..=pos].to_vec(), pos + 1, true) |
| 425 | } else { |
| 426 | (available.to_vec(), available.len(), false) |
| 427 | } |
| 428 | }; |
| 429 | if consumed > 0 { |
| 430 | reader.consume(consumed); |
| 431 | } |
| 432 | out.extend_from_slice(&chunk); |
| 433 | if out.len() > max { |
| 434 | return Err(std::io::Error::new( |
| 435 | std::io::ErrorKind::InvalidData, |
| 436 | format!("MCP stdio line exceeded {max} bytes"), |
| 437 | )); |
| 438 | } |
| 439 | if done { |
| 440 | break; |
| 441 | } |
| 442 | } |
| 443 | Ok(out.len()) |
| 444 | } |
| 445 | |
| 446 | #[cfg(test)] |
| 447 | mod read_cap_tests { |
| 448 | use super::read_line_capped; |
| 449 | |
| 450 | #[tokio::test] |
| 451 | async fn cancelled_partial_read_preserves_next_frame() { |
| 452 | use futures_util::FutureExt; |
| 453 | use tokio::io::AsyncWriteExt; |
| 454 | let (mut writer, reader) = tokio::io::duplex(4096); |
| 455 | let mut reader = tokio::io::BufReader::new(reader); |
| 456 | let prefix = br#"{"jsonrpc":"2.0","id":"1","result":"#; |
| 457 | writer.write_all(prefix).await.unwrap(); |
| 458 | let mut pending = Vec::new(); |
| 459 | // Poll through the consumed prefix to Pending, then drop the future. |
| 460 | assert!( |
| 461 | read_line_capped(&mut reader, &mut pending, 1024) |
| 462 | .now_or_never() |
| 463 | .is_none() |
| 464 | ); |
| 465 | assert_eq!(pending, prefix); |
| 466 | writer.write_all(b"null}\n").await.unwrap(); |
| 467 | read_line_capped(&mut reader, &mut pending, 1024) |
| 468 | .await |
| 469 | .unwrap(); |
| 470 | let first: serde_json::Value = |
| 471 | serde_json::from_slice(&std::mem::take(&mut pending)).unwrap(); |
| 472 | assert_eq!(first["id"], "1"); |
| 473 | writer |
| 474 | .write_all(b"{\"id\":\"2\",\"result\":true}\n") |
| 475 | .await |
| 476 | .unwrap(); |
| 477 | read_line_capped(&mut reader, &mut pending, 1024) |
| 478 | .await |
| 479 | .unwrap(); |
| 480 | let second: serde_json::Value = serde_json::from_slice(&pending).unwrap(); |
| 481 | assert_eq!(second["id"], "2"); |
| 482 | assert_eq!(second["result"], true); |
| 483 | } |
| 484 | |
| 485 | #[tokio::test] |
| 486 | async fn resumed_frame_still_enforces_cap_at_newline() { |
| 487 | use futures_util::FutureExt; |
| 488 | use tokio::io::AsyncWriteExt; |
| 489 | let (mut writer, reader) = tokio::io::duplex(4096); |
| 490 | let mut reader = tokio::io::BufReader::new(reader); |
| 491 | let mut pending = Vec::new(); |
| 492 | writer.write_all(b"1234").await.unwrap(); |
| 493 | assert!( |
| 494 | read_line_capped(&mut reader, &mut pending, 6) |
| 495 | .now_or_never() |
| 496 | .is_none() |
| 497 | ); |
| 498 | writer.write_all(b"567\n").await.unwrap(); |
| 499 | assert_eq!( |
| 500 | read_line_capped(&mut reader, &mut pending, 6) |
| 501 | .await |
| 502 | .unwrap_err() |
| 503 | .kind(), |
| 504 | std::io::ErrorKind::InvalidData |
| 505 | ); |
| 506 | } |
| 507 | |
| 508 | #[tokio::test] |
| 509 | async fn reads_a_line_and_reports_eof() { |
| 510 | let data = b"hello\nworld\n".to_vec(); |
| 511 | let mut reader = tokio::io::BufReader::new(std::io::Cursor::new(data)); |
| 512 | let mut out = Vec::new(); |
| 513 | assert_eq!( |
| 514 | read_line_capped(&mut reader, &mut out, 1024).await.unwrap(), |
| 515 | 6 |
| 516 | ); |
| 517 | assert_eq!(out, b"hello\n"); |
| 518 | out.clear(); |
| 519 | assert_eq!( |
| 520 | read_line_capped(&mut reader, &mut out, 1024).await.unwrap(), |
| 521 | 6 |
| 522 | ); |
| 523 | assert_eq!(out, b"world\n"); |
| 524 | out.clear(); |
| 525 | // EOF. |
| 526 | assert_eq!( |
| 527 | read_line_capped(&mut reader, &mut out, 1024).await.unwrap(), |
| 528 | 0 |
| 529 | ); |
| 530 | } |
| 531 | |
| 532 | #[tokio::test] |
| 533 | async fn aborts_on_newline_free_line_over_cap() { |
| 534 | let data = vec![b'x'; 4096]; // no newline |
| 535 | let mut reader = tokio::io::BufReader::new(std::io::Cursor::new(data)); |
| 536 | let mut out = Vec::new(); |
| 537 | let err = read_line_capped(&mut reader, &mut out, 1024) |
| 538 | .await |
| 539 | .unwrap_err(); |
| 540 | assert_eq!(err.kind(), std::io::ErrorKind::InvalidData); |
| 541 | } |
| 542 | } |
| 543 |