use crate::{PixelsError, Result};
pub trait Source: Send {
fn read(&mut self, buf: &mut [u8]) -> Result<usize>;
fn read_exact(&mut self, buf: &mut [u8]) -> Result<()> {
let wanted = buf.len();
let mut filled = 0;
while filled < wanted {
let Some(rest) = buf.get_mut(filled..) else {
break;
};
match self.read(rest)? {
0 => {
return Err(PixelsError::malformed(
"stream",
format!("stream ended after {filled} of {wanted} expected bytes"),
));
}
n => filled += n,
}
}
Ok(())
}
}
pub trait Sink: Send {
fn write_all(&mut self, buf: &[u8]) -> Result<()>;
fn flush(&mut self) -> Result<()>;
}
impl<R: std::io::Read + Send> Source for R {
fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
loop {
return match std::io::Read::read(self, buf) {
Ok(n) => Ok(n),
Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
Err(e) => Err(PixelsError::io("reading from source", e)),
};
}
}
}
impl<W: std::io::Write + Send> Sink for W {
fn write_all(&mut self, buf: &[u8]) -> Result<()> {
std::io::Write::write_all(self, buf).map_err(|e| PixelsError::io("writing to sink", e))
}
fn flush(&mut self) -> Result<()> {
std::io::Write::flush(self).map_err(|e| PixelsError::io("flushing sink", e))
}
}
#[derive(Debug)]
pub struct Prefixed<S: Source> {
prefix: Vec<u8>,
consumed: usize,
rest: S,
}
impl<S: Source> Prefixed<S> {
#[must_use]
pub const fn new(prefix: Vec<u8>, rest: S) -> Self {
Self {
prefix,
consumed: 0,
rest,
}
}
}
impl<S: Source> Source for Prefixed<S> {
fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
let Some(remaining) = self.prefix.get(self.consumed..) else {
return self.rest.read(buf);
};
if remaining.is_empty() {
return self.rest.read(buf);
}
let take = remaining.len().min(buf.len());
let (Some(from), Some(into)) = (remaining.get(..take), buf.get_mut(..take)) else {
return Ok(0);
};
into.copy_from_slice(from);
self.consumed += take;
Ok(take)
}
}
#[cfg(test)]
#[allow(
clippy::unwrap_used,
clippy::indexing_slicing,
reason = "tests operate on known-good values and assert shapes directly"
)]
mod tests {
use super::*;
use crate::ErrorCode;
struct Trickle {
data: Vec<u8>,
pos: usize,
chunk: usize,
}
impl Source for Trickle {
fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
let remaining = self.data.len() - self.pos;
let n = remaining.min(buf.len()).min(self.chunk);
buf[..n].copy_from_slice(&self.data[self.pos..self.pos + n]);
self.pos += n;
Ok(n)
}
}
#[test]
fn read_exact_reassembles_short_reads() {
let mut src = Trickle {
data: (0..10).collect(),
pos: 0,
chunk: 3,
};
let mut buf = [0_u8; 10];
src.read_exact(&mut buf).unwrap();
assert_eq!(buf, [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
}
#[test]
fn a_truncated_stream_is_malformed_not_an_io_error() {
let mut src = Trickle {
data: vec![1, 2, 3],
pos: 0,
chunk: 2,
};
let mut buf = [0_u8; 8];
let err = src.read_exact(&mut buf).unwrap_err();
assert_eq!(err.code(), ErrorCode::Malformed);
assert!(err.to_string().contains("3 of 8"), "{err}");
}
#[test]
fn slices_and_vecs_work_without_adapters() {
let mut src: &[u8] = b"hello";
let mut buf = [0_u8; 5];
src.read_exact(&mut buf).unwrap();
assert_eq!(&buf, b"hello");
let mut sink: Vec<u8> = Vec::new();
sink.write_all(b"out").unwrap();
sink.flush().unwrap();
assert_eq!(sink, b"out");
}
#[test]
fn read_exact_of_nothing_succeeds() {
let mut src: &[u8] = b"";
src.read_exact(&mut []).unwrap();
}
#[test]
fn sink_errors_are_reported_as_io() {
struct Broken;
impl std::io::Write for Broken {
fn write(&mut self, _: &[u8]) -> std::io::Result<usize> {
Err(std::io::Error::other("disk full"))
}
fn flush(&mut self) -> std::io::Result<()> {
Err(std::io::Error::other("disk full"))
}
}
let err = Sink::write_all(&mut Broken, b"x").unwrap_err();
assert_eq!(err.code(), ErrorCode::Io);
assert_eq!(Sink::flush(&mut Broken).unwrap_err().code(), ErrorCode::Io);
}
#[test]
fn a_prefixed_source_replays_then_delegates() {
let mut source = Prefixed::new(vec![1, 2, 3], &b"456"[..]);
let mut all = Vec::new();
let mut buf = [0_u8; 2];
loop {
match source.read(&mut buf).unwrap() {
0 => break,
n => all.extend_from_slice(&buf[..n]),
}
}
assert_eq!(all, vec![1, 2, 3, b'4', b'5', b'6']);
}
#[test]
fn a_prefixed_source_never_mixes_the_prefix_into_one_read() {
let mut source = Prefixed::new(vec![9, 9], &b"xyz"[..]);
let mut buf = [0_u8; 16];
assert_eq!(source.read(&mut buf).unwrap(), 2);
assert_eq!(&buf[..2], &[9, 9]);
assert_eq!(source.read(&mut buf).unwrap(), 3);
assert_eq!(&buf[..3], b"xyz");
assert_eq!(source.read(&mut buf).unwrap(), 0);
}
#[test]
fn an_empty_prefix_is_a_pass_through() {
let mut source = Prefixed::new(Vec::new(), &b"data"[..]);
let mut buf = [0_u8; 8];
assert_eq!(source.read(&mut buf).unwrap(), 4);
assert_eq!(&buf[..4], b"data");
}
#[test]
fn read_exact_spans_the_prefix_boundary() {
let mut source = Prefixed::new(vec![0xDE, 0xAD], &b"\xBE\xEF"[..]);
let mut buf = [0_u8; 4];
source.read_exact(&mut buf).unwrap();
assert_eq!(buf, [0xDE, 0xAD, 0xBE, 0xEF]);
}
}