use crate::crypto::{Argon2Params, CipherId, Fek, Kek, SecretKey, content_hash, header_mac, verify_header_mac};
use crate::error::{Error, Result};
use crate::payload;
use crate::header::{
COMPRESS_NONE, COMPRESS_ZSTD, DEFAULT_CHUNK_SIZE, FIXED_HEADER_LEN, FixedHeader,
HEADER_MAC_LEN, KDF_ARGON2ID, SLOT_AREA_LEN, SLOT_AREA_OFFSET, SLOT_COUNT, SLOT_COUNT_U8,
TLV_AREA_OFFSET, VERSION_MAJOR, VERSION_MINOR, flags, validate_chunk_size,
};
use crate::slot::{build_slot_area, unwrap_fek};
use crate::tlv::{
CompressionIndexEntry, TlvEntry, TlvSet, decode_compression_index, encode_compression_index,
encrypt_entry, pad_filename, tlv_flags, types, unpad_filename,
};
use crate::util::{Writer, random_16};
#[derive(Debug, Clone)]
pub struct EncryptOptions {
pub filename: Option<String>,
pub preserve_extension: bool,
pub compress: bool,
pub zstd_level: i32,
pub chunk_size: u32,
pub cipher: CipherId,
pub argon2: Argon2Params,
pub write_content_hash: bool,
pub thumbnail: Option<Vec<u8>>,
pub media_meta: Option<Vec<u8>>,
pub moov_cache: Option<Vec<u8>>,
pub folder_index: Option<Vec<u8>>,
pub slot_directory: Option<Vec<u8>>,
}
impl Default for EncryptOptions {
fn default() -> Self {
Self {
filename: None,
preserve_extension: false,
compress: false,
zstd_level: 3,
chunk_size: DEFAULT_CHUNK_SIZE,
cipher: CipherId::ChaCha20Poly1305,
argon2: Argon2Params::default(),
write_content_hash: true,
thumbnail: None,
media_meta: None,
moov_cache: None,
folder_index: None,
slot_directory: None,
}
}
}
#[derive(Debug, Clone)]
pub struct RandomMaterial {
pub file_uuid: [u8; 16],
pub base_nonce: [u8; 7],
pub slot_padding: Option<Vec<u8>>,
}
impl RandomMaterial {
#[must_use]
pub fn generate() -> Self {
let mut bn = [0u8; 7];
crate::util::fill_random(&mut bn);
Self { file_uuid: random_16(), base_nonce: bn, slot_padding: None }
}
}
#[derive(Debug)]
pub struct EncryptedFile {
pub bytes: Vec<u8>,
pub header: FixedHeader,
}
pub fn encrypt(
plaintext: &[u8],
keks: &[Kek],
vault_salt: &[u8; 16],
opts: &EncryptOptions,
rnd: &RandomMaterial,
) -> Result<EncryptedFile> {
encrypt_with_progress(plaintext, keks, vault_salt, opts, rnd, None)
}
pub fn encrypt_with_progress(
plaintext: &[u8],
keks: &[Kek],
vault_salt: &[u8; 16],
opts: &EncryptOptions,
rnd: &RandomMaterial,
progress: Option<payload::ProgressFn<'_>>,
) -> Result<EncryptedFile> {
if keks.is_empty() {
return Err(Error::TooManySlots { got: 0, max: SLOT_COUNT });
}
validate_chunk_size(opts.chunk_size)?;
let fek = Fek::random();
encrypt_with_fek_and_progress(plaintext, keks, vault_salt, opts, rnd, &fek, progress)
}
pub fn encrypt_with_fek(
plaintext: &[u8],
keks: &[Kek],
vault_salt: &[u8; 16],
opts: &EncryptOptions,
rnd: &RandomMaterial,
fek: &Fek,
) -> Result<EncryptedFile> {
encrypt_with_fek_and_progress(plaintext, keks, vault_salt, opts, rnd, fek, None)
}
pub fn encrypt_with_fek_and_progress(
plaintext: &[u8],
keks: &[Kek],
vault_salt: &[u8; 16],
opts: &EncryptOptions,
rnd: &RandomMaterial,
fek: &Fek,
progress: Option<payload::ProgressFn<'_>>,
) -> Result<EncryptedFile> {
if opts.chunk_size == 0 {
return Err(Error::MalformedHeader { reason: "chunk_size must not be zero" });
}
let mut header = new_header(opts, rnd, vault_salt, plaintext.len() as u64);
let payload_key = fek.derive_payload_key(&header.file_uuid);
let (payload, index) = payload::encrypt_payload_with_progress(
&header,
&payload_key,
plaintext,
opts.zstd_level,
progress,
)?;
let hash = opts.write_content_hash.then(|| content_hash(plaintext));
let tlvs = build_tlvs(opts, fek, hash.as_ref(), &index, opts.folder_index.as_deref())?;
let head = seal_header(&mut header, &tlvs, keks, fek, rnd, opts.cipher)?;
let mut out = Writer::with_capacity(head.len().saturating_add(payload.len()));
out.bytes(&head).bytes(&payload);
let bytes = out.into_vec();
debug_assert_eq!(
bytes.len(),
(header.header_len as usize).saturating_add(payload.len()),
"输出长度与 header_len + payload 不符"
);
Ok(EncryptedFile { bytes, header })
}
fn file_flags(opts: &EncryptOptions) -> u32 {
let mut file_flags = 0u32;
if opts.filename.is_some() {
file_flags |= flags::FILENAME_ENCRYPTED;
}
if opts.preserve_extension {
file_flags |= flags::EXT_PRESERVED;
}
if opts.compress {
file_flags |= flags::COMPRESSED;
}
if opts.thumbnail.is_some() {
file_flags |= flags::HAS_THUMBNAIL;
}
if opts.slot_directory.is_some() {
file_flags |= flags::SLOT_DIRECTORY;
}
if opts.folder_index.is_some() {
file_flags |= flags::CONTAINER;
}
file_flags
}
fn new_header(
opts: &EncryptOptions,
rnd: &RandomMaterial,
vault_salt: &[u8; 16],
plaintext_size: u64,
) -> FixedHeader {
FixedHeader {
version_major: VERSION_MAJOR,
version_minor: VERSION_MINOR,
header_len: 0,
file_uuid: rnd.file_uuid,
vault_salt: *vault_salt,
flags: file_flags(opts),
cipher_id: opts.cipher,
kdf_id: KDF_ARGON2ID,
compress_id: if opts.compress { COMPRESS_ZSTD } else { COMPRESS_NONE },
slot_count: SLOT_COUNT_U8,
argon2_m_kib: opts.argon2.m_kib,
argon2_t: opts.argon2.t,
argon2_p: opts.argon2.p,
chunk_size: opts.chunk_size,
plaintext_size,
base_nonce: rnd.base_nonce,
chunk_version: 0,
tlv_len: 0,
}
}
fn build_tlvs(
opts: &EncryptOptions,
fek: &Fek,
hash: Option<&[u8; 32]>,
index: &[CompressionIndexEntry],
folder_index: Option<&[u8]>,
) -> Result<TlvSet> {
let mut tlvs = TlvSet::new();
if let Some(name) = &opts.filename {
let padded = pad_filename(name)?;
tlvs.push(encrypt_entry(
types::FILENAME,
tlv_flags::CRITICAL,
&padded,
fek,
opts.cipher,
)?);
}
if opts.preserve_extension {
let ext = opts
.filename
.as_deref()
.and_then(|n| n.rsplit_once('.').map(|(_, e)| e))
.unwrap_or("");
tlvs.push(TlvEntry::new(types::PLAIN_EXT, 0, ext.as_bytes().to_vec()));
}
if let Some(t) = &opts.thumbnail {
tlvs.push(encrypt_entry(types::THUMBNAIL, 0, t, fek, opts.cipher)?);
}
if let Some(m) = &opts.media_meta {
tlvs.push(encrypt_entry(types::MEDIA_META, 0, m, fek, opts.cipher)?);
}
if let Some(mv) = &opts.moov_cache {
tlvs.push(encrypt_entry(types::MOOV_CACHE, 0, mv, fek, opts.cipher)?);
}
if let Some(idx) = folder_index {
tlvs.push(encrypt_entry(
types::FOLDER_INDEX,
tlv_flags::CRITICAL,
idx,
fek,
opts.cipher,
)?);
}
if let Some(dir) = &opts.slot_directory {
tlvs.push(encrypt_entry(
types::SLOT_DIRECTORY,
tlv_flags::CRITICAL,
dir,
fek,
opts.cipher,
)?);
}
if let Some(h) = hash {
tlvs.push(encrypt_entry(types::CONTENT_HASH, 0, h, fek, opts.cipher)?);
}
if opts.compress {
let enc = encode_compression_index(index);
tlvs.push(encrypt_entry(
types::COMPRESSION_INDEX,
tlv_flags::CRITICAL,
&enc,
fek,
opts.cipher,
)?);
}
Ok(tlvs)
}
fn seal_header(
header: &mut FixedHeader,
tlvs: &TlvSet,
keks: &[Kek],
fek: &Fek,
rnd: &RandomMaterial,
cipher: CipherId,
) -> Result<Vec<u8>> {
let tlv_blob = tlvs.to_bytes();
header.tlv_len = u32::try_from(tlv_blob.len())
.map_err(|_| Error::MalformedHeader { reason: "tlv area exceeds u32" })?;
header.header_len = u32::try_from(
TLV_AREA_OFFSET
.saturating_add(tlv_blob.len())
.saturating_add(HEADER_MAC_LEN),
)
.map_err(|_| Error::MalformedHeader { reason: "header_len exceeds u32" })?;
let fixed = header.to_bytes()?;
let slot_area = if let Some(pad) = &rnd.slot_padding {
build_slot_area_with_padding(keks, fek, &header.file_uuid, cipher, pad)?
} else {
build_slot_area(keks, fek, &header.file_uuid, cipher)?
};
let mut covered = Writer::with_capacity(header.header_len as usize);
covered.bytes(&fixed).bytes(&slot_area).bytes(&tlv_blob);
let mac_key = fek.derive_header_mac_key(&header.file_uuid);
let mac = header_mac(&mac_key, covered.as_slice());
let mut out = covered.into_vec();
out.extend_from_slice(&mac);
Ok(out)
}
pub type FinalizeFolderIndex<'a> = &'a mut dyn FnMut() -> Result<Vec<u8>>;
#[derive(Debug, Clone)]
pub struct StreamedFile {
pub header: FixedHeader,
pub head: Vec<u8>,
pub len: u64,
}
#[allow(clippy::too_many_arguments)]
pub fn encrypt_stream<R: std::io::Read, W: std::io::Write + std::io::Seek>(
input: &mut R,
plaintext_len: u64,
keks: &[Kek],
vault_salt: &[u8; 16],
opts: &EncryptOptions,
rnd: &RandomMaterial,
out: &mut W,
progress: Option<payload::ProgressFn<'_>>,
finalize_folder_index: Option<FinalizeFolderIndex<'_>>,
) -> Result<StreamedFile> {
if keks.is_empty() {
return Err(Error::TooManySlots { got: 0, max: SLOT_COUNT });
}
validate_chunk_size(opts.chunk_size)?;
let fek = Fek::random();
encrypt_stream_with_fek(
input,
plaintext_len,
keks,
vault_salt,
opts,
rnd,
&fek,
out,
progress,
finalize_folder_index,
)
}
#[allow(clippy::too_many_arguments)]
pub fn encrypt_stream_with_fek<R: std::io::Read, W: std::io::Write + std::io::Seek>(
input: &mut R,
plaintext_len: u64,
keks: &[Kek],
vault_salt: &[u8; 16],
opts: &EncryptOptions,
rnd: &RandomMaterial,
fek: &Fek,
out: &mut W,
progress: Option<payload::ProgressFn<'_>>,
finalize_folder_index: Option<FinalizeFolderIndex<'_>>,
) -> Result<StreamedFile> {
use std::io::SeekFrom;
if opts.chunk_size == 0 {
return Err(Error::MalformedHeader { reason: "chunk_size must not be zero" });
}
let mut header = new_header(opts, rnd, vault_salt, plaintext_len);
let n_chunks = usize::try_from(header.n_chunks())
.map_err(|_| Error::MalformedHeader { reason: "chunk count exceeds usize" })?;
let zero_hash = [0u8; 32];
let placeholder_index = vec![
CompressionIndexEntry { ct_offset: 0, ct_len: 0, plain_len: 0 };
if opts.compress { n_chunks } else { 0 }
];
let placeholder = build_tlvs(
opts,
fek,
opts.write_content_hash.then_some(&zero_hash),
&placeholder_index,
opts.folder_index.as_deref(),
)?;
let reserved = TLV_AREA_OFFSET
.saturating_add(placeholder.to_bytes().len())
.saturating_add(HEADER_MAC_LEN);
let base = out.stream_position().map_err(Error::Io)?;
out.write_all(&vec![0u8; reserved]).map_err(Error::Io)?;
let payload_key = fek.derive_payload_key(&header.file_uuid);
let (hash, index) = payload::encrypt_payload_stream(
&header,
&payload_key,
input,
opts.zstd_level,
&mut |ct| out.write_all(ct).map_err(Error::Io),
progress,
)?;
let folder_index = match (finalize_folder_index, opts.folder_index.as_deref()) {
(Some(f), Some(placeholder_idx)) => {
let real = f()?;
if real.len() != placeholder_idx.len() {
return Err(Error::MalformedHeader {
reason: "finalized folder index length differs from placeholder",
});
}
Some(real)
}
(_, other) => other.map(<[u8]>::to_vec),
};
let tlvs = build_tlvs(
opts,
fek,
opts.write_content_hash.then_some(&hash),
&index,
folder_index.as_deref(),
)?;
let head = seal_header(&mut header, &tlvs, keks, fek, rnd, opts.cipher)?;
if head.len() != reserved {
return Err(Error::MalformedHeader { reason: "streamed header length changed" });
}
let end = out.stream_position().map_err(Error::Io)?;
out.seek(SeekFrom::Start(base)).map_err(Error::Io)?;
out.write_all(&head).map_err(Error::Io)?;
out.seek(SeekFrom::Start(end)).map_err(Error::Io)?;
out.flush().map_err(Error::Io)?;
Ok(StreamedFile { header, head, len: end.saturating_sub(base) })
}
fn build_slot_area_with_padding(
keks: &[Kek],
fek: &Fek,
file_uuid: &[u8; 16],
cipher: CipherId,
padding: &[u8],
) -> Result<Vec<u8>> {
use crate::crypto::ZERO_NONCE;
use crate::header::SLOT_LEN;
if keks.len() > SLOT_COUNT {
return Err(Error::TooManySlots { got: keks.len(), max: SLOT_COUNT });
}
let mut w = Writer::with_capacity(SLOT_AREA_LEN);
for (i, kek) in keks.iter().enumerate() {
let idx = u16::try_from(i)
.map_err(|_| Error::TooManySlots { got: keks.len(), max: SLOT_COUNT })?;
let wk = kek.derive_slot_key(file_uuid, idx);
w.bytes(&cipher.encrypt(&wk, &ZERO_NONCE, fek.as_key().as_bytes(), &[])?);
}
let need = SLOT_AREA_LEN.saturating_sub(keks.len().saturating_mul(SLOT_LEN));
let pad = padding.get(..need).ok_or(Error::MalformedHeader {
reason: "provided slot padding shorter than required",
})?;
w.bytes(pad);
Ok(w.into_vec())
}
#[derive(Debug)]
pub struct OpenedFile {
pub header: FixedHeader,
pub tlvs: TlvSet,
pub kek_index: usize,
pub slot_index: u16,
fek: Fek,
payload_key: SecretKey,
}
impl OpenedFile {
#[must_use]
pub const fn payload_key(&self) -> &SecretKey {
&self.payload_key
}
#[must_use]
pub const fn fek(&self) -> &Fek {
&self.fek
}
pub fn filename(&self) -> Result<String> {
let padded = self.tlvs.decrypt_value(types::FILENAME, &self.fek, self.header.cipher_id)?;
unpad_filename(&padded)
}
#[must_use]
pub fn plain_extension(&self) -> Option<String> {
self.tlvs
.find(types::PLAIN_EXT)
.map(|e| String::from_utf8_lossy(&e.value).into_owned())
}
pub fn thumbnail(&self) -> Result<Vec<u8>> {
self.tlvs.decrypt_value(types::THUMBNAIL, &self.fek, self.header.cipher_id)
}
pub fn media_meta(&self) -> Result<Vec<u8>> {
self.tlvs.decrypt_value(types::MEDIA_META, &self.fek, self.header.cipher_id)
}
pub fn moov_cache(&self) -> Result<Vec<u8>> {
self.tlvs.decrypt_value(types::MOOV_CACHE, &self.fek, self.header.cipher_id)
}
#[must_use]
pub fn has_moov_cache(&self) -> bool {
self.tlvs.find(types::MOOV_CACHE).is_some()
}
#[must_use]
pub fn has_media_meta(&self) -> bool {
self.tlvs.find(types::MEDIA_META).is_some()
}
#[must_use]
pub fn has_thumbnail(&self) -> bool {
self.tlvs.find(types::THUMBNAIL).is_some()
}
pub fn stored_content_hash(&self) -> Result<[u8; 32]> {
let v = self.tlvs.decrypt_value(types::CONTENT_HASH, &self.fek, self.header.cipher_id)?;
let s = v.get(..32).ok_or(Error::MalformedTlv {
tlv_type: types::CONTENT_HASH,
reason: "content hash must be 32 bytes",
})?;
let mut out = [0u8; 32];
out.copy_from_slice(s);
Ok(out)
}
pub fn compression_index(&self) -> Result<Vec<CompressionIndexEntry>> {
if !self.header.has_flag(flags::COMPRESSED) {
return Ok(Vec::new());
}
let raw =
self.tlvs.decrypt_value(types::COMPRESSION_INDEX, &self.fek, self.header.cipher_id)?;
decode_compression_index(&raw)
}
pub fn folder_index(&self) -> Result<crate::container::ContainerIndex> {
if !self.header.has_flag(flags::CONTAINER) {
return Err(Error::MissingTlv { tlv_type: types::FOLDER_INDEX });
}
let raw = self
.tlvs
.decrypt_value(types::FOLDER_INDEX, &self.fek, self.header.cipher_id)?;
crate::container::ContainerIndex::parse(&raw)
}
#[must_use]
pub const fn is_container(&self) -> bool {
self.header.has_flag(flags::CONTAINER)
}
#[must_use]
pub const fn is_slot_managed(&self) -> bool {
self.header.has_flag(flags::SLOT_DIRECTORY)
}
pub fn slot_directory(&self) -> Result<crate::slotdir::SlotDirectory> {
if !self.is_slot_managed() {
return Err(Error::MissingTlv { tlv_type: types::SLOT_DIRECTORY });
}
let raw = self
.tlvs
.decrypt_value(types::SLOT_DIRECTORY, &self.fek, self.header.cipher_id)?;
crate::slotdir::SlotDirectory::decode(&raw)
}
pub fn raw_slot_directory(&self) -> Result<Vec<u8>> {
if !self.is_slot_managed() {
return Err(Error::MissingTlv { tlv_type: types::SLOT_DIRECTORY });
}
self.tlvs
.decrypt_value(types::SLOT_DIRECTORY, &self.fek, self.header.cipher_id)
}
pub fn raw_folder_index(&self) -> Result<Vec<u8>> {
if !self.header.has_flag(flags::CONTAINER) {
return Err(Error::MissingTlv { tlv_type: types::FOLDER_INDEX });
}
self.tlvs.decrypt_value(types::FOLDER_INDEX, &self.fek, self.header.cipher_id)
}
pub fn decrypt_all(&self, file_bytes: &[u8]) -> Result<Vec<u8>> {
self.decrypt_all_with_progress(file_bytes, None)
}
pub fn decrypt_all_with_progress(
&self,
file_bytes: &[u8],
progress: Option<payload::ProgressFn<'_>>,
) -> Result<Vec<u8>> {
let start = self.header.header_len as usize;
let payload = file_bytes.get(start..).ok_or(Error::Truncated {
context: "payload",
need: start,
got: file_bytes.len(),
})?;
let idx = self.compression_index()?;
let idx_ref = if idx.is_empty() { None } else { Some(idx.as_slice()) };
let plain = payload::decrypt_payload_with_progress(
&self.header,
&self.payload_key,
payload,
idx_ref,
progress,
)?;
if let Ok(expected) = self.stored_content_hash() {
use subtle::ConstantTimeEq;
let actual = content_hash(&plain);
if !bool::from(actual.ct_eq(&expected)) {
return Err(Error::ContentHashMismatch);
}
}
Ok(plain)
}
}
pub fn peek_header(data: &[u8]) -> Result<FixedHeader> {
FixedHeader::parse(data)
}
#[must_use]
pub fn is_omy_file(data: &[u8]) -> bool {
data.get(..8).is_some_and(|m| m == crate::header::MAGIC_FILE)
}
pub fn open(data: &[u8], keks: &[Kek]) -> Result<OpenedFile> {
let header = FixedHeader::parse(data)?;
let hlen = header.header_len as usize;
if data.len() < hlen {
return Err(Error::Truncated { context: "header", need: hlen, got: data.len() });
}
let slot_end = SLOT_AREA_OFFSET.saturating_add(SLOT_AREA_LEN);
let slot_area = data.get(SLOT_AREA_OFFSET..slot_end).ok_or(Error::Truncated {
context: "key slot area",
need: slot_end,
got: data.len(),
})?;
let tlv_end = TLV_AREA_OFFSET.saturating_add(header.tlv_len as usize);
let tlv_blob = data.get(TLV_AREA_OFFSET..tlv_end).ok_or(Error::Truncated {
context: "tlv area",
need: tlv_end,
got: data.len(),
})?;
let mac_end = tlv_end.saturating_add(HEADER_MAC_LEN);
let stored_mac = data.get(tlv_end..mac_end).ok_or(Error::Truncated {
context: "header mac",
need: mac_end,
got: data.len(),
})?;
let (fek, kek_index, slot_index) =
unwrap_fek(slot_area, keks, &header.file_uuid, header.cipher_id)?;
let covered = data.get(..tlv_end).ok_or(Error::Truncated {
context: "mac-covered region",
need: tlv_end,
got: data.len(),
})?;
let mac_key = fek.derive_header_mac_key(&header.file_uuid);
let mut expected = [0u8; HEADER_MAC_LEN];
let src = stored_mac.get(..HEADER_MAC_LEN).ok_or(Error::Truncated {
context: "header mac",
need: HEADER_MAC_LEN,
got: stored_mac.len(),
})?;
expected.copy_from_slice(src);
if !verify_header_mac(&mac_key, covered, &expected) {
return Err(Error::HeaderMacMismatch);
}
let tlvs = TlvSet::parse(tlv_blob)?;
tlvs.check_critical()?;
if header.has_flag(flags::COMPRESSED) && tlvs.find(types::COMPRESSION_INDEX).is_none() {
return Err(Error::MissingTlv { tlv_type: types::COMPRESSION_INDEX });
}
let payload_key = fek.derive_payload_key(&header.file_uuid);
Ok(OpenedFile { header, tlvs, kek_index, slot_index, fek, payload_key })
}
pub fn open_with_password(data: &[u8], password: &[u8]) -> Result<OpenedFile> {
let header = FixedHeader::parse(data)?;
let params = Argon2Params {
m_kib: header.argon2_m_kib,
t: header.argon2_t,
p: header.argon2_p,
};
let kek = Kek::from_password(password, &header.vault_salt, params)?;
open(data, &[kek])
}
pub const PEEK_SIZE: usize = FIXED_HEADER_LEN;