326 lines
9.7 KiB
Rust
326 lines
9.7 KiB
Rust
use crate::{Reader, Writer};
|
|
use pin_project_lite::pin_project;
|
|
use rustfs_utils::{HashAlgorithm, read_full, write_all};
|
|
use tokio::io::{AsyncRead, AsyncReadExt};
|
|
|
|
pin_project! {
|
|
/// BitrotReader reads (hash+data) blocks from an async reader and verifies hash integrity.
|
|
pub struct BitrotReader {
|
|
#[pin]
|
|
inner: Box<dyn Reader>,
|
|
hash_algo: HashAlgorithm,
|
|
shard_size: usize,
|
|
buf: Vec<u8>,
|
|
hash_buf: Vec<u8>,
|
|
hash_read: usize,
|
|
data_buf: Vec<u8>,
|
|
data_read: usize,
|
|
hash_checked: bool,
|
|
}
|
|
}
|
|
|
|
impl BitrotReader {
|
|
/// Get a reference to the underlying reader.
|
|
pub fn get_ref(&self) -> &dyn Reader {
|
|
&*self.inner
|
|
}
|
|
|
|
/// Create a new BitrotReader.
|
|
pub fn new(inner: Box<dyn Reader>, shard_size: usize, algo: HashAlgorithm) -> Self {
|
|
let hash_size = algo.size();
|
|
Self {
|
|
inner,
|
|
hash_algo: algo,
|
|
shard_size,
|
|
buf: Vec::new(),
|
|
hash_buf: vec![0u8; hash_size],
|
|
hash_read: 0,
|
|
data_buf: Vec::new(),
|
|
data_read: 0,
|
|
hash_checked: false,
|
|
}
|
|
}
|
|
|
|
/// Read a single (hash+data) block, verify hash, and return the number of bytes read into `out`.
|
|
/// Returns an error if hash verification fails or data exceeds shard_size.
|
|
pub async fn read(&mut self, out: &mut [u8]) -> std::io::Result<usize> {
|
|
if out.len() > self.shard_size {
|
|
return Err(std::io::Error::new(
|
|
std::io::ErrorKind::InvalidInput,
|
|
format!("data size {} exceeds shard size {}", out.len(), self.shard_size),
|
|
));
|
|
}
|
|
|
|
let hash_size = self.hash_algo.size();
|
|
// Read hash
|
|
let mut hash_buf = vec![0u8; hash_size];
|
|
if hash_size > 0 {
|
|
self.inner.read_exact(&mut hash_buf).await?;
|
|
}
|
|
|
|
let data_len = read_full(&mut self.inner, out).await?;
|
|
|
|
// // Read data
|
|
// let mut data_len = 0;
|
|
// while data_len < out.len() {
|
|
// let n = self.inner.read(&mut out[data_len..]).await?;
|
|
// if n == 0 {
|
|
// break;
|
|
// }
|
|
// data_len += n;
|
|
// // Only read up to one shard_size block
|
|
// if data_len >= self.shard_size {
|
|
// break;
|
|
// }
|
|
// }
|
|
|
|
if hash_size > 0 {
|
|
let actual_hash = self.hash_algo.hash_encode(&out[..data_len]);
|
|
if actual_hash != hash_buf {
|
|
return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, "bitrot hash mismatch"));
|
|
}
|
|
}
|
|
Ok(data_len)
|
|
}
|
|
}
|
|
|
|
pin_project! {
|
|
/// BitrotWriter writes (hash+data) blocks to an async writer.
|
|
pub struct BitrotWriter {
|
|
#[pin]
|
|
inner: Writer,
|
|
hash_algo: HashAlgorithm,
|
|
shard_size: usize,
|
|
buf: Vec<u8>,
|
|
finished: bool,
|
|
}
|
|
}
|
|
|
|
impl BitrotWriter {
|
|
/// Create a new BitrotWriter.
|
|
pub fn new(inner: Writer, shard_size: usize, algo: HashAlgorithm) -> Self {
|
|
let hash_algo = algo;
|
|
Self {
|
|
inner,
|
|
hash_algo,
|
|
shard_size,
|
|
buf: Vec::new(),
|
|
finished: false,
|
|
}
|
|
}
|
|
|
|
pub fn into_inner(self) -> Writer {
|
|
self.inner
|
|
}
|
|
|
|
/// Write a (hash+data) block. Returns the number of data bytes written.
|
|
/// Returns an error if called after a short write or if data exceeds shard_size.
|
|
pub async fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
|
|
if buf.is_empty() {
|
|
return Ok(0);
|
|
}
|
|
|
|
if self.finished {
|
|
return Err(std::io::Error::new(std::io::ErrorKind::InvalidInput, "bitrot writer already finished"));
|
|
}
|
|
|
|
if buf.len() > self.shard_size {
|
|
return Err(std::io::Error::new(
|
|
std::io::ErrorKind::InvalidInput,
|
|
format!("data size {} exceeds shard size {}", buf.len(), self.shard_size),
|
|
));
|
|
}
|
|
|
|
if buf.len() < self.shard_size {
|
|
self.finished = true;
|
|
}
|
|
|
|
let hash_algo = &self.hash_algo;
|
|
|
|
if hash_algo.size() > 0 {
|
|
let hash = hash_algo.hash_encode(buf);
|
|
self.buf.extend_from_slice(&hash);
|
|
}
|
|
|
|
self.buf.extend_from_slice(buf);
|
|
|
|
// Write hash+data in one call
|
|
let mut n = write_all(&mut self.inner, &self.buf).await?;
|
|
|
|
if n < hash_algo.size() {
|
|
return Err(std::io::Error::new(
|
|
std::io::ErrorKind::WriteZero,
|
|
"short write: not enough bytes written",
|
|
));
|
|
}
|
|
|
|
n -= hash_algo.size();
|
|
|
|
self.buf.clear();
|
|
|
|
Ok(n)
|
|
}
|
|
}
|
|
|
|
pub fn bitrot_shard_file_size(size: usize, shard_size: usize, algo: HashAlgorithm) -> usize {
|
|
if algo != HashAlgorithm::HighwayHash256S {
|
|
return size;
|
|
}
|
|
size.div_ceil(shard_size) * algo.size() + size
|
|
}
|
|
|
|
pub async fn bitrot_verify<R: AsyncRead + Unpin + Send>(
|
|
mut r: R,
|
|
want_size: usize,
|
|
part_size: usize,
|
|
algo: HashAlgorithm,
|
|
_want: Vec<u8>,
|
|
mut shard_size: usize,
|
|
) -> std::io::Result<()> {
|
|
let mut hash_buf = vec![0; algo.size()];
|
|
let mut left = want_size;
|
|
|
|
if left != bitrot_shard_file_size(part_size, shard_size, algo.clone()) {
|
|
return Err(std::io::Error::other("bitrot shard file size mismatch"));
|
|
}
|
|
|
|
while left > 0 {
|
|
let n = r.read_exact(&mut hash_buf).await?;
|
|
left -= n;
|
|
|
|
if left < shard_size {
|
|
shard_size = left;
|
|
}
|
|
|
|
let mut buf = vec![0; shard_size];
|
|
let read = r.read_exact(&mut buf).await?;
|
|
|
|
let actual_hash = algo.hash_encode(&buf);
|
|
if actual_hash != hash_buf[0..n] {
|
|
return Err(std::io::Error::other("bitrot hash mismatch"));
|
|
}
|
|
|
|
left -= read;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
|
|
use crate::{BitrotReader, BitrotWriter, Writer};
|
|
use rustfs_utils::HashAlgorithm;
|
|
use std::io::Cursor;
|
|
|
|
#[tokio::test]
|
|
async fn test_bitrot_read_write_ok() {
|
|
let data = b"hello world! this is a test shard.";
|
|
let data_size = data.len();
|
|
let shard_size = 8;
|
|
|
|
let buf = Vec::new();
|
|
let writer = Cursor::new(buf);
|
|
let mut bitrot_writer = BitrotWriter::new(Writer::from_cursor(writer), shard_size, HashAlgorithm::HighwayHash256);
|
|
|
|
let mut n = 0;
|
|
for chunk in data.chunks(shard_size) {
|
|
n += bitrot_writer.write(chunk).await.unwrap();
|
|
}
|
|
assert_eq!(n, data.len());
|
|
|
|
// 读
|
|
let reader = Cursor::new(bitrot_writer.into_inner().into_cursor_inner().unwrap());
|
|
let reader = Box::new(reader);
|
|
let mut bitrot_reader = BitrotReader::new(reader, shard_size, HashAlgorithm::HighwayHash256);
|
|
let mut out = Vec::new();
|
|
let mut n = 0;
|
|
while n < data_size {
|
|
let mut buf = vec![0u8; shard_size];
|
|
let m = bitrot_reader.read(&mut buf).await.unwrap();
|
|
assert_eq!(&buf[..m], &data[n..n + m]);
|
|
|
|
out.extend_from_slice(&buf[..m]);
|
|
n += m;
|
|
}
|
|
|
|
assert_eq!(n, data_size);
|
|
assert_eq!(data, &out[..]);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_bitrot_read_hash_mismatch() {
|
|
let data = b"test data for bitrot";
|
|
let data_size = data.len();
|
|
let shard_size = 8;
|
|
let buf = Vec::new();
|
|
let writer = Cursor::new(buf);
|
|
let mut bitrot_writer = BitrotWriter::new(Writer::from_cursor(writer), shard_size, HashAlgorithm::HighwayHash256);
|
|
for chunk in data.chunks(shard_size) {
|
|
let _ = bitrot_writer.write(chunk).await.unwrap();
|
|
}
|
|
let mut written = bitrot_writer.into_inner().into_cursor_inner().unwrap();
|
|
// change the last byte to make hash mismatch
|
|
let pos = written.len() - 1;
|
|
written[pos] ^= 0xFF;
|
|
let reader = Cursor::new(written);
|
|
let reader = Box::new(reader);
|
|
let mut bitrot_reader = BitrotReader::new(reader, shard_size, HashAlgorithm::HighwayHash256);
|
|
|
|
let count = data_size.div_ceil(shard_size);
|
|
|
|
let mut idx = 0;
|
|
let mut n = 0;
|
|
while n < data_size {
|
|
let mut buf = vec![0u8; shard_size];
|
|
let res = bitrot_reader.read(&mut buf).await;
|
|
|
|
if idx == count - 1 {
|
|
// 最后一个块，应该返回错误
|
|
assert!(res.is_err());
|
|
assert_eq!(res.unwrap_err().kind(), std::io::ErrorKind::InvalidData);
|
|
break;
|
|
}
|
|
|
|
let m = res.unwrap();
|
|
|
|
assert_eq!(&buf[..m], &data[n..n + m]);
|
|
|
|
n += m;
|
|
idx += 1;
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_bitrot_read_write_none_hash() {
|
|
let data = b"bitrot none hash test data!";
|
|
let data_size = data.len();
|
|
let shard_size = 8;
|
|
|
|
let buf = Vec::new();
|
|
let writer = Cursor::new(buf);
|
|
let mut bitrot_writer = BitrotWriter::new(Writer::from_cursor(writer), shard_size, HashAlgorithm::None);
|
|
|
|
let mut n = 0;
|
|
for chunk in data.chunks(shard_size) {
|
|
n += bitrot_writer.write(chunk).await.unwrap();
|
|
}
|
|
assert_eq!(n, data.len());
|
|
|
|
let reader = Cursor::new(bitrot_writer.into_inner().into_cursor_inner().unwrap());
|
|
let reader = Box::new(reader);
|
|
let mut bitrot_reader = BitrotReader::new(reader, shard_size, HashAlgorithm::None);
|
|
let mut out = Vec::new();
|
|
let mut n = 0;
|
|
while n < data_size {
|
|
let mut buf = vec![0u8; shard_size];
|
|
let m = bitrot_reader.read(&mut buf).await.unwrap();
|
|
assert_eq!(&buf[..m], &data[n..n + m]);
|
|
out.extend_from_slice(&buf[..m]);
|
|
n += m;
|
|
}
|
|
assert_eq!(n, data_size);
|
|
assert_eq!(data, &out[..]);
|
|
}
|
|
}
|