use crate::crypto::{CipherId, SecretKey, TAG_LEN, chunk_aad, chunk_nonce};
use crate::error::{Error, Result};
use crate::header::{COMPRESS_ZSTD, FixedHeader, flags};
use crate::tlv::CompressionIndexEntry;
pub type ProgressFn<'a> = &'a mut dyn FnMut(u64, u64);
pub fn encrypt_chunk(
payload_key: &SecretKey,
cipher: CipherId,
base_nonce: &[u8; 7],
file_uuid: &[u8; 16],
index: u32,
is_final: bool,
plaintext: &[u8],
) -> Result<Vec<u8>> {
let nonce = chunk_nonce(base_nonce, index, is_final);
let aad = chunk_aad(file_uuid, index);
cipher.encrypt(payload_key, &nonce, plaintext, &aad)
}
pub fn decrypt_chunk(
payload_key: &SecretKey,
cipher: CipherId,
base_nonce: &[u8; 7],
file_uuid: &[u8; 16],
index: u32,
is_final: bool,
ciphertext: &[u8],
) -> Result<Vec<u8>> {
let nonce = chunk_nonce(base_nonce, index, is_final);
let aad = chunk_aad(file_uuid, index);
cipher
.decrypt(payload_key, &nonce, ciphertext, &aad)?
.ok_or(Error::ChunkAuthFailed { index: u64::from(index) })
}
fn decompress_chunk(body: &[u8], plain_len: u32, chunk_size: u32) -> Result<Vec<u8>> {
if plain_len > chunk_size {
return Err(Error::MalformedTlv {
tlv_type: crate::tlv::types::COMPRESSION_INDEX,
reason: "index plain_len exceeds chunk_size",
});
}
let out = zstd::bulk::decompress(body, plain_len as usize)
.map_err(|err| Error::Compression { reason: err.to_string() })?;
if out.len() != plain_len as usize {
return Err(Error::Compression {
reason: format!(
"decompressed length {} does not match index plain_len {plain_len}",
out.len()
),
});
}
Ok(out)
}
pub fn encrypt_payload(
header: &FixedHeader,
payload_key: &SecretKey,
plaintext: &[u8],
zstd_level: i32,
) -> Result<(Vec<u8>, Vec<CompressionIndexEntry>)> {
encrypt_payload_with_progress(header, payload_key, plaintext, zstd_level, None)
}
pub fn encrypt_payload_with_progress(
header: &FixedHeader,
payload_key: &SecretKey,
plaintext: &[u8],
zstd_level: i32,
mut progress: Option<ProgressFn<'_>>,
) -> Result<(Vec<u8>, Vec<CompressionIndexEntry>)> {
let compressed = header.compress_id == COMPRESS_ZSTD;
let chunk_size = header.chunk_size as usize;
let n_chunks = header.n_chunks();
let n_chunks_u32 = u32::try_from(n_chunks)
.map_err(|_| Error::MalformedHeader { reason: "chunk count exceeds u32" })?;
let mut out = Vec::new();
let mut index = Vec::new();
for i in 0..n_chunks_u32 {
let start = (i as usize).saturating_mul(chunk_size);
let end = start.saturating_add(chunk_size).min(plaintext.len());
let raw = plaintext.get(start..end).unwrap_or(&[]);
let is_final = i.saturating_add(1) == n_chunks_u32;
let body = if compressed {
zstd::bulk::compress(raw, zstd_level)
.map_err(|e| Error::Compression { reason: e.to_string() })?
} else {
raw.to_vec()
};
let ct = encrypt_chunk(
payload_key,
header.cipher_id,
&header.base_nonce,
&header.file_uuid,
i,
is_final,
&body,
)?;
if compressed {
index.push(CompressionIndexEntry {
ct_offset: out.len() as u64,
ct_len: u32::try_from(ct.len()).map_err(|_| Error::MalformedHeader {
reason: "chunk ciphertext exceeds u32",
})?,
plain_len: u32::try_from(raw.len()).map_err(|_| Error::MalformedHeader {
reason: "chunk plaintext exceeds u32",
})?,
});
}
out.extend_from_slice(&ct);
if let Some(cb) = progress.as_deref_mut() {
cb(end as u64, plaintext.len() as u64);
}
}
Ok((out, index))
}
pub fn encrypt_payload_stream<R: std::io::Read>(
header: &FixedHeader,
payload_key: &SecretKey,
input: &mut R,
zstd_level: i32,
sink: &mut dyn FnMut(&[u8]) -> Result<()>,
mut progress: Option<ProgressFn<'_>>,
) -> Result<([u8; 32], Vec<CompressionIndexEntry>)> {
use blake2::digest::{Digest, consts::U32};
use std::io::Read as _;
let compressed = header.compress_id == COMPRESS_ZSTD;
let chunk_size = u64::from(header.chunk_size);
let total = header.plaintext_size;
let n_chunks_u32 = u32::try_from(header.n_chunks())
.map_err(|_| Error::MalformedHeader { reason: "chunk count exceeds u32" })?;
let mut hasher = blake2::Blake2b::<U32>::new();
let mut index = Vec::new();
let mut ct_offset = 0u64;
let mut done = 0u64;
let mut raw = Vec::with_capacity(usize::try_from(chunk_size).unwrap_or(0));
for i in 0..n_chunks_u32 {
let want = total.saturating_sub(done).min(chunk_size);
raw.clear();
let got = (&mut *input).take(want).read_to_end(&mut raw).map_err(Error::Io)?;
if got as u64 != want {
return Err(Error::Io(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
"source shrank while encrypting",
)));
}
hasher.update(&raw);
let is_final = i.saturating_add(1) == n_chunks_u32;
let body = if compressed {
zstd::bulk::compress(&raw, zstd_level)
.map_err(|e| Error::Compression { reason: e.to_string() })?
} else {
raw.clone()
};
let ct = encrypt_chunk(
payload_key,
header.cipher_id,
&header.base_nonce,
&header.file_uuid,
i,
is_final,
&body,
)?;
if compressed {
index.push(CompressionIndexEntry {
ct_offset,
ct_len: u32::try_from(ct.len()).map_err(|_| Error::MalformedHeader {
reason: "chunk ciphertext exceeds u32",
})?,
plain_len: u32::try_from(raw.len()).map_err(|_| Error::MalformedHeader {
reason: "chunk plaintext exceeds u32",
})?,
});
}
sink(&ct)?;
ct_offset = ct_offset.saturating_add(ct.len() as u64);
done = done.saturating_add(want);
if let Some(cb) = progress.as_deref_mut() {
cb(done, total);
}
}
let mut probe = [0u8; 1];
if input.read(&mut probe).map_err(Error::Io)? != 0 {
return Err(Error::Io(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"source grew while encrypting",
)));
}
Ok((hasher.finalize().into(), index))
}
pub fn decrypt_payload(
header: &FixedHeader,
payload_key: &SecretKey,
ciphertext: &[u8],
index: Option<&[CompressionIndexEntry]>,
) -> Result<Vec<u8>> {
decrypt_payload_with_progress(header, payload_key, ciphertext, index, None)
}
pub fn decrypt_payload_with_progress(
header: &FixedHeader,
payload_key: &SecretKey,
ciphertext: &[u8],
index: Option<&[CompressionIndexEntry]>,
mut progress: Option<ProgressFn<'_>>,
) -> Result<Vec<u8>> {
let compressed = header.has_flag(flags::COMPRESSED);
let n_chunks = header.n_chunks();
let n_chunks_u32 = u32::try_from(n_chunks)
.map_err(|_| Error::MalformedHeader { reason: "chunk count exceeds u32" })?;
let mut out = usize::try_from(header.plaintext_size)
.map_or_else(|_| Vec::new(), Vec::with_capacity);
if compressed {
let idx = index.ok_or(Error::MissingTlv {
tlv_type: crate::tlv::types::COMPRESSION_INDEX,
})?;
if idx.len() as u64 != n_chunks {
return Err(Error::MalformedTlv {
tlv_type: crate::tlv::types::COMPRESSION_INDEX,
reason: "index entry count does not match chunk count",
});
}
for (i, e) in idx.iter().enumerate() {
let i_u32 = u32::try_from(i)
.map_err(|_| Error::MalformedHeader { reason: "chunk index exceeds u32" })?;
let start = usize::try_from(e.ct_offset)
.map_err(|_| Error::MalformedTlv {
tlv_type: crate::tlv::types::COMPRESSION_INDEX,
reason: "ct_offset exceeds addressable range",
})?;
let end = start.saturating_add(e.ct_len as usize);
let ct = ciphertext.get(start..end).ok_or(Error::Truncated {
context: "compressed chunk",
need: end,
got: ciphertext.len(),
})?;
let is_final = i_u32.saturating_add(1) == n_chunks_u32;
let body = decrypt_chunk(
payload_key,
header.cipher_id,
&header.base_nonce,
&header.file_uuid,
i_u32,
is_final,
ct,
)?;
let plain = decompress_chunk(&body, e.plain_len, header.chunk_size)?;
out.extend_from_slice(&plain);
if let Some(cb) = progress.as_deref_mut() {
cb(out.len() as u64, header.plaintext_size);
}
}
} else {
for i in 0..n_chunks_u32 {
let (off, len) = header.chunk_ct_range(u64::from(i))?;
let rel = off
.checked_sub(u64::from(header.header_len))
.ok_or(Error::MalformedHeader { reason: "chunk offset precedes payload" })?;
let start = usize::try_from(rel).map_err(|_| Error::MalformedHeader {
reason: "chunk offset exceeds addressable range",
})?;
let end = start.saturating_add(usize::try_from(len).map_err(|_| {
Error::MalformedHeader { reason: "chunk length exceeds addressable range" }
})?);
let ct = ciphertext.get(start..end).ok_or(Error::Truncated {
context: "payload chunk",
need: end,
got: ciphertext.len(),
})?;
let is_final = i.saturating_add(1) == n_chunks_u32;
let plain = decrypt_chunk(
payload_key,
header.cipher_id,
&header.base_nonce,
&header.file_uuid,
i,
is_final,
ct,
)?;
out.extend_from_slice(&plain);
if let Some(cb) = progress.as_deref_mut() {
cb(out.len() as u64, header.plaintext_size);
}
}
}
if out.len() as u64 != header.plaintext_size {
return Err(Error::MalformedHeader {
reason: "decrypted length does not match plaintext_size",
});
}
Ok(out)
}
pub fn read_range<F>(
header: &FixedHeader,
payload_key: &SecretKey,
index: Option<&[CompressionIndexEntry]>,
offset: u64,
length: u64,
mut fetch_ct: F,
) -> Result<Vec<u8>>
where
F: FnMut(u64, u64) -> Result<Vec<u8>>,
{
if length == 0 {
return Ok(Vec::new());
}
let end = offset
.checked_add(length)
.ok_or(Error::MalformedHeader { reason: "range end overflows" })?
.min(header.plaintext_size);
if offset >= header.plaintext_size {
return Ok(Vec::new());
}
let n_chunks = header.n_chunks();
let compressed = header.has_flag(flags::COMPRESSED);
let cs = u64::from(header.chunk_size);
let first = header.chunk_index_of(offset);
let last = header.chunk_index_of(end.saturating_sub(1));
let mut acc = Vec::with_capacity(usize::try_from(end.saturating_sub(offset)).unwrap_or(0));
for i in first..=last {
if i >= n_chunks {
return Err(Error::ChunkOutOfRange { index: i, total: n_chunks });
}
let i_u32 =
u32::try_from(i).map_err(|_| Error::MalformedHeader { reason: "chunk index exceeds u32" })?;
let is_final = i.saturating_add(1) == n_chunks;
let plain = if compressed {
let idx = index.ok_or(Error::MissingTlv {
tlv_type: crate::tlv::types::COMPRESSION_INDEX,
})?;
let e = idx.get(usize::try_from(i).unwrap_or(usize::MAX)).ok_or(
Error::MalformedTlv {
tlv_type: crate::tlv::types::COMPRESSION_INDEX,
reason: "index shorter than chunk count",
},
)?;
let ct = fetch_ct(e.ct_offset, u64::from(e.ct_len))?;
let body = decrypt_chunk(
payload_key,
header.cipher_id,
&header.base_nonce,
&header.file_uuid,
i_u32,
is_final,
&ct,
)?;
decompress_chunk(&body, e.plain_len, header.chunk_size)?
} else {
let (abs_off, ct_len) = header.chunk_ct_range(i)?;
let rel = abs_off
.checked_sub(u64::from(header.header_len))
.ok_or(Error::MalformedHeader { reason: "chunk offset precedes payload" })?;
let ct = fetch_ct(rel, ct_len)?;
decrypt_chunk(
payload_key,
header.cipher_id,
&header.base_nonce,
&header.file_uuid,
i_u32,
is_final,
&ct,
)?
};
let chunk_start = i.saturating_mul(cs);
let want_from = offset.saturating_sub(chunk_start);
let want_to = end.saturating_sub(chunk_start).min(plain.len() as u64);
if want_from < want_to {
let a = usize::try_from(want_from).unwrap_or(0);
let b = usize::try_from(want_to).unwrap_or(plain.len());
if let Some(s) = plain.get(a..b) {
acc.extend_from_slice(s);
}
}
}
Ok(acc)
}
#[must_use]
pub fn payload_ct_size(header: &FixedHeader) -> u64 {
let n = header.n_chunks();
header
.plaintext_size
.saturating_add(n.saturating_mul(TAG_LEN as u64))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::crypto::{Argon2Params, Fek, KEY_LEN, Kek, SecretKey};
use crate::file::{EncryptOptions, RandomMaterial, encrypt, open};
use crate::header::{MIN_CHUNK_SIZE, SLOT_AREA_LEN, SLOT_LEN};
fn make(plain: &[u8]) -> (Vec<u8>, Kek) {
let salt = [0x31u8; 16];
let params = Argon2Params::TEST_WEAK;
let kek = Kek::from_password(b"payload-test", &salt, params).unwrap();
let kek2 = Kek::from_password(b"payload-test", &salt, params).unwrap();
let opts = EncryptOptions {
chunk_size: MIN_CHUNK_SIZE,
argon2: params,
..EncryptOptions::default()
};
let enc = encrypt(plain, &[kek], &salt, &opts, &RandomMaterial::generate()).unwrap();
(enc.bytes, kek2)
}
#[test]
fn fetch_ct_offset_is_payload_relative() {
let plain: Vec<u8> = (0..150_000u32).map(|i| (i % 253) as u8).collect();
let (bytes, kek) = make(&plain);
let opened = open(&bytes, &[kek]).unwrap();
let header_len = u64::from(opened.header.header_len);
assert!(header_len > 0, "头部长度应为正");
let seen: std::cell::RefCell<Vec<(u64, u64)>> = std::cell::RefCell::new(Vec::new());
let got = read_range(
&opened.header,
opened.payload_key(),
None,
0,
plain.len() as u64,
|off, len| {
seen.borrow_mut().push((off, len));
let a = usize::try_from(off + header_len).unwrap();
let b = usize::try_from(off + header_len + len).unwrap();
Ok(bytes[a..b].to_vec())
},
)
.unwrap();
assert_eq!(got, plain, "正确加回 header_len 后应完整还原");
let calls = seen.borrow();
assert_eq!(calls[0].0, 0, "首个 fetch_ct 偏移应为 0,即相对载荷起点");
let wrong = read_range(
&opened.header,
opened.payload_key(),
None,
0,
1000,
|off, len| {
let a = usize::try_from(off).unwrap();
let b = usize::try_from(off + len).unwrap();
Ok(bytes[a..b].to_vec())
},
);
assert!(
matches!(wrong, Err(Error::ChunkAuthFailed { .. })),
"把偏移当绝对值使用必须触发认证失败,实际: {wrong:?}"
);
}
#[test]
fn read_range_only_fetches_needed_chunks() {
let cs = u64::from(MIN_CHUNK_SIZE);
let plain: Vec<u8> = (0..(cs * 5) as u32).map(|i| (i % 251) as u8).collect();
let (bytes, kek) = make(&plain);
let opened = open(&bytes, &[kek]).unwrap();
let header_len = u64::from(opened.header.header_len);
assert_eq!(opened.header.n_chunks(), 5, "应为 5 块");
let n_calls = std::cell::Cell::new(0u32);
let total_ct = std::cell::Cell::new(0u64);
let off = cs * 3 + 500;
let got = read_range(
&opened.header,
opened.payload_key(),
None,
off,
100,
|o, l| {
n_calls.set(n_calls.get() + 1);
total_ct.set(total_ct.get() + l);
let a = usize::try_from(o + header_len).unwrap();
let b = usize::try_from(o + header_len + l).unwrap();
Ok(bytes[a..b].to_vec())
},
)
.unwrap();
assert_eq!(got, &plain[off as usize..(off + 100) as usize]);
assert_eq!(n_calls.get(), 1, "只跨 1 块就只该取 1 次");
assert!(
total_ct.get() <= cs + TAG_LEN as u64,
"实读密文 {} 超过单块上限,说明退化成了多块读取",
total_ct.get()
);
assert!(
total_ct.get() * 2 < plain.len() as u64,
"实读密文不应接近整个文件"
);
}
#[test]
fn read_range_across_chunk_boundary() {
let cs = u64::from(MIN_CHUNK_SIZE);
let plain: Vec<u8> = (0..(cs * 3) as u32).map(|i| (i % 249) as u8).collect();
let (bytes, kek) = make(&plain);
let opened = open(&bytes, &[kek]).unwrap();
let header_len = u64::from(opened.header.header_len);
let fetch = |o: u64, l: u64| -> Result<Vec<u8>> {
let a = usize::try_from(o + header_len).unwrap();
let b = usize::try_from(o + header_len + l).unwrap();
Ok(bytes[a..b].to_vec())
};
let off = cs - 50;
let got = read_range(&opened.header, opened.payload_key(), None, off, 100, fetch).unwrap();
assert_eq!(got, &plain[off as usize..(off + 100) as usize]);
let off2 = cs * 3 - 30;
let got2 = read_range(&opened.header, opened.payload_key(), None, off2, 30, fetch).unwrap();
assert_eq!(got2, &plain[off2 as usize..]);
let got3 = read_range(&opened.header, opened.payload_key(), None, off2, 9999, fetch).unwrap();
assert_eq!(got3, &plain[off2 as usize..], "超出部分应被裁剪");
let got4 = read_range(
&opened.header,
opened.payload_key(),
None,
plain.len() as u64,
10,
fetch,
)
.unwrap();
assert!(got4.is_empty(), "起点超界应返回空而非报错");
}
#[test]
fn progress_is_monotonic_and_reaches_total() {
let cs = u64::from(MIN_CHUNK_SIZE);
let plain: Vec<u8> = (0..(cs * 3 + 123) as u32).map(|i| (i % 251) as u8).collect();
let salt = [7u8; 16];
let params = Argon2Params::TEST_WEAK;
let kek = Kek::from_password(b"progress-test", &salt, params).unwrap();
let kek2 = Kek::from_password(b"progress-test", &salt, params).unwrap();
let opts = EncryptOptions {
chunk_size: MIN_CHUNK_SIZE,
argon2: params,
..EncryptOptions::default()
};
let mut seen: Vec<(u64, u64)> = Vec::new();
let enc = {
let mut cb = |done: u64, total: u64| seen.push((done, total));
crate::file::encrypt_with_progress(
&plain,
&[kek],
&salt,
&opts,
&RandomMaterial::generate(),
Some(&mut cb),
)
.unwrap()
};
assert_eq!(seen.len(), 4, "3 个整块 + 1 个尾块,应回调 4 次");
assert!(
seen.windows(2).all(|w| w[1].0 >= w[0].0),
"进度必须单调不减,实际: {seen:?}"
);
assert!(
seen.iter().all(|&(_, total)| total == plain.len() as u64),
"总长必须始终是明文长度"
);
assert_eq!(
seen.last().map(|&(d, _)| d),
Some(plain.len() as u64),
"最后一次必须报满,否则界面停在 99%"
);
let opened = open(&enc.bytes, &[kek2]).unwrap();
let mut dseen: Vec<u64> = Vec::new();
let out = {
let mut dcb = |done: u64, _total: u64| dseen.push(done);
opened
.decrypt_all_with_progress(&enc.bytes, Some(&mut dcb))
.unwrap()
};
assert_eq!(out, plain, "带进度回调不能改变解密结果");
assert_eq!(
dseen.last().copied(),
Some(plain.len() as u64),
"解密进度也必须报满"
);
}
#[test]
fn progress_reports_plaintext_bytes_when_compressed() {
let cs = u64::from(MIN_CHUNK_SIZE);
let plain = vec![0u8; (cs * 3) as usize];
let salt = [9u8; 16];
let params = Argon2Params::TEST_WEAK;
let kek = Kek::from_password(b"zip-progress", &salt, params).unwrap();
let opts = EncryptOptions {
chunk_size: MIN_CHUNK_SIZE,
argon2: params,
compress: true,
..EncryptOptions::default()
};
let mut last = 0u64;
let enc = {
let mut cb = |done: u64, _total: u64| last = done;
crate::file::encrypt_with_progress(
&plain,
&[kek],
&salt,
&opts,
&RandomMaterial::generate(),
Some(&mut cb),
)
.unwrap()
};
assert_eq!(last, plain.len() as u64, "压缩时进度仍应按明文计");
assert!(
enc.bytes.len() < plain.len(),
"前提校验:全零数据应当被显著压缩,否则这个测试没测到东西"
);
}
#[test]
fn progress_does_not_change_output() {
let plain: Vec<u8> = (0..5000u32).map(|i| (i % 251) as u8).collect();
let salt = [3u8; 16];
let params = Argon2Params::TEST_WEAK;
let kek = Kek::from_password(b"same-out", &salt, params).unwrap();
let opts = EncryptOptions {
chunk_size: MIN_CHUNK_SIZE,
argon2: params,
..EncryptOptions::default()
};
let key_bytes: [u8; KEY_LEN] = core::array::from_fn(|i| (i as u8).wrapping_mul(7) ^ 0x3D);
let fek = Fek::from_key(SecretKey::from_bytes(key_bytes));
let fek2 = Fek::from_key(SecretKey::from_bytes(key_bytes));
let rnd = RandomMaterial {
file_uuid: core::array::from_fn(|i| (i as u8).wrapping_mul(11).wrapping_add(2)),
base_nonce: core::array::from_fn(|i| (i as u8).wrapping_mul(13).wrapping_add(5)),
slot_padding: Some(
(0..SLOT_AREA_LEN - SLOT_LEN).map(|i| (i % 251) as u8).collect(),
),
};
let a = crate::file::encrypt_with_fek(
&plain,
&[kek.duplicate()],
&salt,
&opts,
&rnd,
&fek,
)
.unwrap();
let mut hits = 0u32;
let b = {
let mut cb = |_d: u64, _t: u64| hits += 1;
crate::file::encrypt_with_fek_and_progress(
&plain,
&[kek],
&salt,
&opts,
&rnd,
&fek2,
Some(&mut cb),
)
.unwrap()
};
assert_eq!(a.bytes, b.bytes, "带不带进度回调,产出必须逐字节相同");
assert!(hits > 0, "前提校验:回调确实被调用过,否则上面的相等是白比的");
}
}