use std::io;
use std::io::Read;
use std::path::Path;
pub(crate) const MAX_FILE_BYTES: u64 = 16 * 1024 * 1024;
pub(crate) fn read_to_string_capped(path: impl AsRef<Path>) -> io::Result<String> {
read_to_string_capped_with(path.as_ref(), MAX_FILE_BYTES)
}
fn read_to_string_capped_with(path: &Path, max: u64) -> io::Result<String> {
let file = std::fs::File::open(path)?;
read_capped_from(file, max, &path.display().to_string())
}
pub(crate) fn read_stdin_to_string_capped() -> io::Result<String> {
read_capped_from(io::stdin().lock(), MAX_FILE_BYTES, "standard input")
}
fn read_capped_from(reader: impl Read, max: u64, label: &str) -> io::Result<String> {
let mut buf = String::new();
let read = reader.take(max + 1).read_to_string(&mut buf)?;
if read as u64 > max {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("{label} is larger than the {max} byte limit"),
));
}
Ok(buf)
}
pub(crate) fn read_capped(path: impl AsRef<Path>) -> io::Result<Vec<u8>> {
read_capped_with(path.as_ref(), MAX_FILE_BYTES)
}
fn read_capped_with(path: &Path, max: u64) -> io::Result<Vec<u8>> {
let file = std::fs::File::open(path)?;
let mut buf = Vec::new();
let read = file.take(max + 1).read_to_end(&mut buf)?;
if read as u64 > max {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("{} is larger than the {max} byte limit", path.display()),
));
}
Ok(buf)
}
#[cfg(test)]
mod tests {
use super::{read_capped_from, read_capped_with, read_to_string_capped_with};
#[test]
fn read_capped_from_reads_within_limit() {
let out = read_capped_from(std::io::Cursor::new(b"version: 1"), 64, "stdin").unwrap();
assert_eq!(out, "version: 1");
}
#[test]
fn read_capped_from_rejects_over_limit() {
let err = read_capped_from(std::io::Cursor::new(vec![b'x'; 32]), 16, "stdin").unwrap_err();
assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
assert!(err.to_string().contains("stdin"));
}
#[test]
fn reads_file_within_limit() {
let dir = tempfile::tempdir().expect("tempdir");
let f = dir.path().join("ok");
std::fs::write(&f, b"hello").expect("write");
assert_eq!(read_to_string_capped_with(&f, 16).unwrap(), "hello");
}
#[test]
fn rejects_file_over_limit() {
let dir = tempfile::tempdir().expect("tempdir");
let f = dir.path().join("big");
std::fs::write(&f, vec![b'x'; 32]).expect("write");
let err = read_to_string_capped_with(&f, 16).unwrap_err();
assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
}
#[test]
fn missing_file_is_error() {
assert!(read_to_string_capped_with(std::path::Path::new("/no/such/file"), 16).is_err());
}
#[test]
fn read_capped_reads_bytes_within_limit() {
let dir = tempfile::tempdir().expect("tempdir");
let f = dir.path().join("ok");
std::fs::write(&f, [0xff, 0x00, 0xfe]).expect("write");
assert_eq!(read_capped_with(&f, 16).unwrap(), vec![0xff, 0x00, 0xfe]);
}
#[test]
fn read_capped_rejects_bytes_over_limit() {
let dir = tempfile::tempdir().expect("tempdir");
let f = dir.path().join("big");
std::fs::write(&f, vec![b'x'; 32]).expect("write");
let err = read_capped_with(&f, 16).unwrap_err();
assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
}
#[test]
fn read_capped_missing_file_is_error() {
assert!(read_capped_with(std::path::Path::new("/no/such/file"), 16).is_err());
}
}