use crate::crypto::{Fek, Kek, header_mac};
use crate::error::{Error, Result};
use crate::file::{EncryptOptions, EncryptedFile, RandomMaterial};
use crate::header::{FixedHeader, HEADER_MAC_LEN, SLOT_AREA_LEN, SLOT_AREA_OFFSET, flags};
use crate::payload::ProgressFn;
use crate::tlv::types;
#[derive(Debug)]
pub struct RotateOutcome {
pub bytes: Vec<u8>,
pub opened_slot: u16,
pub slot_used: usize,
pub plaintext_size: u64,
}
#[derive(Debug, Clone, Default)]
pub struct ReplacementMetadata {
pub thumbnail: Option<Vec<u8>>,
pub media_meta: Option<Vec<u8>>,
pub moov_cache: Option<Vec<u8>>,
}
pub fn replace_plaintext_preserving_slots(
data: &[u8],
unlock: &[Kek],
plaintext: &[u8],
metadata: &ReplacementMetadata,
new_base_nonce: [u8; 7],
) -> Result<EncryptedFile> {
let opened = crate::file::open(data, unlock)?;
if opened.is_container() {
return Err(Error::MalformedHeader {
reason: "container payload cannot be replaced as a single file",
});
}
let old = FixedHeader::parse(data)?;
let slot_end = SLOT_AREA_OFFSET.saturating_add(SLOT_AREA_LEN);
let old_slots = data
.get(SLOT_AREA_OFFSET..slot_end)
.ok_or(Error::Truncated {
context: "key slot area",
need: slot_end,
got: data.len(),
})?;
let mut opts = rebuild_options(&opened, &old, 1)?;
opts.slot_directory = if opened.is_slot_managed() {
Some(opened.raw_slot_directory()?)
} else {
None
};
opts.thumbnail = metadata.thumbnail.clone();
opts.media_meta = metadata.media_meta.clone();
opts.moov_cache = metadata.moov_cache.clone();
let keeper = unlock
.get(opened.kek_index)
.ok_or(Error::NoMatchingSlot)?
.duplicate();
let rnd = RandomMaterial {
file_uuid: old.file_uuid,
base_nonce: new_base_nonce,
slot_padding: None,
};
let mut out = crate::file::encrypt_with_fek(
plaintext,
&[keeper],
&old.vault_salt,
&opts,
&rnd,
opened.fek(),
)?;
let new_header = FixedHeader::parse(&out.bytes)?;
let new_slot_end = SLOT_AREA_OFFSET.saturating_add(SLOT_AREA_LEN);
let output_len = out.bytes.len();
let dst = out
.bytes
.get_mut(SLOT_AREA_OFFSET..new_slot_end)
.ok_or(Error::Truncated {
context: "replacement key slot area",
need: new_slot_end,
got: output_len,
})?;
dst.copy_from_slice(old_slots);
let mac_pos = (new_header.header_len as usize).saturating_sub(HEADER_MAC_LEN);
let covered = out.bytes.get(..mac_pos).ok_or(Error::Truncated {
context: "replacement header",
need: mac_pos,
got: out.bytes.len(),
})?;
let mac_key = opened.fek().derive_header_mac_key(&old.file_uuid);
let mac = header_mac(&mac_key, covered);
let total_len = out.bytes.len();
let mac_dst = out
.bytes
.get_mut(mac_pos..new_header.header_len as usize)
.ok_or(Error::Truncated {
context: "replacement header mac",
need: new_header.header_len as usize,
got: total_len,
})?;
mac_dst.copy_from_slice(&mac);
out.header = new_header;
Ok(out)
}
pub fn rotate_fek(
data: &[u8],
unlock: &[Kek],
keep: &[Kek],
rnd: &RandomMaterial,
) -> Result<RotateOutcome> {
rotate_fek_with_progress(data, unlock, keep, rnd, None, None)
}
pub fn rotate_fek_with_progress(
data: &[u8],
unlock: &[Kek],
keep: &[Kek],
rnd: &RandomMaterial,
decrypting: Option<ProgressFn<'_>>,
encrypting: Option<ProgressFn<'_>>,
) -> Result<RotateOutcome> {
if keep.is_empty() {
return Err(Error::MalformedHeader {
reason: "refusing to leave a file with zero key slots; it could never be opened again",
});
}
let opened = crate::file::open(data, unlock)?;
let header = FixedHeader::parse(data)?;
let plain = opened.decrypt_all_with_progress(data, decrypting)?;
let opts = rebuild_options(&opened, &header, keep.len())?;
let fek = Fek::random();
let enc: EncryptedFile = crate::file::encrypt_with_fek_and_progress(
&plain,
keep,
&header.vault_salt,
&opts,
rnd,
&fek,
encrypting,
)?;
Ok(RotateOutcome {
bytes: enc.bytes,
opened_slot: opened.slot_index,
slot_used: keep.len(),
plaintext_size: plain.len() as u64,
})
}
fn rebuild_options(
opened: &crate::file::OpenedFile,
header: &FixedHeader,
kept: usize,
) -> Result<EncryptOptions> {
let filename = if header.has_flag(flags::FILENAME_ENCRYPTED) {
Some(opened.filename()?)
} else {
None
};
let thumbnail = if opened.has_thumbnail() {
Some(opened.thumbnail()?)
} else {
None
};
let media_meta = if opened.has_media_meta() {
Some(opened.media_meta()?)
} else {
None
};
let moov_cache = if opened.has_moov_cache() {
Some(opened.moov_cache()?)
} else {
None
};
let folder_index = if opened.is_container() {
Some(opened.raw_folder_index()?)
} else {
None
};
let slot_directory = if opened.is_slot_managed() {
let mut dir = crate::slotdir::SlotDirectory::new();
for i in 0..kept.min(crate::header::SLOT_COUNT) {
dir.set(
i,
crate::slotdir::SlotEntry::of(crate::slotdir::SlotKind::Vault),
)?;
}
Some(dir.encode())
} else {
None
};
Ok(EncryptOptions {
filename,
slot_directory,
preserve_extension: header.has_flag(flags::EXT_PRESERVED),
compress: header.has_flag(flags::COMPRESSED),
zstd_level: EncryptOptions::default().zstd_level,
chunk_size: header.chunk_size,
cipher: header.cipher_id,
argon2: header.argon2_params(),
write_content_hash: opened.tlvs.find(types::CONTENT_HASH).is_some(),
thumbnail,
media_meta,
moov_cache,
folder_index,
})
}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used, clippy::indexing_slicing, clippy::panic)]
use super::*;
use crate::container::EntryMeta;
use crate::crypto::Argon2Params;
use crate::file::{RandomMaterial, encrypt};
fn kek(pw: &str, salt: &[u8; 16]) -> Kek {
Kek::from_password(pw.as_bytes(), salt, Argon2Params::TEST_WEAK).unwrap()
}
fn sample(salt: &[u8; 16], pw: &[&str], plain: &[u8]) -> Vec<u8> {
let keks: Vec<Kek> = pw.iter().map(|p| kek(p, salt)).collect();
let opts = EncryptOptions {
filename: Some(String::from("报告.txt")),
preserve_extension: true,
argon2: Argon2Params::TEST_WEAK,
thumbnail: Some((0..64u32).map(|i| (i % 251) as u8).collect()),
media_meta: Some(br#"{"w":1920}"#.to_vec()),
..EncryptOptions::default()
};
encrypt(plain, &keks, salt, &opts, &RandomMaterial::generate())
.unwrap()
.bytes
}
#[test]
fn rotated_file_opens_with_kept_password_and_returns_same_plaintext() {
let salt = [0x31u8; 16];
let plain = b"rotate me".repeat(500);
let data = sample(&salt, &["pw-old"], &plain);
let out = rotate_fek(
&data,
&[kek("pw-old", &salt)],
&[kek("pw-new", &salt)],
&RandomMaterial::generate(),
)
.unwrap();
let re = crate::file::open(&out.bytes, &[kek("pw-new", &salt)]).unwrap();
assert_eq!(
re.decrypt_all(&out.bytes).unwrap(),
plain,
"轮换后必须还原出原明文"
);
assert_eq!(out.plaintext_size, plain.len() as u64);
}
#[test]
fn rotation_replaces_fek_and_uuid_so_old_ciphertext_is_gone() {
let salt = [0x32u8; 16];
let data = sample(&salt, &["pw-old"], b"payload bytes here");
let before = FixedHeader::parse(&data).unwrap();
let out = rotate_fek(
&data,
&[kek("pw-old", &salt)],
&[kek("pw-old", &salt)],
&RandomMaterial::generate(),
)
.unwrap();
let after = FixedHeader::parse(&out.bytes).unwrap();
assert_ne!(before.file_uuid, after.file_uuid, "必须换 file_uuid");
assert_ne!(before.base_nonce, after.base_nonce, "必须换 base_nonce");
assert_ne!(
data.get(before.header_len as usize..),
out.bytes.get(before.header_len as usize..),
"载荷密文必须全部改变,否则等于没轮换"
);
}
#[test]
fn old_fek_cannot_derive_new_payload_key() {
let salt = [0x3Eu8; 16];
let data = sample(&salt, &["pw-old"], b"payload bytes here");
let old_opened = crate::file::open(&data, &[kek("pw-old", &salt)]).unwrap();
let out = rotate_fek(
&data,
&[kek("pw-old", &salt)],
&[kek("pw-new", &salt)],
&RandomMaterial::generate(),
)
.unwrap();
let new_opened = crate::file::open(&out.bytes, &[kek("pw-new", &salt)]).unwrap();
let forged = old_opened
.fek()
.derive_payload_key(&new_opened.header.file_uuid);
assert_ne!(
forged.as_bytes(),
new_opened.payload_key().as_bytes(),
"旧 FEK 不得能推出新文件的载荷密钥,否则持有旧副本的人不需要新密码就能解开新文件"
);
}
#[test]
fn old_password_cannot_open_rotated_file() {
let salt = [0x33u8; 16];
let data = sample(&salt, &["pw-old"], b"secret");
let out = rotate_fek(
&data,
&[kek("pw-old", &salt)],
&[kek("pw-new", &salt)],
&RandomMaterial::generate(),
)
.unwrap();
assert!(crate::file::open(&out.bytes, &[kek("pw-old", &salt)]).is_err());
}
#[test]
fn rotation_preserves_filename_thumbnail_and_media_meta() {
let salt = [0x34u8; 16];
let data = sample(&salt, &["pw"], b"x");
let out = rotate_fek(
&data,
&[kek("pw", &salt)],
&[kek("pw", &salt)],
&RandomMaterial::generate(),
)
.unwrap();
let re = crate::file::open(&out.bytes, &[kek("pw", &salt)]).unwrap();
assert_eq!(re.filename().unwrap(), "报告.txt");
assert_eq!(re.plain_extension().as_deref(), Some("txt"));
assert_eq!(
re.thumbnail().unwrap(),
(0..64u32).map(|i| (i % 251) as u8).collect::<Vec<u8>>()
);
assert_eq!(re.media_meta().unwrap(), br#"{"w":1920}"#.to_vec());
assert!(
re.header.has_flag(flags::HAS_THUMBNAIL),
"flag 也要跟着保留"
);
}
#[test]
fn rotation_preserves_format_parameters() {
let salt = [0x35u8; 16];
let data = sample(&salt, &["pw"], b"y");
let before = FixedHeader::parse(&data).unwrap();
let out = rotate_fek(
&data,
&[kek("pw", &salt)],
&[kek("pw", &salt)],
&RandomMaterial::generate(),
)
.unwrap();
let after = FixedHeader::parse(&out.bytes).unwrap();
assert_eq!(before.chunk_size, after.chunk_size);
assert_eq!(before.cipher_id, after.cipher_id);
assert_eq!(
before.vault_salt, after.vault_salt,
"同一个 vault,salt 不能变"
);
assert_eq!(before.argon2_m_kib, after.argon2_m_kib);
assert_eq!(before.argon2_t, after.argon2_t);
assert_eq!(before.argon2_p, after.argon2_p);
assert_eq!(before.flags, after.flags, "flags 整体必须一致");
}
#[test]
fn rotation_can_change_password_set_at_the_same_time() {
let salt = [0x36u8; 16];
let data = sample(&salt, &["pw-a"], b"z");
let out = rotate_fek(
&data,
&[kek("pw-a", &salt)],
&[kek("pw-b", &salt), kek("pw-c", &salt)],
&RandomMaterial::generate(),
)
.unwrap();
assert_eq!(out.slot_used, 2);
for p in ["pw-b", "pw-c"] {
assert!(
crate::file::open(&out.bytes, &[kek(p, &salt)]).is_ok(),
"{p} 应能打开"
);
}
assert!(crate::file::open(&out.bytes, &[kek("pw-a", &salt)]).is_err());
}
#[test]
fn empty_keep_is_rejected() {
let salt = [0x37u8; 16];
let data = sample(&salt, &["pw"], b"w");
assert!(rotate_fek(&data, &[kek("pw", &salt)], &[], &RandomMaterial::generate()).is_err());
}
#[test]
fn wrong_unlock_password_fails() {
let salt = [0x38u8; 16];
let data = sample(&salt, &["pw"], b"v");
assert!(
rotate_fek(
&data,
&[kek("nope", &salt)],
&[kek("pw", &salt)],
&RandomMaterial::generate()
)
.is_err()
);
}
#[test]
fn tampered_header_is_rejected_before_rewriting() {
let salt = [0x39u8; 16];
let mut data = sample(&salt, &["pw"], b"u");
data[64] ^= 0xFF; assert!(
rotate_fek(
&data,
&[kek("pw", &salt)],
&[kek("pw", &salt)],
&RandomMaterial::generate()
)
.is_err()
);
}
#[test]
fn rotation_preserves_compression() {
let salt = [0x3Au8; 16];
let keks = [kek("pw", &salt)];
let plain = b"aaaaaaaaaaaaaaaa".repeat(4096);
let opts = EncryptOptions {
compress: true,
argon2: Argon2Params::TEST_WEAK,
..EncryptOptions::default()
};
let data = encrypt(&plain, &keks, &salt, &opts, &RandomMaterial::generate())
.unwrap()
.bytes;
assert!(
FixedHeader::parse(&data)
.unwrap()
.has_flag(flags::COMPRESSED)
);
let out = rotate_fek(&data, &keks, &keks, &RandomMaterial::generate()).unwrap();
let after = FixedHeader::parse(&out.bytes).unwrap();
assert!(after.has_flag(flags::COMPRESSED), "压缩设置必须保留");
let re = crate::file::open(&out.bytes, &keks).unwrap();
assert_eq!(re.decrypt_all(&out.bytes).unwrap(), plain);
}
#[test]
fn rotation_preserves_container_flag_and_index() {
let salt = [0x3Bu8; 16];
let keks = [kek("pw", &salt)];
let idx_bytes = {
let mut b = crate::container::ContainerBuilder::new("folder");
b.add_file(vec![String::from("a.txt")], 3, None, EntryMeta::default())
.unwrap();
b.finish().unwrap().encode()
};
let opts = EncryptOptions {
folder_index: Some(idx_bytes.clone()),
argon2: Argon2Params::TEST_WEAK,
..EncryptOptions::default()
};
let data = encrypt(b"abc", &keks, &salt, &opts, &RandomMaterial::generate())
.unwrap()
.bytes;
let out = rotate_fek(&data, &keks, &keks, &RandomMaterial::generate()).unwrap();
let re = crate::file::open(&out.bytes, &keks).unwrap();
assert!(re.is_container(), "CONTAINER flag 必须保留");
assert_eq!(
re.raw_folder_index().unwrap(),
idx_bytes,
"索引字节必须逐字节一致"
);
}
#[test]
fn rotation_keeps_content_hash_absent_when_original_had_none() {
let salt = [0x3Cu8; 16];
let keks = [kek("pw", &salt)];
let opts = EncryptOptions {
write_content_hash: false,
argon2: Argon2Params::TEST_WEAK,
..EncryptOptions::default()
};
let data = encrypt(b"t", &keks, &salt, &opts, &RandomMaterial::generate())
.unwrap()
.bytes;
let out = rotate_fek(&data, &keks, &keks, &RandomMaterial::generate()).unwrap();
let re = crate::file::open(&out.bytes, &keks).unwrap();
assert!(
re.stored_content_hash().is_err(),
"原文件没有哈希,轮换后也不该有"
);
}
#[test]
fn replacement_preserves_all_slots_and_uses_new_plaintext() {
let salt = core::array::from_fn(|i| (i as u8).wrapping_mul(11) ^ 0x4D);
let keys = [kek("pw-a", &salt), kek("pw-b", &salt)];
let opts = EncryptOptions {
filename: Some(String::from("报告.txt")),
argon2: Argon2Params::TEST_WEAK,
..EncryptOptions::default()
};
let original = encrypt(b"before", &keys, &salt, &opts, &RandomMaterial::generate())
.unwrap()
.bytes;
let before = FixedHeader::parse(&original).unwrap();
let slots_before = original
.get(SLOT_AREA_OFFSET..SLOT_AREA_OFFSET + SLOT_AREA_LEN)
.unwrap()
.to_vec();
let nonce = core::array::from_fn(|i| (i as u8).wrapping_mul(17) ^ 0xA3);
let replaced = replace_plaintext_preserving_slots(
&original,
&keys,
b"after external edit",
&ReplacementMetadata::default(),
nonce,
)
.unwrap();
let after = FixedHeader::parse(&replaced.bytes).unwrap();
assert_eq!(
before.file_uuid, after.file_uuid,
"槽位派生依赖 file_uuid,不能变"
);
assert_eq!(after.base_nonce, nonce, "必须采用新 nonce");
assert_ne!(before.base_nonce, after.base_nonce, "不得复用旧 nonce");
assert_eq!(
replaced
.bytes
.get(SLOT_AREA_OFFSET..SLOT_AREA_OFFSET + SLOT_AREA_LEN)
.unwrap(),
slots_before,
"所有已知和未知槽都必须逐字节保留",
);
for key in &keys {
let opened = crate::file::open(&replaced.bytes, &[key.duplicate()]).unwrap();
assert_eq!(
opened.decrypt_all(&replaced.bytes).unwrap(),
b"after external edit",
"每把旧钥匙都要能打开修改后的新内容",
);
assert_eq!(opened.filename().unwrap(), "报告.txt");
}
}
#[test]
fn replacement_rejects_container() {
let salt = core::array::from_fn(|i| (i as u8).wrapping_mul(5) ^ 0x71);
let keys = [kek("pw", &salt)];
let mut b = crate::container::ContainerBuilder::new("folder");
b.add_file(vec![String::from("a.txt")], 3, None, EntryMeta::default())
.unwrap();
let opts = EncryptOptions {
folder_index: Some(b.finish().unwrap().encode()),
argon2: Argon2Params::TEST_WEAK,
..EncryptOptions::default()
};
let original = encrypt(b"abc", &keys, &salt, &opts, &RandomMaterial::generate())
.unwrap()
.bytes;
assert!(
replace_plaintext_preserving_slots(
&original,
&keys,
b"flat file",
&ReplacementMetadata::default(),
[9; 7],
)
.is_err(),
"容器不能被第三方单文件编辑静默拍平",
);
}
#[test]
fn corrupted_payload_is_rejected_instead_of_being_re_signed() {
let salt = [0x3Du8; 16];
let keks = [kek("pw", &salt)];
let mut data = sample(&salt, &["pw"], &b"payload".repeat(100));
let last = data.len().saturating_sub(1);
data[last] ^= 0xFF;
assert!(rotate_fek(&data, &keks, &keks, &RandomMaterial::generate()).is_err());
}
}