use std::io;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, ReadBuf};
pub trait IntegrityCheck: Send {
fn update(&mut self, chunk: &[u8]);
fn finalize(self: Box<Self>, total: u64) -> Result<(), String>;
}
#[derive(Debug, Clone, Copy)]
pub struct LengthCheck {
expected: u64,
}
impl LengthCheck {
pub fn new(expected: u64) -> Self {
Self { expected }
}
}
impl IntegrityCheck for LengthCheck {
fn update(&mut self, _chunk: &[u8]) {}
fn finalize(self: Box<Self>, total: u64) -> Result<(), String> {
if total == self.expected {
Ok(())
} else {
Err(format!(
"body length mismatch: object store advertised {} byte(s) but read {} \
(truncated or corrupted transfer)",
self.expected, total
))
}
}
}
pub struct VerifyingReader<R> {
inner: R,
checks: Vec<Box<dyn IntegrityCheck>>,
total: u64,
finalized: bool,
}
impl<R> VerifyingReader<R> {
pub fn new(inner: R, checks: Vec<Box<dyn IntegrityCheck>>) -> Self {
Self {
inner,
checks,
total: 0,
finalized: false,
}
}
}
impl<R: AsyncRead + Unpin> AsyncRead for VerifyingReader<R> {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
let this = self.get_mut();
let before = buf.filled().len();
match Pin::new(&mut this.inner).poll_read(cx, buf) {
Poll::Ready(Ok(())) => {
let filled = buf.filled();
let n = filled.len() - before;
if n > 0 {
let chunk = &filled[before..];
this.total += n as u64;
for check in &mut this.checks {
check.update(chunk);
}
Poll::Ready(Ok(()))
} else if this.finalized {
Poll::Ready(Ok(()))
} else {
this.finalized = true;
for check in std::mem::take(&mut this.checks) {
if let Err(reason) = check.finalize(this.total) {
return Poll::Ready(Err(io::Error::new(
io::ErrorKind::InvalidData,
reason,
)));
}
}
Poll::Ready(Ok(()))
}
}
other => other,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use tokio::io::AsyncReadExt as _;
struct SumCheck {
expected: u64,
running: u64,
}
impl SumCheck {
fn new(expected: u64) -> Self {
Self {
expected,
running: 0,
}
}
}
impl IntegrityCheck for SumCheck {
fn update(&mut self, chunk: &[u8]) {
self.running += chunk.iter().map(|b| *b as u64).sum::<u64>();
}
fn finalize(self: Box<Self>, _total: u64) -> Result<(), String> {
if self.running == self.expected {
Ok(())
} else {
Err(format!(
"checksum mismatch: expected {}, computed {}",
self.expected, self.running
))
}
}
}
struct OneByteAtATime {
data: Vec<u8>,
pos: usize,
}
impl AsyncRead for OneByteAtATime {
fn poll_read(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
let this = self.get_mut();
if this.pos < this.data.len() {
buf.put_slice(&this.data[this.pos..this.pos + 1]);
this.pos += 1;
}
Poll::Ready(Ok(()))
}
}
#[tokio::test]
async fn digest_mismatch_fails() {
let data: &[u8] = &[1, 2, 3]; let mut reader = VerifyingReader::new(data, vec![Box::new(SumCheck::new(99))]);
let mut out = Vec::new();
let err = reader
.read_to_end(&mut out)
.await
.expect_err("wrong checksum must fail the read");
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
assert!(err.to_string().contains("checksum mismatch"), "got: {err}");
}
#[tokio::test]
async fn digest_match_passes() {
let data: &[u8] = &[1, 2, 3]; let mut reader = VerifyingReader::new(data, vec![Box::new(SumCheck::new(6))]);
let mut out = Vec::new();
reader
.read_to_end(&mut out)
.await
.expect("matching checksum must pass");
assert_eq!(out, vec![1, 2, 3]);
}
#[tokio::test]
async fn accumulates_across_one_byte_reads() {
let body = vec![10u8, 20, 30, 40]; let reader = OneByteAtATime {
data: body.clone(),
pos: 0,
};
let mut reader = VerifyingReader::new(
reader,
vec![Box::new(LengthCheck::new(4)), Box::new(SumCheck::new(100))],
);
let mut out = Vec::new();
reader
.read_to_end(&mut out)
.await
.expect("incremental reads must verify");
assert_eq!(out, body);
}
#[tokio::test]
async fn first_failing_check_surfaces() {
let data: &[u8] = &[1, 2, 3];
let mut reader = VerifyingReader::new(
data,
vec![Box::new(LengthCheck::new(3)), Box::new(SumCheck::new(0))],
);
let mut out = Vec::new();
let err = reader.read_to_end(&mut out).await.expect_err("must fail");
assert!(err.to_string().contains("checksum mismatch"), "got: {err}");
}
#[tokio::test]
async fn empty_body_expecting_zero_passes() {
let data: &[u8] = &[];
let mut reader = VerifyingReader::new(data, vec![Box::new(LengthCheck::new(0))]);
let mut out = Vec::new();
let n = reader
.read_to_end(&mut out)
.await
.expect("empty body, len 0");
assert_eq!(n, 0);
}
#[tokio::test]
async fn empty_body_expecting_nonzero_fails() {
let data: &[u8] = &[];
let mut reader = VerifyingReader::new(data, vec![Box::new(LengthCheck::new(5))]);
let mut out = Vec::new();
let err = reader
.read_to_end(&mut out)
.await
.expect_err("an empty body advertised as 5 bytes must fail");
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
}
#[tokio::test]
async fn overlong_read_fails() {
let data: &[u8] = b"helloworld"; let mut reader = VerifyingReader::new(data, vec![Box::new(LengthCheck::new(4))]);
let mut out = Vec::new();
let err = reader
.read_to_end(&mut out)
.await
.expect_err("reading more than advertised must fail");
assert!(err.to_string().contains("length mismatch"), "got: {err}");
}
#[tokio::test]
async fn no_checks_is_transparent_passthrough() {
let data: &[u8] = b"passthrough";
let mut reader = VerifyingReader::new(data, Vec::new());
let mut out = Vec::new();
reader.read_to_end(&mut out).await.expect("no checks = ok");
assert_eq!(out, b"passthrough");
}
#[tokio::test]
async fn short_read_fails_length_check() {
let data: &[u8] = b"hello"; let mut reader = VerifyingReader::new(data, vec![Box::new(LengthCheck::new(10))]);
let mut out = Vec::new();
let err = reader
.read_to_end(&mut out)
.await
.expect_err("a 5-byte body against an advertised 10 must fail");
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
assert!(
err.to_string().contains("length mismatch"),
"unexpected error: {err}"
);
}
#[tokio::test]
async fn exact_length_passes() {
let data: &[u8] = b"helloworld"; let mut reader = VerifyingReader::new(data, vec![Box::new(LengthCheck::new(10))]);
let mut out = Vec::new();
let n = reader
.read_to_end(&mut out)
.await
.expect("exact-length body must succeed");
assert_eq!(n, 10);
assert_eq!(out, b"helloworld");
}
}