use std::io;
pub trait RegionSink {
fn data(&mut self, offset: u64, bytes: &[u8]) -> io::Result<()>;
fn zeros(&mut self, offset: u64, len: u64) -> io::Result<()>;
}
impl<S: RegionSink + ?Sized> RegionSink for &mut S {
fn data(&mut self, offset: u64, bytes: &[u8]) -> io::Result<()> {
(**self).data(offset, bytes)
}
fn zeros(&mut self, offset: u64, len: u64) -> io::Result<()> {
(**self).zeros(offset, len)
}
}
#[derive(Debug, Default)]
pub struct VecSink {
pub buf: Vec<u8>,
}
impl RegionSink for VecSink {
fn data(&mut self, offset: u64, bytes: &[u8]) -> io::Result<()> {
let end = offset as usize + bytes.len();
if self.buf.len() < end {
self.buf.resize(end, 0);
}
self.buf[offset as usize..end].copy_from_slice(bytes);
Ok(())
}
fn zeros(&mut self, offset: u64, len: u64) -> io::Result<()> {
let end = (offset + len) as usize;
if self.buf.len() < end {
self.buf.resize(end, 0);
}
self.buf[offset as usize..end].fill(0);
Ok(())
}
}
#[cfg(unix)]
#[derive(Debug)]
pub struct FileSink {
file: std::fs::File,
len: u64,
}
#[cfg(unix)]
impl FileSink {
pub fn create(path: &std::path::Path, image_len: u64) -> io::Result<FileSink> {
let file = std::fs::File::create(path)?;
file.set_len(image_len)?;
Ok(FileSink {
file,
len: image_len,
})
}
pub fn into_file(self) -> std::fs::File {
self.file
}
}
#[cfg(unix)]
impl RegionSink for FileSink {
fn data(&mut self, offset: u64, bytes: &[u8]) -> io::Result<()> {
std::os::unix::fs::FileExt::write_all_at(&self.file, bytes, offset)
}
fn zeros(&mut self, offset: u64, len: u64) -> io::Result<()> {
if offset + len > self.len {
return Err(io::Error::other("zeros past image end"));
}
Ok(())
}
}
#[derive(Debug)]
pub struct CheckingSink<S> {
inner: S,
covered: Vec<(u64, u64)>,
}
impl<S: RegionSink> CheckingSink<S> {
pub fn new(inner: S) -> Self {
CheckingSink {
inner,
covered: Vec::new(),
}
}
fn record(&mut self, offset: u64, len: u64) -> io::Result<()> {
if len == 0 {
return Ok(());
}
let end = offset + len;
let i = self.covered.partition_point(|&(s, _)| s < offset);
if i > 0 {
let (ps, pe) = self.covered[i - 1];
if pe > offset {
return Err(io::Error::other(format!(
"overlap: [{offset:#x},{end:#x}) with [{ps:#x},{pe:#x})"
)));
}
}
if i < self.covered.len() {
let (ns, ne) = self.covered[i];
if ns < end {
return Err(io::Error::other(format!(
"overlap: [{offset:#x},{end:#x}) with [{ns:#x},{ne:#x})"
)));
}
}
let merge_prev = i > 0 && self.covered[i - 1].1 == offset;
let merge_next = i < self.covered.len() && self.covered[i].0 == end;
match (merge_prev, merge_next) {
(true, true) => {
self.covered[i - 1].1 = self.covered[i].1;
self.covered.remove(i);
}
(true, false) => self.covered[i - 1].1 = end,
(false, true) => self.covered[i].0 = offset,
(false, false) => self.covered.insert(i, (offset, end)),
}
Ok(())
}
pub fn finish(self, image_len: u64) -> io::Result<S> {
if self.covered != [(0, image_len)] {
return Err(io::Error::other(format!(
"coverage is not exactly [0, {image_len:#x}): {:x?}",
&self.covered[..self.covered.len().min(8)]
)));
}
Ok(self.inner)
}
pub fn inner(&self) -> &S {
&self.inner
}
}
impl<S: RegionSink> RegionSink for CheckingSink<S> {
fn data(&mut self, offset: u64, bytes: &[u8]) -> io::Result<()> {
self.record(offset, bytes.len() as u64)?;
self.inner.data(offset, bytes)
}
fn zeros(&mut self, offset: u64, len: u64) -> io::Result<()> {
self.record(offset, len)?;
self.inner.zeros(offset, len)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn checking_sink_accepts_exact_tiling() {
let mut s = CheckingSink::new(VecSink::default());
s.data(0, &[1, 2, 3]).unwrap();
s.zeros(3, 5).unwrap();
s.data(8, &[9]).unwrap();
let inner = s.finish(9).unwrap();
assert_eq!(inner.buf, [1, 2, 3, 0, 0, 0, 0, 0, 9]);
}
#[test]
fn checking_sink_rejects_overlap() {
let mut s = CheckingSink::new(VecSink::default());
s.data(0, &[0; 8]).unwrap();
assert!(s.zeros(4, 8).is_err());
}
#[test]
fn checking_sink_rejects_gap() {
let mut s = CheckingSink::new(VecSink::default());
s.data(0, &[0; 4]).unwrap();
s.data(8, &[0; 4]).unwrap();
assert!(s.finish(12).is_err());
}
#[test]
fn checking_sink_out_of_order_ok() {
let mut s = CheckingSink::new(VecSink::default());
s.data(8, &[0; 4]).unwrap();
s.data(0, &[0; 8]).unwrap();
s.finish(12).unwrap();
}
}