use std::{
fs,
io::{IsTerminal as _, Read},
path::PathBuf,
sync::mpsc,
thread,
time::Duration,
};
pub(super) const STDIN_BYTE_LIMIT: usize = 512 * 1024;
const STDIN_IDLE_TIMEOUT: Duration = Duration::from_secs(30);
fn read_bounded_progress<R: Read>(
reader: &mut R,
limit: usize,
mut on_progress: impl FnMut(),
) -> Result<String, StdinReadError> {
let mut bytes: Vec<u8> = Vec::new();
let mut buf = [0u8; 8192];
loop {
let n = reader.read(&mut buf).map_err(StdinReadError::Io)?;
if n == 0 {
break;
}
bytes.extend_from_slice(&buf[..n]);
if bytes.len() > limit {
return Err(StdinReadError::OverLimit { limit });
}
on_progress();
}
String::from_utf8(bytes).map_err(StdinReadError::Utf8)
}
#[cfg(test)]
fn read_bounded<R: Read>(reader: &mut R, limit: usize) -> Result<String, StdinReadError> {
read_bounded_progress(reader, limit, || {})
}
fn read_stdin_bounded_with_deadline<R: Read + Send + 'static>(
reader: R,
limit: usize,
idle: Duration,
) -> Result<String, StdinReadError> {
enum Signal {
Progress,
Done(Result<String, StdinReadError>),
}
let (tx, rx) = mpsc::channel::<Signal>();
thread::spawn(move || {
let mut reader = reader;
let progress = || {
let _ = tx.send(Signal::Progress);
};
let _ = tx.send(Signal::Done(read_bounded_progress(
&mut reader,
limit,
progress,
)));
});
loop {
match rx.recv_timeout(idle) {
Ok(Signal::Progress) => continue,
Ok(Signal::Done(result)) => return result,
Err(mpsc::RecvTimeoutError::Timeout) => return Err(StdinReadError::Idle),
Err(mpsc::RecvTimeoutError::Disconnected) => {
return Err(StdinReadError::Io(std::io::Error::other(
"stdin reader failed",
)));
}
}
}
}
#[derive(Debug)]
enum StdinReadError {
OverLimit { limit: usize },
Idle,
Io(std::io::Error),
Utf8(std::string::FromUtf8Error),
}
impl std::fmt::Display for StdinReadError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::OverLimit { limit } => write!(
f,
"stdin input exceeds the {limit}-byte limit; use --file for larger input"
),
Self::Idle => write!(
f,
"no input arrived on stdin within {:?}; pass a prompt, use --file, or pipe data",
STDIN_IDLE_TIMEOUT
),
Self::Io(e) => write!(f, "reading stdin: {e}"),
Self::Utf8(e) => write!(f, "stdin is not valid UTF-8: {e}"),
}
}
}
impl std::error::Error for StdinReadError {}
pub(super) fn input(
value: Option<String>,
file: Option<PathBuf>,
) -> Result<String, Box<dyn std::error::Error>> {
match (value, file) {
(Some(value), None) => Ok(value),
(None, Some(path)) => Ok(fs::read_to_string(path)?),
(Some(_), Some(_)) => Err("provide a prompt or --file, not both".into()),
(None, None) if !std::io::stdin().is_terminal() => {
let buffer = read_stdin_bounded_with_deadline(
std::io::stdin(),
STDIN_BYTE_LIMIT,
STDIN_IDLE_TIMEOUT,
)?;
let trimmed = buffer.trim();
if trimmed.is_empty() {
Err("a prompt or --file is required".into())
} else {
Ok(trimmed.to_string())
}
}
(None, None) => Err("a prompt or --file is required".into()),
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Cursor;
#[test]
fn explicit_inputs_win_and_conflicts_error() {
assert_eq!(input(Some("hi".into()), None).unwrap(), "hi");
let path = std::env::temp_dir().join(format!("saya-qin-{}.txt", std::process::id()));
std::fs::write(&path, "from file").unwrap();
assert_eq!(input(None, Some(path.clone())).unwrap(), "from file");
let _ = std::fs::remove_file(&path);
assert!(input(Some("a".into()), Some("b".into())).is_err());
}
#[test]
fn over_limit_input_is_refused_naming_the_limit() {
let over: Vec<u8> = vec![b'x'; STDIN_BYTE_LIMIT + 1];
let mut reader = Cursor::new(over);
let err = read_bounded(&mut reader, STDIN_BYTE_LIMIT)
.expect_err("one byte over the limit must be refused");
let StdinReadError::OverLimit { limit } = &err else {
panic!("expected OverLimit, got {err:?}");
};
assert_eq!(*limit, STDIN_BYTE_LIMIT);
let rendered = err.to_string();
assert!(
rendered.contains(&STDIN_BYTE_LIMIT.to_string()),
"error must name the limit in bytes: {rendered}"
);
}
#[test]
fn at_limit_input_is_accepted() {
let exactly: Vec<u8> = vec![b'y'; STDIN_BYTE_LIMIT];
let mut reader = Cursor::new(exactly);
let got = read_bounded(&mut reader, STDIN_BYTE_LIMIT).expect("at-limit is allowed");
assert_eq!(got.len(), STDIN_BYTE_LIMIT);
}
#[test]
fn empty_input_is_eof() {
let mut reader = Cursor::new(Vec::<u8>::new());
assert_eq!(read_bounded(&mut reader, STDIN_BYTE_LIMIT).unwrap(), "");
}
#[test]
fn multibyte_utf8_split_across_reads_decodes() {
let bytes = "café".as_bytes().to_vec();
let mut reader = Cursor::new(bytes);
let got = read_bounded(&mut reader, STDIN_BYTE_LIMIT).unwrap();
assert_eq!(got, "café");
}
#[test]
fn invalid_utf8_is_an_error() {
let mut reader = Cursor::new(vec![0xff, 0xfe, 0xfd]);
assert!(matches!(
read_bounded(&mut reader, STDIN_BYTE_LIMIT),
Err(StdinReadError::Utf8(_))
));
}
#[test]
fn silent_reader_gives_up_within_idle_deadline() {
struct Silent;
impl Read for Silent {
fn read(&mut self, _: &mut [u8]) -> std::io::Result<usize> {
std::thread::park(); Ok(0)
}
}
let started = std::time::Instant::now();
let result =
read_stdin_bounded_with_deadline(Silent, STDIN_BYTE_LIMIT, Duration::from_millis(80));
let elapsed = started.elapsed();
assert!(
matches!(result, Err(StdinReadError::Idle)),
"expected Idle, got {result:?}"
);
assert!(elapsed >= Duration::from_millis(70));
assert!(elapsed < Duration::from_secs(2));
}
#[test]
fn slow_but_steady_reader_is_not_killed() {
struct TwoSlowChunks(usize);
impl Read for TwoSlowChunks {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
self.0 += 1;
if self.0 > 2 {
return Ok(0);
}
std::thread::sleep(Duration::from_millis(30));
buf[0] = b'k';
Ok(1)
}
}
let result = read_stdin_bounded_with_deadline(
TwoSlowChunks(0),
STDIN_BYTE_LIMIT,
Duration::from_millis(200),
);
assert_eq!(result.unwrap(), "kk");
}
#[test]
fn deadlined_path_refuses_over_limit() {
let over = Cursor::new(vec![b'x'; STDIN_BYTE_LIMIT + 5]);
let err = read_stdin_bounded_with_deadline(over, STDIN_BYTE_LIMIT, Duration::from_secs(5))
.expect_err("over-limit must be refused");
assert!(matches!(err, StdinReadError::OverLimit { .. }));
let rendered = err.to_string();
assert!(
rendered.contains(&STDIN_BYTE_LIMIT.to_string()),
"error must name the limit: {rendered}"
);
}
}