1use crate::error::{Context, Error, bail, ensure};
2use crate::format::{ContentType, RecordKind, Urma, chunk_count};
3use aes::cipher::{KeyIvInit, StreamCipher};
4use hkdf::Hkdf;
5use hmac::{Hmac, Mac};
6use rand::{CryptoRng, RngCore};
7use sha2::{Digest, Sha256};
8use std::collections::BTreeMap;
9use zeroize::Zeroizing;
10
11type Cipher = ctr::Ctr128BE<aes::Aes256>;
12type Authenticator = Hmac<Sha256>;
13
14pub fn random_secret<R: RngCore + CryptoRng>(rng: &mut R) -> Result<Zeroizing<[u8; 32]>, Error> {
15 let mut key = Zeroizing::new([0; 32]);
16 rng.try_fill_bytes(key.as_mut())?;
17 Ok(key)
18}
19fn derive(root: &[u8; 32], id: &[u8; 32], purpose: &[u8], output: &mut [u8]) -> Result<(), Error> {
20 Hkdf::<Sha256>::new(Some(id), root)
21 .expand(purpose, output)
22 .map_err(Error::Kdf)?;
23 Ok(())
24}
25fn discovery_tag(root: &[u8; 32], id: &[u8; 32]) -> Result<[u8; 16], Error> {
26 let mut tag = [0; 16];
27 derive(root, id, Urma::DISCOVERY_DOMAIN, &mut tag)?;
28 Ok(tag)
29}
30fn cipher(root: &[u8; 32], id: &[u8; 32], iv: &[u8; 16]) -> Result<Cipher, Error> {
31 let mut key = Zeroizing::new([0; 32]);
32 derive(root, id, Urma::CONTENT_DOMAIN, key.as_mut())?;
33 Ok(Cipher::new((&*key).into(), iv.into()))
34}
35fn nonce(index: u32) -> [u8; 16] {
36 let mut iv = [0; 16];
37 iv[..8].copy_from_slice(&u64::from(index).to_be_bytes());
38 iv
39}
40fn authenticator(root: &[u8; 32], id: &[u8; 32]) -> Result<Authenticator, Error> {
41 let mut key = Zeroizing::new([0; 32]);
42 derive(root, id, Urma::AUTHENTICATION_DOMAIN, key.as_mut())?;
43 Ok(Authenticator::new_from_slice(key.as_ref())?)
44}
45pub fn seal<R: RngCore + CryptoRng>(
46 root: &[u8; 32],
47 plaintext: &[u8],
48 content_type: ContentType,
49 rng: &mut R,
50) -> Result<Vec<Vec<u8>>, Error> {
51 let count = chunk_count(u64::try_from(plaintext.len())?)?;
52 let id = *random_secret(rng)?;
53 let tag = discovery_tag(root, &id)?;
54 let digest: [u8; 32] = Sha256::digest(plaintext).into();
55 let mut records = Vec::new();
56 records.try_reserve_exact(usize::try_from(count)?)?;
57 for (position, chunk) in plaintext.chunks(Urma::CHUNK_BYTES).enumerate() {
58 let index = u32::try_from(position)?;
59 let mut record = Vec::with_capacity(Urma::PRIVATE_RECORD_BYTES);
60 record.extend_from_slice(&RecordKind::Private.prefix());
61 record.extend_from_slice(&id);
62 record.extend_from_slice(&tag);
63 record.extend_from_slice(&index.to_le_bytes());
64 record.extend_from_slice(&count.to_le_bytes());
65 let iv = nonce(index);
66 record.extend_from_slice(&iv);
67 let mut body = Zeroizing::new(vec![0; Urma::BODY_BYTES]);
68 body[..32].copy_from_slice(&digest);
69 body[32..40].copy_from_slice(&u64::try_from(plaintext.len())?.to_le_bytes());
70 body[40..44].copy_from_slice(&content_type.code().to_le_bytes());
71 body[Urma::METADATA_BYTES..Urma::METADATA_BYTES + chunk.len()].copy_from_slice(chunk);
72 rng.try_fill_bytes(&mut body[Urma::METADATA_BYTES + chunk.len()..])?;
73 cipher(root, &id, &iv)?.apply_keystream(&mut body);
74 record.extend_from_slice(&body);
75 let mut mac = authenticator(root, &id)?;
76 mac.update(&record);
77 record.extend_from_slice(&mac.finalize().into_bytes());
78 records.push(record);
79 }
80 Ok(records)
81}
82pub struct PrivateChunk {
83 pub id: [u8; 32],
84 pub index: u32,
85 pub count: u32,
86 pub total: u64,
87 pub digest: [u8; 32],
88 pub content_type: ContentType,
89 pub bytes: Zeroizing<Vec<u8>>,
90 record_digest: [u8; 32],
91}
92#[derive(Clone, Copy, Debug, PartialEq, Eq)]
93pub struct PrivateRecordHeader {
94 pub id: [u8; 32],
95 pub discovery_tag: [u8; 16],
96 pub index: u32,
97 pub count: u32,
98}
99pub enum RecordMatch {
100 Unrelated,
101 Authenticated(PrivateChunk),
102}
103pub fn inspect_header(record: &[u8]) -> Result<PrivateRecordHeader, Error> {
104 ensure!(
105 record.len() == Urma::PRIVATE_RECORD_BYTES,
106 "invalid private record length"
107 );
108 ensure!(
109 RecordKind::parse(record)? == RecordKind::Private,
110 "not a private record"
111 );
112 let header = PrivateRecordHeader {
113 id: record[8..40].try_into()?,
114 discovery_tag: record[40..56].try_into()?,
115 index: u32::from_le_bytes(record[56..60].try_into()?),
116 count: u32::from_le_bytes(record[60..64].try_into()?),
117 };
118 ensure!(
119 header.count > 0 && header.index < header.count,
120 "invalid chunk bounds"
121 );
122 Ok(header)
123}
124pub fn open_record(root: &[u8; 32], record: &[u8]) -> Result<RecordMatch, Error> {
125 if !record.starts_with(&Urma::MAGIC) {
126 return Ok(RecordMatch::Unrelated);
127 }
128 match RecordKind::parse(record)? {
129 RecordKind::Private => {}
130 RecordKind::Post
131 | RecordKind::Reply
132 | RecordKind::Profile
133 | RecordKind::Avatar
134 | RecordKind::Container
135 | RecordKind::DataPart
136 | RecordKind::LeafManifest
137 | RecordKind::WirePost
138 | RecordKind::WireReply
139 | RecordKind::ProfileRecord
140 | RecordKind::RootManifest => return Ok(RecordMatch::Unrelated),
141 }
142 let header = inspect_header(record)?;
143 if header.discovery_tag != discovery_tag(root, &header.id)? {
144 return Ok(RecordMatch::Unrelated);
145 }
146 let mut mac = authenticator(root, &header.id)?;
147 mac.update(&record[..Urma::MAC_OFFSET]);
148 mac.verify_slice(&record[Urma::MAC_OFFSET..])
149 .map_err(|error| Error::Invalid(format!("record authentication failed: {error}")))?;
150 let iv = nonce(header.index);
151 ensure!(
152 record[Urma::PRIVATE_HEADER_BYTES..Urma::BODY_OFFSET] == iv,
153 "noncanonical counter IV"
154 );
155 let mut body = Zeroizing::new(record[Urma::BODY_OFFSET..Urma::MAC_OFFSET].to_vec());
156 cipher(root, &header.id, &iv)?.apply_keystream(&mut body);
157 let total = u64::from_le_bytes(body[32..40].try_into()?);
158 ensure!(
159 chunk_count(total)? == header.count,
160 "chunk count does not match encrypted file length"
161 );
162 let content_type = ContentType::from_code(u32::from_le_bytes(body[40..44].try_into()?))?;
163 if body[44..48] != [0; 4] {
164 return Err(Error::Unsupported(
165 "unsupported private metadata flags".into(),
166 ));
167 }
168 let remaining = total
169 .checked_sub(u64::from(header.index) * u64::try_from(Urma::CHUNK_BYTES)?)
170 .context("chunk offset exceeds original length")?;
171 let size = usize::try_from(remaining.min(u64::try_from(Urma::CHUNK_BYTES)?))?;
172 Ok(RecordMatch::Authenticated(PrivateChunk {
173 id: header.id,
174 index: header.index,
175 count: header.count,
176 total,
177 digest: body[..32].try_into()?,
178 content_type,
179 bytes: Zeroizing::new(body[Urma::METADATA_BYTES..Urma::METADATA_BYTES + size].to_vec()),
180 record_digest: Sha256::digest(record).into(),
181 }))
182}
183pub struct PrivateObject {
184 pub id: [u8; 32],
185 pub count: u32,
186 pub total: u64,
187 pub content_type: ContentType,
188 digest: [u8; 32],
189 chunks: BTreeMap<u32, PrivateChunk>,
190 conflict: bool,
191}
192impl PrivateObject {
193 pub fn new(chunk: PrivateChunk) -> Self {
194 Self {
195 id: chunk.id,
196 count: chunk.count,
197 total: chunk.total,
198 content_type: chunk.content_type,
199 digest: chunk.digest,
200 chunks: BTreeMap::from([(chunk.index, chunk)]),
201 conflict: false,
202 }
203 }
204 pub fn insert(&mut self, chunk: PrivateChunk) -> Result<(), Error> {
205 if self.id != chunk.id
206 || self.count != chunk.count
207 || self.total != chunk.total
208 || self.digest != chunk.digest
209 || self.content_type != chunk.content_type
210 {
211 self.conflict = true;
212 bail!("conflicting authenticated object metadata");
213 }
214 match self.chunks.entry(chunk.index) {
215 std::collections::btree_map::Entry::Occupied(previous) => {
216 if previous.get().record_digest != chunk.record_digest {
217 self.conflict = true;
218 bail!("conflicting authenticated record");
219 }
220 }
221 std::collections::btree_map::Entry::Vacant(slot) => {
222 slot.insert(chunk);
223 }
224 }
225 Ok(())
226 }
227 pub fn received(&self) -> usize {
228 self.chunks.len()
229 }
230 pub fn is_conflicted(&self) -> bool {
231 self.conflict
232 }
233 pub fn is_complete(&self) -> bool {
234 match u32::try_from(self.chunks.len()) {
235 Ok(count) => !self.conflict && count == self.count,
236 Err(error) => panic!("chunk map exceeds u32 key space: {error}"),
237 }
238 }
239 pub fn finish(&self) -> Result<Zeroizing<Vec<u8>>, Error> {
240 ensure!(!self.conflict, "conflicting authenticated object");
241 ensure!(
242 u32::try_from(self.chunks.len())? == self.count,
243 "incomplete object: {} of {} chunks",
244 self.chunks.len(),
245 self.count
246 );
247 let mut bytes = Zeroizing::new(Vec::new());
248 bytes.try_reserve_exact(
249 usize::try_from(self.total).context("object exceeds host capacity")?,
250 )?;
251 for index in 0..self.count {
252 bytes.extend_from_slice(&self.chunks.get(&index).context("missing chunk")?.bytes);
253 }
254 ensure!(
255 u64::try_from(bytes.len())? == self.total
256 && Sha256::digest(&bytes).as_slice() == self.digest,
257 "whole-file integrity failed"
258 );
259 Ok(bytes)
260 }
261}
262pub fn pack(records: &[Vec<u8>]) -> Result<Vec<u8>, Error> {
263 let count = u32::try_from(records.len())?;
264 ensure!(count > 0, "empty container");
265 let first = inspect_header(&records[0])?;
266 let capacity = records
267 .len()
268 .checked_mul(Urma::PRIVATE_RECORD_BYTES + 4)
269 .and_then(|n| n.checked_add(Urma::CONTAINER_HEADER_BYTES))
270 .context("container size overflow")?;
271 let mut bytes = Vec::new();
272 bytes.try_reserve_exact(capacity)?;
273 bytes.extend_from_slice(&RecordKind::Container.prefix());
274 bytes.extend_from_slice(&count.to_le_bytes());
275 for record in records {
276 ensure!(inspect_header(record)?.id == first.id, "mixed objects");
277 bytes.extend_from_slice(&u32::try_from(record.len())?.to_le_bytes());
278 bytes.extend_from_slice(record);
279 }
280 Ok(bytes)
281}
282pub fn unpack(bytes: &[u8]) -> Result<Vec<Vec<u8>>, Error> {
283 let entries = framed_records(bytes)?;
284 let mut records = Vec::new();
285 records.try_reserve_exact(entries.len())?;
286 for entry in entries {
287 let record = framed_record(entry)?;
288 inspect_header(record)?;
289 records.push(record.to_vec());
290 }
291 let id = inspect_header(&records[0])?.id;
292 for record in &records {
293 ensure!(inspect_header(record)?.id == id, "mixed objects");
294 }
295 Ok(records)
296}
297fn framed_records(bytes: &[u8]) -> Result<&[[u8; Urma::PRIVATE_RECORD_BYTES + 4]], Error> {
298 ensure!(
299 RecordKind::parse(bytes)? == RecordKind::Container,
300 "not an offline container"
301 );
302 ensure!(
303 bytes.len() >= Urma::CONTAINER_HEADER_BYTES,
304 "truncated container"
305 );
306 let count = usize::try_from(u32::from_le_bytes(bytes[8..12].try_into()?))?;
307 ensure!(count > 0, "empty container");
308 let size = count
309 .checked_mul(Urma::PRIVATE_RECORD_BYTES + 4)
310 .and_then(|n| n.checked_add(Urma::CONTAINER_HEADER_BYTES))
311 .context("container size overflow")?;
312 ensure!(bytes.len() == size, "truncated or trailing container data");
313 Ok(bytes[Urma::CONTAINER_HEADER_BYTES..]
314 .as_chunks::<{ Urma::PRIVATE_RECORD_BYTES + 4 }>()
315 .0)
316}
317fn framed_record(entry: &[u8; Urma::PRIVATE_RECORD_BYTES + 4]) -> Result<&[u8], Error> {
318 ensure!(
319 u32::from_le_bytes(entry[..4].try_into()?) == u32::try_from(Urma::PRIVATE_RECORD_BYTES)?,
320 "invalid container record length"
321 );
322 Ok(&entry[4..])
323}
324fn required_chunk(root: &[u8; 32], record: &[u8]) -> Result<PrivateChunk, Error> {
325 match open_record(root, record)? {
326 RecordMatch::Authenticated(chunk) => Ok(chunk),
327 RecordMatch::Unrelated => bail!("wrong key or unrelated record"),
328 }
329}
330pub fn open(root: &[u8; 32], records: &[Vec<u8>]) -> Result<Zeroizing<Vec<u8>>, Error> {
331 let (first, rest) = records.split_first().context("no records")?;
332 let mut object = PrivateObject::new(required_chunk(root, first)?);
333 for record in rest {
334 object.insert(required_chunk(root, record)?)?;
335 }
336 object.finish()
337}
338pub struct ContainerRecovery {
339 pub id: [u8; 32],
340 pub count: u32,
341 pub total: u64,
342 pub digest: [u8; 32],
343 pub content_type: ContentType,
344 pub bytes: Zeroizing<Vec<u8>>,
345 pub skipped_unrelated: usize,
346 pub rejected_records: usize,
347}
348enum RecoveryRecord {
349 Matched(RecordMatch),
350 Rejected(String),
351}
352fn recovery_record(root: &[u8; 32], record: &[u8]) -> Result<RecoveryRecord, Error> {
353 let candidate = match open_record(root, record) {
354 Ok(matched) => RecoveryRecord::Matched(matched),
355 Err(error) => classify_rejection(error)?,
356 };
357 Ok(candidate)
358}
359fn classify_rejection(error: Error) -> Result<RecoveryRecord, Error> {
360 match error {
361 Error::Invalid(reason) | Error::Unsupported(reason) => Ok(RecoveryRecord::Rejected(reason)),
362 fatal => Err(fatal),
363 }
364}
365pub fn recover_container(root: &[u8; 32], bytes: &[u8]) -> Result<ContainerRecovery, Error> {
366 let mut objects = BTreeMap::new();
367 let mut skipped_unrelated = 0;
368 let mut rejected_records = 0;
369 let entries = framed_records(bytes)?;
370 for entry in entries {
371 framed_record(entry)?;
372 }
373 for entry in entries {
374 match recovery_record(root, &entry[4..])? {
375 RecoveryRecord::Matched(RecordMatch::Unrelated) => skipped_unrelated += 1,
376 RecoveryRecord::Rejected(reason) => {
377 drop(reason);
378 rejected_records += 1;
379 }
380 RecoveryRecord::Matched(RecordMatch::Authenticated(chunk)) => {
381 match objects.entry(chunk.id) {
382 std::collections::btree_map::Entry::Vacant(slot) => {
383 slot.insert(PrivateObject::new(chunk));
384 }
385 std::collections::btree_map::Entry::Occupied(mut slot) => {
386 slot.get_mut().insert(chunk)?;
387 }
388 }
389 ensure!(objects.len() == 1, "multiple authenticated objects");
390 }
391 }
392 }
393 let object = objects
394 .into_values()
395 .next()
396 .ok_or_else(|| Error::Invalid("no authenticated object".into()))?;
397 Ok(ContainerRecovery {
398 bytes: object.finish()?,
399 id: object.id,
400 count: object.count,
401 total: object.total,
402 digest: object.digest,
403 content_type: object.content_type,
404 skipped_unrelated,
405 rejected_records,
406 })
407}