ahu/
bundle.rs

1use sha2::{Digest, Sha256};
2use std::io::Cursor;
3use std::path::Path;
4
5use crate::error::{AhuError, Result};
6use crate::header::{FileHeader, HEADER_SIZE};
7use crate::index::{INDEX_RECORD_SIZE, IndexRecord};
8use crate::manifest::Manifest;
9
10/// A parsed ahu bundle, loaded into memory.
11#[derive(Debug)]
12pub struct Bundle {
13    pub header: FileHeader,
14    pub manifest: Manifest,
15    pub manifest_bytes: Vec<u8>,
16    pub seal_bytes: Vec<u8>,
17    pub index: Vec<IndexRecord>,
18    pub data: Vec<u8>,
19}
20
21impl Bundle {
22    /// Read a bundle from a file path.
23    pub fn from_file(path: &Path) -> Result<Self> {
24        let data = std::fs::read(path)?;
25        Self::from_bytes(&data)
26    }
27
28    /// Read a bundle from a byte slice.
29    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
30        let file_size = bytes.len() as u64;
31        let mut cursor = Cursor::new(bytes);
32
33        let header = FileHeader::read_from(&mut cursor)?;
34        header.validate_bounds(file_size)?;
35
36        let manifest_bytes =
37            Self::read_section(bytes, header.manifest_offset, header.manifest_length as u64)?;
38        let manifest = Manifest::from_cbor(&manifest_bytes)?;
39
40        let seal_bytes = Self::read_section(bytes, header.seal_offset, header.seal_length as u64)?;
41
42        let index_bytes = Self::read_section(bytes, header.index_offset, header.index_length)?;
43        let index = Self::parse_index(&index_bytes)?;
44
45        let data = Self::read_section(bytes, header.data_offset, header.data_length)?;
46
47        Ok(Bundle {
48            header,
49            manifest,
50            manifest_bytes,
51            seal_bytes,
52            index,
53            data,
54        })
55    }
56
57    /// Serialize the already parsed sections without changing the signed manifest.
58    pub fn to_bytes(&self) -> Result<Vec<u8>> {
59        let mut index = Vec::new();
60        for record in &self.index {
61            record.write_to(&mut index)?;
62        }
63        let mut header = self.header.clone();
64        header.manifest_offset = HEADER_SIZE as u64;
65        header.manifest_length = u32::try_from(self.manifest_bytes.len())
66            .map_err(|_| AhuError::InvalidOperation("manifest too large".into()))?;
67        header.seal_offset = header.manifest_offset + u64::from(header.manifest_length);
68        header.seal_length = u32::try_from(self.seal_bytes.len())
69            .map_err(|_| AhuError::InvalidOperation("seal too large".into()))?;
70        header.index_offset = header.seal_offset + u64::from(header.seal_length);
71        header.index_length = index.len() as u64;
72        header.data_offset = header
73            .index_offset
74            .checked_add(header.index_length)
75            .ok_or_else(|| AhuError::InvalidOperation("bundle too large".into()))?;
76        header.data_length = self.data.len() as u64;
77        let mut out = Vec::new();
78        header.write_to(&mut out)?;
79        out.extend_from_slice(&self.manifest_bytes);
80        out.extend_from_slice(&self.seal_bytes);
81        out.extend_from_slice(&index);
82        out.extend_from_slice(&self.data);
83        Ok(out)
84    }
85
86    /// Look up an entry by its entry key (SHA-256 of DER CertID).
87    /// Returns the default (discriminator=0) entry.
88    pub fn lookup(&self, entry_key: &[u8; 32]) -> Option<&[u8]> {
89        let idx = crate::index::binary_search(&self.index, entry_key)?;
90        self.entry_at(idx)
91    }
92
93    /// Look up the best-matching entry given a list of preferred discriminators.
94    pub fn lookup_preferred(&self, entry_key: &[u8; 32], preferences: &[u16]) -> Option<&[u8]> {
95        let idx = crate::index::binary_search_preferred(&self.index, entry_key, preferences)?;
96        self.entry_at(idx)
97    }
98
99    pub fn entry_at(&self, idx: usize) -> Option<&[u8]> {
100        let record = self.index.get(idx)?;
101        if record.is_tombstone() {
102            return None;
103        }
104        let start = usize::try_from(record.data_offset).ok()?;
105        let length = usize::try_from(record.data_length).ok()?;
106        let end = start.checked_add(length)?;
107        self.data.get(start..end)
108    }
109
110    /// Get the raw response bytes for an index record.
111    pub fn entry_bytes(&self, record: &IndexRecord) -> Option<&[u8]> {
112        if record.is_tombstone() {
113            return None;
114        }
115        let start = usize::try_from(record.data_offset).ok()?;
116        let length = usize::try_from(record.data_length).ok()?;
117        let end = start.checked_add(length)?;
118        self.data.get(start..end)
119    }
120
121    fn read_section(bytes: &[u8], offset: u64, length: u64) -> Result<Vec<u8>> {
122        let range = usize::try_from(offset).ok().and_then(|start| {
123            usize::try_from(length)
124                .ok()
125                .and_then(|len| start.checked_add(len).map(|end| start..end))
126        });
127        range
128            .and_then(|r| bytes.get(r))
129            .map(<[u8]>::to_vec)
130            .ok_or(AhuError::HeaderOutOfBounds {
131                field: "section",
132                offset,
133                length,
134                file_size: bytes.len() as u64,
135            })
136    }
137
138    fn parse_index(index_bytes: &[u8]) -> Result<Vec<IndexRecord>> {
139        if index_bytes.len() % INDEX_RECORD_SIZE != 0 {
140            return Err(AhuError::IndexSizeMismatch {
141                size: index_bytes.len() as u64,
142                record_size: INDEX_RECORD_SIZE,
143            });
144        }
145
146        let count = index_bytes.len() / INDEX_RECORD_SIZE;
147        let mut records = Vec::with_capacity(count);
148        let mut cursor = Cursor::new(index_bytes);
149
150        for _ in 0..count {
151            records.push(IndexRecord::read_from(&mut cursor)?);
152        }
153        Ok(records)
154    }
155}
156
157/// Builder for constructing ahu bundles.
158pub struct BundleBuilder {
159    pub manifest: Manifest,
160    entries: Vec<(IndexRecord, Vec<u8>)>,
161}
162
163impl BundleBuilder {
164    pub fn new(manifest: Manifest) -> Self {
165        BundleBuilder {
166            manifest,
167            entries: Vec::new(),
168        }
169    }
170
171    /// Add an entry. `certid_der` is the DER encoding of the CertID;
172    /// `response_der` is the complete DER-encoded OCSPResponse.
173    pub fn add_entry(&mut self, entry_key: [u8; 32], response_der: Vec<u8>) {
174        self.add_entry_with_discriminator(entry_key, 0, response_der);
175    }
176
177    pub fn add_entry_with_discriminator(
178        &mut self,
179        entry_key: [u8; 32],
180        discriminator: u16,
181        response_der: Vec<u8>,
182    ) {
183        self.entries.push((
184            IndexRecord {
185                entry_key,
186                data_offset: 0,
187                data_length: response_der.len() as u32,
188                flags: crate::index::IndexFlags::empty(),
189                discriminator,
190            },
191            response_der,
192        ));
193    }
194
195    /// Add a dual-CertID entry: one response payload indexed under two keys.
196    ///
197    /// Both records carry the full payload so that sorting doesn't break
198    /// the pairing. During build, we deduplicate by data content so the
199    /// payload is stored only once in the data section.
200    pub fn add_dual_entry(
201        &mut self,
202        entry_key_1: [u8; 32],
203        entry_key_2: [u8; 32],
204        response_der: Vec<u8>,
205    ) {
206        self.add_dual_entry_with_discriminator(entry_key_1, entry_key_2, response_der, 0);
207    }
208
209    pub fn add_dual_entry_with_discriminator(
210        &mut self,
211        entry_key_1: [u8; 32],
212        entry_key_2: [u8; 32],
213        response_der: Vec<u8>,
214        discriminator: u16,
215    ) {
216        let len = response_der.len() as u32;
217        let flags = crate::index::IndexFlags::ALIAS | crate::index::IndexFlags::MULTI;
218        self.entries.push((
219            IndexRecord {
220                entry_key: entry_key_1,
221                data_offset: 0,
222                data_length: len,
223                flags,
224                discriminator,
225            },
226            response_der.clone(),
227        ));
228        self.entries.push((
229            IndexRecord {
230                entry_key: entry_key_2,
231                data_offset: 0,
232                data_length: len,
233                flags,
234                discriminator,
235            },
236            response_der,
237        ));
238    }
239
240    /// Add a tombstone (delta only). The discriminator must match the
241    /// target record's discriminator to suppress the correct algorithm variant.
242    pub fn add_tombstone(&mut self, entry_key: [u8; 32], discriminator: u16) {
243        self.entries.push((
244            IndexRecord {
245                entry_key,
246                data_offset: 0,
247                data_length: 0,
248                flags: crate::index::IndexFlags::TOMBSTONE,
249                discriminator,
250            },
251            Vec::new(),
252        ));
253    }
254
255    /// Build the bundle. `seal_fn` is called with the manifest bytes and
256    /// must return a CMS SignedData (detached) as DER bytes.
257    pub fn build<F>(mut self, seal_fn: F) -> Result<Vec<u8>>
258    where
259        F: FnOnce(&[u8]) -> Result<Vec<u8>>,
260    {
261        // Sort entries by (key, discriminator), resolving alias data sharing.
262        self.entries
263            .sort_by_key(|a| (a.0.entry_key, a.0.discriminator));
264
265        // Build data section and fix up offsets.
266        // For ALIAS entries, deduplicate identical payloads so the data
267        // is stored once even though two index records point at it.
268        let mut data_section = Vec::new();
269        let mut index_records = Vec::with_capacity(self.entries.len());
270        let mut payload_offsets: std::collections::HashMap<[u8; 32], (u64, u32)> =
271            std::collections::HashMap::new();
272
273        for (mut record, payload) in self.entries {
274            if record.is_tombstone() {
275                record.data_offset = 0;
276                record.data_length = 0;
277                index_records.push(record);
278                continue;
279            }
280
281            if record.is_alias() {
282                let digest: [u8; 32] = Sha256::digest(&payload).into();
283                if let Some(&(offset, len)) = payload_offsets.get(&digest) {
284                    record.data_offset = offset;
285                    record.data_length = len;
286                    index_records.push(record);
287                    continue;
288                }
289                let offset = data_section.len() as u64;
290                let len = payload.len() as u32;
291                data_section.extend_from_slice(&payload);
292                payload_offsets.insert(digest, (offset, len));
293                record.data_offset = offset;
294                record.data_length = len;
295            } else {
296                let offset = data_section.len() as u64;
297                record.data_offset = offset;
298                record.data_length = payload.len() as u32;
299                data_section.extend_from_slice(&payload);
300            }
301
302            index_records.push(record);
303        }
304
305        // Build index section.
306        let mut index_section = Vec::with_capacity(index_records.len() * INDEX_RECORD_SIZE);
307        for record in &index_records {
308            record.write_to(&mut index_section)?;
309        }
310
311        // Compute integrity digests.
312        let index_digest: [u8; 32] = Sha256::digest(&index_section).into();
313        let data_digest: [u8; 32] = Sha256::digest(&data_section).into();
314
315        // Update manifest with computed values.
316        self.manifest.integrity.index_digest = index_digest;
317        self.manifest.integrity.data_digest = data_digest;
318        self.manifest.entry_count = index_records.len() as u64;
319
320        // Encode manifest to deterministic CBOR.
321        let manifest_bytes = self.manifest.to_cbor();
322
323        // Generate seal over manifest.
324        let seal_bytes = seal_fn(&manifest_bytes)?;
325
326        // Compute section layout.
327        let manifest_offset = HEADER_SIZE as u64;
328        let manifest_length = manifest_bytes.len() as u32;
329        let seal_offset = manifest_offset + manifest_length as u64;
330        let seal_length = seal_bytes.len() as u32;
331        let index_offset = seal_offset + seal_length as u64;
332        let index_length = index_section.len() as u64;
333        let data_offset = index_offset + index_length;
334        let data_length = data_section.len() as u64;
335
336        let header = FileHeader {
337            format_major: crate::header::FORMAT_MAJOR,
338            format_minor: crate::header::FORMAT_MINOR,
339            manifest_offset,
340            manifest_length,
341            seal_offset,
342            seal_length,
343            index_offset,
344            index_length,
345            data_offset,
346            data_length,
347        };
348
349        // Write the complete bundle.
350        let total_size = data_offset as usize + data_section.len();
351        let mut output = Vec::with_capacity(total_size);
352        header.write_to(&mut output)?;
353        output.extend_from_slice(&manifest_bytes);
354        output.extend_from_slice(&seal_bytes);
355        output.extend_from_slice(&index_section);
356        output.extend_from_slice(&data_section);
357
358        Ok(output)
359    }
360}
361
362#[cfg(test)]
363mod tests {
364    use super::*;
365    use crate::manifest::*;
366    use uuid::Uuid;
367
368    fn test_manifest() -> Manifest {
369        Manifest {
370            format_version: 1,
371            bundle_id: Uuid::nil(),
372            producer_id: "test".into(),
373            created_at: 1700000000,
374            bundle_type: BundleType::Full,
375            ca_scopes: vec![CaScope {
376                hash_algorithm: vec![0x01],
377                issuer_name_hash: vec![0xAA; 32],
378                issuer_key_hash: vec![0xBB; 32],
379                epoch: 1,
380                responder_id: ResponderId {
381                    id_type: ResponderIdType::ByKey,
382                    value: vec![0xCC; 20],
383                },
384                responder_chain: None,
385                signature_algorithm: vec![0x02],
386                completeness: Completeness::AuthoritativeComplete,
387            }],
388            window: Window {
389                produced_at: 1700000000,
390                this_update_min: 1700000000,
391                next_update_min: 1700086400,
392                next_update_max: 1700093600,
393            },
394            integrity: Integrity {
395                index_digest: [0; 32],
396                data_digest: [0; 32],
397            },
398            entry_count: 0,
399            continuity: Continuity {
400                prev_manifest_digest: None,
401                base_manifest_digest: None,
402                chain_length: 0,
403            },
404            shard: None,
405            compression: None,
406            extensions: None,
407        }
408    }
409
410    #[test]
411    fn build_and_read_empty_bundle() {
412        let manifest = test_manifest();
413        let builder = BundleBuilder::new(manifest);
414
415        let bytes = builder
416            .build(|manifest_bytes| {
417                // Dummy seal: just echo the manifest hash as a "seal"
418                Ok(Sha256::digest(manifest_bytes).to_vec())
419            })
420            .unwrap();
421
422        let bundle = Bundle::from_bytes(&bytes).unwrap();
423        assert_eq!(bundle.index.len(), 0);
424        assert_eq!(bundle.manifest.entry_count, 0);
425        assert_eq!(bundle.manifest.producer_id, "test");
426    }
427
428    #[test]
429    fn build_and_lookup() {
430        let manifest = test_manifest();
431        let mut builder = BundleBuilder::new(manifest);
432
433        let fake_certid = b"fake-certid-for-testing-1234567";
434        let entry_key = crate::index::compute_entry_key(fake_certid);
435        let response = b"fake-ocsp-response-bytes".to_vec();
436
437        builder.add_entry(entry_key, response.clone());
438
439        let bytes = builder.build(|m| Ok(Sha256::digest(m).to_vec())).unwrap();
440
441        let bundle = Bundle::from_bytes(&bytes).unwrap();
442        assert_eq!(bundle.index.len(), 1);
443        assert_eq!(bundle.manifest.entry_count, 1);
444
445        let found = bundle.lookup(&entry_key).expect("entry should be found");
446        assert_eq!(found, &response[..]);
447
448        let missing = [0xFF; 32];
449        assert!(bundle.lookup(&missing).is_none());
450    }
451
452    #[test]
453    fn entries_are_sorted() {
454        let manifest = test_manifest();
455        let mut builder = BundleBuilder::new(manifest);
456
457        // Add in reverse order.
458        let key_b = [0xBB; 32];
459        let key_a = [0xAA; 32];
460
461        builder.add_entry(key_b, b"response-b".to_vec());
462        builder.add_entry(key_a, b"response-a".to_vec());
463
464        let bytes = builder.build(|m| Ok(Sha256::digest(m).to_vec())).unwrap();
465
466        let bundle = Bundle::from_bytes(&bytes).unwrap();
467        assert_eq!(bundle.index[0].entry_key, key_a);
468        assert_eq!(bundle.index[1].entry_key, key_b);
469    }
470}
ahu/bundle.rs

ahu/
bundle.rs
