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| Error::Invalid(format!("HKDF expansion: {error}")))?;
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::RootManifest => return Ok(RecordMatch::Unrelated),
138 }
139 let header = inspect_header(record)?;
140 if header.discovery_tag != discovery_tag(root, &header.id)? {
141 return Ok(RecordMatch::Unrelated);
142 }
143 let mut mac = authenticator(root, &header.id)?;
144 mac.update(&record[..Urma::MAC_OFFSET]);
145 mac.verify_slice(&record[Urma::MAC_OFFSET..])
146 .map_err(|error| Error::Invalid(format!("record authentication failed: {error}")))?;
147 let iv = nonce(header.index);
148 ensure!(
149 record[Urma::PRIVATE_HEADER_BYTES..Urma::BODY_OFFSET] == iv,
150 "noncanonical counter IV"
151 );
152 let mut body = Zeroizing::new(record[Urma::BODY_OFFSET..Urma::MAC_OFFSET].to_vec());
153 cipher(root, &header.id, &iv)?.apply_keystream(&mut body);
154 let total = u64::from_le_bytes(body[32..40].try_into()?);
155 ensure!(
156 chunk_count(total)? == header.count,
157 "chunk count does not match encrypted file length"
158 );
159 let content_type = ContentType::from_code(u32::from_le_bytes(body[40..44].try_into()?))?;
160 if body[44..48] != [0; 4] {
161 return Err(Error::Unsupported(
162 "unsupported private metadata flags".into(),
163 ));
164 }
165 let remaining = total
166 .checked_sub(u64::from(header.index) * u64::try_from(Urma::CHUNK_BYTES)?)
167 .context("chunk offset exceeds original length")?;
168 let size = usize::try_from(remaining.min(u64::try_from(Urma::CHUNK_BYTES)?))?;
169 Ok(RecordMatch::Authenticated(PrivateChunk {
170 id: header.id,
171 index: header.index,
172 count: header.count,
173 total,
174 digest: body[..32].try_into()?,
175 content_type,
176 bytes: Zeroizing::new(body[Urma::METADATA_BYTES..Urma::METADATA_BYTES + size].to_vec()),
177 record_digest: Sha256::digest(record).into(),
178 }))
179}
180pub struct PrivateObject {
181 pub id: [u8; 32],
182 pub count: u32,
183 pub total: u64,
184 pub content_type: ContentType,
185 digest: [u8; 32],
186 chunks: BTreeMap<u32, PrivateChunk>,
187 conflict: bool,
188}
189impl PrivateObject {
190 pub fn new(chunk: PrivateChunk) -> Self {
191 Self {
192 id: chunk.id,
193 count: chunk.count,
194 total: chunk.total,
195 content_type: chunk.content_type,
196 digest: chunk.digest,
197 chunks: BTreeMap::from([(chunk.index, chunk)]),
198 conflict: false,
199 }
200 }
201 pub fn insert(&mut self, chunk: PrivateChunk) -> Result<(), Error> {
202 if self.id != chunk.id
203 || self.count != chunk.count
204 || self.total != chunk.total
205 || self.digest != chunk.digest
206 || self.content_type != chunk.content_type
207 {
208 self.conflict = true;
209 bail!("conflicting authenticated object metadata");
210 }
211 match self.chunks.entry(chunk.index) {
212 std::collections::btree_map::Entry::Occupied(previous) => {
213 if previous.get().record_digest != chunk.record_digest {
214 self.conflict = true;
215 bail!("conflicting authenticated record");
216 }
217 }
218 std::collections::btree_map::Entry::Vacant(slot) => {
219 slot.insert(chunk);
220 }
221 }
222 Ok(())
223 }
224 pub fn received(&self) -> usize {
225 self.chunks.len()
226 }
227 pub fn is_conflicted(&self) -> bool {
228 self.conflict
229 }
230 pub fn is_complete(&self) -> bool {
231 match u32::try_from(self.chunks.len()) {
232 Ok(count) => !self.conflict && count == self.count,
233 Err(error) => panic!("chunk map exceeds u32 key space: {error}"),
234 }
235 }
236 pub fn finish(&self) -> Result<Zeroizing<Vec<u8>>, Error> {
237 ensure!(!self.conflict, "conflicting authenticated object");
238 ensure!(
239 u32::try_from(self.chunks.len())? == self.count,
240 "incomplete object: {} of {} chunks",
241 self.chunks.len(),
242 self.count
243 );
244 let mut bytes = Zeroizing::new(Vec::new());
245 bytes.try_reserve_exact(
246 usize::try_from(self.total).context("object exceeds host capacity")?,
247 )?;
248 for index in 0..self.count {
249 bytes.extend_from_slice(&self.chunks.get(&index).context("missing chunk")?.bytes);
250 }
251 ensure!(
252 u64::try_from(bytes.len())? == self.total
253 && Sha256::digest(&bytes).as_slice() == self.digest,
254 "whole-file integrity failed"
255 );
256 Ok(bytes)
257 }
258}
259pub fn pack(records: &[Vec<u8>]) -> Result<Vec<u8>, Error> {
260 let count = u32::try_from(records.len())?;
261 ensure!(count > 0, "empty container");
262 let first = inspect_header(&records[0])?;
263 let capacity = records
264 .len()
265 .checked_mul(Urma::PRIVATE_RECORD_BYTES + 4)
266 .and_then(|n| n.checked_add(Urma::CONTAINER_HEADER_BYTES))
267 .context("container size overflow")?;
268 let mut bytes = Vec::new();
269 bytes.try_reserve_exact(capacity)?;
270 bytes.extend_from_slice(&RecordKind::Container.prefix());
271 bytes.extend_from_slice(&count.to_le_bytes());
272 for record in records {
273 ensure!(inspect_header(record)?.id == first.id, "mixed objects");
274 bytes.extend_from_slice(&u32::try_from(record.len())?.to_le_bytes());
275 bytes.extend_from_slice(record);
276 }
277 Ok(bytes)
278}
279pub fn unpack(bytes: &[u8]) -> Result<Vec<Vec<u8>>, Error> {
280 ensure!(
281 RecordKind::parse(bytes)? == RecordKind::Container,
282 "not an offline container"
283 );
284 ensure!(
285 bytes.len() >= Urma::CONTAINER_HEADER_BYTES,
286 "truncated container"
287 );
288 let count = usize::try_from(u32::from_le_bytes(bytes[8..12].try_into()?))?;
289 ensure!(count > 0, "empty container");
290 let size = count
291 .checked_mul(Urma::PRIVATE_RECORD_BYTES + 4)
292 .and_then(|n| n.checked_add(Urma::CONTAINER_HEADER_BYTES))
293 .context("container size overflow")?;
294 ensure!(bytes.len() == size, "truncated or trailing container data");
295 let mut records = Vec::new();
296 records.try_reserve_exact(count)?;
297 for entry in bytes[Urma::CONTAINER_HEADER_BYTES..]
298 .as_chunks::<{ Urma::PRIVATE_RECORD_BYTES + 4 }>()
299 .0
300 {
301 ensure!(
302 u32::from_le_bytes(entry[..4].try_into()?)
303 == u32::try_from(Urma::PRIVATE_RECORD_BYTES)?,
304 "invalid container record length"
305 );
306 inspect_header(&entry[4..])?;
307 records.push(entry[4..].to_vec());
308 }
309 let id = inspect_header(&records[0])?.id;
310 for record in &records {
311 ensure!(inspect_header(record)?.id == id, "mixed objects");
312 }
313 Ok(records)
314}
315fn required_chunk(root: &[u8; 32], record: &[u8]) -> Result<PrivateChunk, Error> {
316 match open_record(root, record)? {
317 RecordMatch::Authenticated(chunk) => Ok(chunk),
318 RecordMatch::Unrelated => bail!("wrong key or unrelated record"),
319 }
320}
321pub fn open(root: &[u8; 32], records: &[Vec<u8>]) -> Result<Zeroizing<Vec<u8>>, Error> {
322 let (first, rest) = records.split_first().context("no records")?;
323 let mut object = PrivateObject::new(required_chunk(root, first)?);
324 for record in rest {
325 object.insert(required_chunk(root, record)?)?;
326 }
327 object.finish()
328}