use std::io::{
BufRead,
Error,
ErrorKind,
Read,
Result,
};
use std::string::FromUtf8Error;
use crate::util::try_reserve_vec;
pub const READ_TO_END_BUFFER_SIZE: usize = 8 * 1024;
pub fn read_exact_or_eof(
reader: &mut dyn Read,
buffer: &mut [u8],
) -> Result<usize> {
let mut total = 0;
while total < buffer.len() {
match reader.read(&mut buffer[total..]) {
Ok(0) => break,
Ok(count) => total += count,
Err(error) => {
if error.kind() == ErrorKind::Interrupted {
continue;
}
return Err(error);
}
}
}
Ok(total)
}
pub fn read_exact_vec_limited_into(
reader: &mut dyn Read,
output: &mut Vec<u8>,
len: usize,
max_len: usize,
) -> Result<()> {
validate_exact_read_len(len, max_len)?;
let original_len = output.len();
let new_len = match original_len.checked_add(len) {
Some(value) => value,
None => {
return Err(Error::new(
ErrorKind::InvalidInput,
format!("length {original_len} plus {len} overflows usize"),
));
}
};
try_reserve_vec(output, len)?;
output.resize(new_len, 0);
match reader.read_exact(&mut output[original_len..]) {
Ok(()) => Ok(()),
Err(error) => {
output.truncate(original_len);
Err(error)
}
}
}
#[inline(always)]
pub fn validate_exact_read_len(len: usize, max_len: usize) -> Result<()> {
if len > max_len {
return Err(Error::new(
ErrorKind::InvalidData,
format!("requested length {len} exceeds maximum length {max_len}"),
));
}
Ok(())
}
pub fn read_to_end_limited(
reader: &mut dyn Read,
max_len: usize,
) -> Result<Vec<u8>> {
let mut output = Vec::new();
try_reserve_vec(&mut output, max_len.min(READ_TO_END_BUFFER_SIZE))?;
read_to_end_limited_into(reader, &mut output, max_len)?;
Ok(output)
}
pub fn read_to_end_limited_into(
reader: &mut dyn Read,
output: &mut Vec<u8>,
max_len: usize,
) -> Result<usize> {
let original_len = output.len();
let mut buffer = [0; READ_TO_END_BUFFER_SIZE];
let mut appended = 0;
loop {
let remaining = max_len.saturating_sub(appended);
let requested =
remaining.saturating_add(1).min(READ_TO_END_BUFFER_SIZE);
match reader.read(&mut buffer[..requested]) {
Ok(count) => {
if count == 0 {
return Ok(appended);
} else if count <= remaining {
if let Err(error) = try_reserve_vec(output, count) {
output.truncate(original_len);
return Err(error);
}
output.extend_from_slice(&buffer[..count]);
appended += count;
} else {
output.truncate(original_len);
return Err(Error::new(
ErrorKind::InvalidData,
format!(
"input exceeds maximum length of {max_len} bytes"
),
));
}
}
Err(error) => {
if error.kind() == ErrorKind::Interrupted {
continue;
}
output.truncate(original_len);
return Err(error);
}
}
}
}
pub fn invalid_utf8_error(error: FromUtf8Error) -> Error {
Error::new(
ErrorKind::InvalidData,
format!("limited input is not valid UTF-8: {error}"),
)
}
pub fn read_until_limited_into<T>(
reader: &mut T,
delimiter: u8,
output: &mut Vec<u8>,
max_len: usize,
) -> Result<usize>
where
T: BufRead + ?Sized,
{
let original_len = output.len();
match read_until_limited_into_inner(reader, delimiter, output, max_len) {
Ok(count) => Ok(count),
Err(error) => {
output.truncate(original_len);
Err(error)
}
}
}
fn read_until_limited_into_inner<T>(
reader: &mut T,
delimiter: u8,
output: &mut Vec<u8>,
max_len: usize,
) -> Result<usize>
where
T: BufRead + ?Sized,
{
let mut appended = 0;
loop {
let available = reader.fill_buf()?;
if available.is_empty() {
return Ok(appended);
}
let delimiter_position =
available.iter().position(|byte| *byte == delimiter);
let requested =
delimiter_position.map_or(available.len(), |position| position + 1);
let remaining = max_len.saturating_sub(appended);
if requested > remaining {
if remaining > 0 {
reader.consume(remaining);
}
return Err(limit_exceeded_error(max_len, delimiter));
}
try_reserve_vec(output, requested)?;
output.extend_from_slice(&available[..requested]);
reader.consume(requested);
appended += requested;
if delimiter_position.is_some() {
return Ok(appended);
}
}
}
fn limit_exceeded_error(max_len: usize, delimiter: u8) -> Error {
Error::new(
ErrorKind::InvalidData,
format!(
"input exceeds maximum length of {max_len} bytes before delimiter {delimiter}"
),
)
}