Skip to main content

urma_core/
container.rs

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