use crate::error::{Error, Result};
use crate::file::OpenedFile;
use crate::header::FixedHeader;
use crate::payload::read_range;
use crate::tlv::CompressionIndexEntry;
use std::io::{Read, Seek, SeekFrom};
use std::path::{Path, PathBuf};
pub trait BlockSource {
fn header(&self) -> &FixedHeader;
fn read_ct(&self, offset: u64, len: u64) -> Result<Vec<u8>>;
fn ct_len(&self) -> u64;
fn is_remote(&self) -> bool {
false
}
fn describe(&self) -> String {
String::from("<block source>")
}
}
#[derive(Debug)]
pub struct MemorySource {
header: FixedHeader,
payload: Vec<u8>,
}
impl MemorySource {
pub fn new(file_bytes: &[u8]) -> Result<Self> {
let header = crate::file::peek_header(file_bytes)?;
let start = usize::try_from(header.header_len)
.map_err(|_| Error::MalformedHeader { reason: "header_len exceeds usize" })?;
let payload = file_bytes
.get(start..)
.ok_or(Error::Truncated {
context: "payload area",
need: start,
got: file_bytes.len(),
})?
.to_vec();
Ok(Self { header, payload })
}
}
impl BlockSource for MemorySource {
fn header(&self) -> &FixedHeader {
&self.header
}
fn read_ct(&self, offset: u64, len: u64) -> Result<Vec<u8>> {
let start = usize::try_from(offset)
.map_err(|_| Error::ChunkOutOfRange { index: offset, total: self.ct_len() })?;
let n = usize::try_from(len)
.map_err(|_| Error::ChunkOutOfRange { index: len, total: self.ct_len() })?;
let end = start
.checked_add(n)
.ok_or(Error::ChunkOutOfRange { index: offset, total: self.ct_len() })?;
self.payload
.get(start..end)
.map(<[u8]>::to_vec)
.ok_or(Error::Truncated {
context: "memory ciphertext range",
need: end,
got: self.payload.len(),
})
}
fn ct_len(&self) -> u64 {
self.payload.len() as u64
}
fn describe(&self) -> String {
String::from("memory")
}
}
#[derive(Debug)]
pub struct LocalFileSource {
path: PathBuf,
header: FixedHeader,
payload_start: u64,
payload_len: u64,
}
impl LocalFileSource {
pub fn open(path: impl AsRef<Path>) -> Result<Self> {
let path = path.as_ref().to_path_buf();
let mut f = std::fs::File::open(&path)?;
let mut probe = vec![0u8; crate::scan::MIN_PROBE_SIZE];
let n = read_at_most(&mut f, &mut probe)?;
probe.truncate(n);
let partial = crate::file::peek_header(&probe)?;
let hlen = partial.header_len;
let total = f.metadata()?.len();
if total < u64::from(hlen) {
return Err(Error::Truncated {
context: "local file header",
need: hlen as usize,
got: usize::try_from(total).unwrap_or(usize::MAX),
});
}
let mut head = vec![0u8; hlen as usize];
f.seek(SeekFrom::Start(0))?;
f.read_exact(&mut head)?;
let header = crate::file::peek_header(&head)?;
Ok(Self {
path,
payload_start: u64::from(hlen),
payload_len: total.saturating_sub(u64::from(hlen)),
header,
})
}
#[must_use]
pub fn path(&self) -> &Path {
&self.path
}
}
impl BlockSource for LocalFileSource {
fn header(&self) -> &FixedHeader {
&self.header
}
fn read_ct(&self, offset: u64, len: u64) -> Result<Vec<u8>> {
let end = offset
.checked_add(len)
.ok_or(Error::ChunkOutOfRange { index: offset, total: self.payload_len })?;
if end > self.payload_len {
return Err(Error::Truncated {
context: "local file ciphertext range",
need: usize::try_from(end).unwrap_or(usize::MAX),
got: usize::try_from(self.payload_len).unwrap_or(usize::MAX),
});
}
let n = usize::try_from(len)
.map_err(|_| Error::ChunkOutOfRange { index: len, total: self.payload_len })?;
let mut f = std::fs::File::open(&self.path)?;
f.seek(SeekFrom::Start(
self.payload_start
.checked_add(offset)
.ok_or(Error::ChunkOutOfRange { index: offset, total: self.payload_len })?,
))?;
let mut buf = vec![0u8; n];
f.read_exact(&mut buf)?;
Ok(buf)
}
fn ct_len(&self) -> u64 {
self.payload_len
}
fn describe(&self) -> String {
self.path.display().to_string()
}
}
fn read_at_most(f: &mut std::fs::File, buf: &mut [u8]) -> Result<usize> {
let mut total = 0usize;
while total < buf.len() {
let Some(slice) = buf.get_mut(total..) else {
break;
};
match f.read(slice) {
Ok(0) => break,
Ok(n) => total = total.saturating_add(n),
Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {}
Err(e) => return Err(e.into()),
}
}
Ok(total)
}
#[derive(Debug)]
pub struct ContainerEntryView<'a, S: BlockSource> {
inner: &'a S,
base: u64,
len: u64,
}
impl<'a, S: BlockSource> ContainerEntryView<'a, S> {
#[must_use]
pub const fn new(inner: &'a S, offset: u64, size: u64) -> Self {
Self { inner, base: offset, len: size }
}
#[must_use]
pub const fn plaintext_len(&self) -> u64 {
self.len
}
pub fn read(
&self,
opened: &OpenedFile,
offset: u64,
length: u64,
) -> Result<Vec<u8>> {
if offset >= self.len {
return Ok(Vec::new());
}
let avail = self.len.saturating_sub(offset);
let want = length.min(avail);
let abs = self
.base
.checked_add(offset)
.ok_or(Error::MalformedHeader { reason: "entry offset overflow" })?;
read_source_range(self.inner, opened, abs, want)
}
}
pub fn read_source_range<S: BlockSource + ?Sized>(
src: &S,
opened: &OpenedFile,
offset: u64,
length: u64,
) -> Result<Vec<u8>> {
let index: Option<Vec<CompressionIndexEntry>> =
if src.header().has_flag(crate::header::flags::COMPRESSED) {
Some(opened.compression_index()?)
} else {
None
};
read_range(
src.header(),
opened.payload_key(),
index.as_deref(),
offset,
length,
|ct_off, ct_len| src.read_ct(ct_off, ct_len),
)
}
pub fn decrypt_to_sink<S: BlockSource + ?Sized>(
src: &S,
opened: &OpenedFile,
sink: &mut dyn FnMut(&[u8]) -> Result<()>,
mut progress: Option<crate::payload::ProgressFn<'_>>,
) -> Result<u64> {
use blake2::digest::{Digest, consts::U32};
let total = opened.header.plaintext_size;
let chunk = u64::from(opened.header.chunk_size).max(1);
let window = crate::unpack::STREAM_WINDOW
.checked_div(chunk)
.unwrap_or(1)
.max(1)
.saturating_mul(chunk);
let index: Option<Vec<CompressionIndexEntry>> =
if src.header().has_flag(crate::header::flags::COMPRESSED) {
Some(opened.compression_index()?)
} else {
None
};
let mut hasher = blake2::Blake2b::<U32>::new();
let mut done = 0u64;
while done < total {
let n = total.saturating_sub(done).min(window);
let plain = read_range(
src.header(),
opened.payload_key(),
index.as_deref(),
done,
n,
|ct_off, ct_len| src.read_ct(ct_off, ct_len),
)?;
if plain.len() as u64 != n {
return Err(Error::Truncated {
context: "payload",
need: usize::try_from(n).unwrap_or(usize::MAX),
got: plain.len(),
});
}
hasher.update(&plain);
sink(&plain)?;
done = done.saturating_add(n);
if let Some(cb) = progress.as_deref_mut() {
cb(done, total);
}
}
if let Ok(expected) = opened.stored_content_hash() {
use subtle::ConstantTimeEq;
let actual: [u8; 32] = hasher.finalize().into();
if !bool::from(actual.ct_eq(&expected)) {
return Err(Error::ContentHashMismatch);
}
}
Ok(done)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::crypto::{Argon2Params, Kek};
use crate::file::{EncryptOptions, RandomMaterial, encrypt, open};
use crate::header::MIN_CHUNK_SIZE;
fn make(payload: &[u8], compress: bool) -> (Vec<u8>, Kek) {
let salt = [0x22u8; 16];
let params = Argon2Params::TEST_WEAK;
let kek = Kek::from_password(b"src", &salt, params).unwrap();
let opts = EncryptOptions {
filename: Some("s.bin".into()),
chunk_size: MIN_CHUNK_SIZE,
argon2: params,
compress,
..EncryptOptions::default()
};
let k2 = Kek::from_password(b"src", &salt, params).unwrap();
let enc = encrypt(payload, &[kek], &salt, &opts, &RandomMaterial::generate()).unwrap();
(enc.bytes, k2)
}
#[test]
fn memory_source_reads_ranges() {
let plain: Vec<u8> = (0..200_000u32).map(|i| (i % 251) as u8).collect();
let (bytes, kek) = make(&plain, false);
let src = MemorySource::new(&bytes).unwrap();
let opened = open(&bytes, &[kek]).unwrap();
for (off, len) in [(0u64, 10u64), (65_530, 20), (199_990, 10), (0, 200_000)] {
let got = read_source_range(&src, &opened, off, len).unwrap();
let want = &plain[off as usize..(off + len) as usize];
assert_eq!(got, want, "区间 ({off},{len}) 不符");
}
assert!(!src.is_remote());
}
#[test]
fn local_file_source_matches_memory() {
let plain: Vec<u8> = (0..150_000u32).map(|i| (i % 253) as u8).collect();
let (bytes, kek) = make(&plain, false);
let dir = std::env::temp_dir().join(format!("omy_src_{}", std::process::id()));
std::fs::create_dir_all(&dir).unwrap();
let p = dir.join("a.omy");
std::fs::write(&p, &bytes).unwrap();
let src = LocalFileSource::open(&p).unwrap();
let opened = open(&bytes, &[kek]).unwrap();
assert_eq!(src.ct_len(), (bytes.len() - src.header().header_len as usize) as u64);
for (off, len) in [(0u64, 100u64), (65_536, 1000), (149_000, 1000)] {
let got = read_source_range(&src, &opened, off, len).unwrap();
assert_eq!(got, &plain[off as usize..(off + len) as usize]);
}
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn compressed_source_reads_ranges() {
let plain = vec![0x5Au8; 300_000];
let (bytes, kek) = make(&plain, true);
let src = MemorySource::new(&bytes).unwrap();
let opened = open(&bytes, &[kek]).unwrap();
assert!(src.header().has_flag(crate::header::flags::COMPRESSED));
let got = read_source_range(&src, &opened, 100, 5000).unwrap();
assert_eq!(got, &plain[100..5100]);
}
#[test]
fn container_view_clips_to_entry() {
let mut plain = vec![0xAAu8; 1000];
plain.extend_from_slice(&[0xBBu8; 1000]);
let (bytes, kek) = make(&plain, false);
let src = MemorySource::new(&bytes).unwrap();
let opened = open(&bytes, &[kek]).unwrap();
let second = ContainerEntryView::new(&src, 1000, 1000);
assert_eq!(second.plaintext_len(), 1000);
let got = second.read(&opened, 0, 1000).unwrap();
assert_eq!(got.len(), 1000);
assert!(got.iter().all(|&b| b == 0xBB));
let got2 = second.read(&opened, 900, 500).unwrap();
assert_eq!(got2.len(), 100);
assert!(got2.iter().all(|&b| b == 0xBB));
assert!(second.read(&opened, 1000, 10).unwrap().is_empty());
}
#[test]
fn out_of_bounds_ct_read_errors() {
let (bytes, _) = make(b"short", false);
let src = MemorySource::new(&bytes).unwrap();
assert!(src.read_ct(0, src.ct_len() + 1).is_err());
assert!(src.read_ct(u64::MAX, 1).is_err());
}
#[test]
fn describe_is_useful() {
let (bytes, _) = make(b"x", false);
let src = MemorySource::new(&bytes).unwrap();
assert_eq!(src.describe(), "memory");
}
#[test]
fn decrypt_to_sink_matches_decrypt_all() -> Result<()> {
let big: Vec<u8> = (0..(crate::unpack::STREAM_WINDOW as usize * 2 + 12_345))
.map(|i| (i % 251) as u8)
.collect();
for compress in [false, true] {
let (file, kek) = make(&big, compress);
let opened = open(&file, &[kek])?;
let src = MemorySource::new(&file)?;
let mut out = Vec::new();
let n = decrypt_to_sink(&src, &opened, &mut |b| { out.extend_from_slice(b); Ok(()) }, None)?;
assert_eq!(n as usize, big.len());
assert_eq!(out, opened.decrypt_all(&file)?, "compress={compress}");
}
Ok(())
}
}