ahu/
mmap_bundle.rs

1use byteorder::{BigEndian, ByteOrder};
2use memmap2::MmapOptions;
3use std::fs::File;
4use std::path::Path;
5
6use crate::error::{AhuError, Result};
7use crate::header::{FileHeader, HEADER_SIZE};
8use crate::index::{INDEX_RECORD_SIZE, IndexFlags};
9use crate::manifest::Manifest;
10
11/// A memory-mapped ahu bundle for zero-copy serving at scale.
12///
13/// Uses MAP_PRIVATE (copy-on-write). The backing inode must remain immutable:
14/// publish replacements with atomic rename, never truncate an open bundle. The OS demand-pages
15/// only the regions actually accessed — a 45 GB bundle uses ~200 MB RSS
16/// under typical access patterns.
17///
18/// The index is searched via binary search directly in the mapped region.
19/// Response bytes are returned as slices into the mapping — no allocation.
20pub struct MmapBundle {
21    mmap: memmap2::Mmap,
22    pub header: FileHeader,
23    pub manifest: Manifest,
24    manifest_start: usize,
25    manifest_end: usize,
26    seal_start: usize,
27    seal_end: usize,
28    index_offset: usize,
29    index_count: usize,
30    data_offset: usize,
31    data_length: usize,
32}
33
34impl MmapBundle {
35    /// Open a bundle file with MAP_PRIVATE (copy-on-write).
36    ///
37    /// # Safety
38    ///
39    /// Uses `unsafe` for the mmap syscall. Publishers must replace files atomically;
40    /// MAP_PRIVATE does not protect against truncation. Size is validated against header
41    /// offsets before any section is accessed.
42    pub fn open(path: &Path) -> Result<Self> {
43        let file = File::open(path)?;
44        let metadata = file.metadata()?;
45        let file_size = metadata.len();
46
47        if usize::try_from(file_size).is_err() || file_size < HEADER_SIZE as u64 {
48            return Err(AhuError::HeaderOutOfBounds {
49                field: "file",
50                offset: 0,
51                length: HEADER_SIZE as u64,
52                file_size,
53            });
54        }
55
56        // SAFETY: Publishers must keep this inode immutable and replace by rename.
57        // File size is validated before any offset access.
58        let mmap = unsafe {
59            MmapOptions::new()
60                .map_copy_read_only(&file)
61                .map_err(AhuError::Io)?
62        };
63
64        let header = FileHeader::read_from(&mut &mmap[..HEADER_SIZE])?;
65        header.validate_bounds(file_size)?;
66
67        let manifest_start = header.manifest_offset as usize;
68        let manifest_end = manifest_start + header.manifest_length as usize;
69        let manifest = Manifest::from_cbor(&mmap[manifest_start..manifest_end])?;
70
71        let seal_start = header.seal_offset as usize;
72        let seal_end = seal_start + header.seal_length as usize;
73
74        let index_offset = header.index_offset as usize;
75        let index_length = header.index_length as usize;
76
77        if index_length % INDEX_RECORD_SIZE != 0 {
78            return Err(AhuError::IndexSizeMismatch {
79                size: index_length as u64,
80                record_size: INDEX_RECORD_SIZE,
81            });
82        }
83        let index_count = index_length / INDEX_RECORD_SIZE;
84
85        let data_offset = header.data_offset as usize;
86        let data_length = header.data_length as usize;
87
88        Ok(MmapBundle {
89            mmap,
90            header,
91            manifest,
92            manifest_start,
93            manifest_end,
94            seal_start,
95            seal_end,
96            index_offset,
97            index_count,
98            data_offset,
99            data_length,
100        })
101    }
102
103    /// Look up an entry by its entry key (SHA-256 of DER CertID).
104    /// Returns the default (discriminator=0) entry.
105    pub fn lookup(&self, entry_key: &[u8; 32]) -> Option<&[u8]> {
106        let idx = self.binary_search_disc(entry_key, 0)?;
107        self.entry_at_mmap(idx)
108    }
109
110    /// Look up the best-matching entry given preferred discriminators.
111    pub fn lookup_preferred(&self, entry_key: &[u8; 32], preferences: &[u16]) -> Option<&[u8]> {
112        for &disc in preferences {
113            if let Some(idx) = self.binary_search_disc(entry_key, disc) {
114                return self.entry_at_mmap(idx);
115            }
116        }
117        if !preferences.contains(&0) {
118            if let Some(idx) = self.binary_search_disc(entry_key, 0) {
119                return self.entry_at_mmap(idx);
120            }
121        }
122        None
123    }
124
125    fn entry_at_mmap(&self, idx: usize) -> Option<&[u8]> {
126        let (offset, length, flags, _disc) = self.read_record_fields(idx);
127        if flags.contains(IndexFlags::TOMBSTONE) {
128            return None;
129        }
130        let relative = usize::try_from(offset).ok()?;
131        let relative_end = relative.checked_add(usize::try_from(length).ok()?)?;
132        if relative_end > self.data_length {
133            return None;
134        }
135        let start = self.data_offset.checked_add(relative)?;
136        let end = self.data_offset.checked_add(relative_end)?;
137        self.mmap.get(start..end)
138    }
139
140    /// Binary search for `(entry_key, discriminator)` in the mmap'd index.
141    fn binary_search_disc(&self, entry_key: &[u8; 32], disc: u16) -> Option<usize> {
142        let mut lo = 0usize;
143        let mut hi = self.index_count;
144
145        while lo < hi {
146            let mid = lo + (hi - lo) / 2;
147            let key = self.read_entry_key(mid);
148            let mid_disc = self.read_discriminator(mid);
149
150            match key.cmp(entry_key).then(mid_disc.cmp(&disc)) {
151                std::cmp::Ordering::Equal => return Some(mid),
152                std::cmp::Ordering::Less => lo = mid + 1,
153                std::cmp::Ordering::Greater => hi = mid,
154            }
155        }
156        None
157    }
158
159    /// Read the 32-byte entry key at index position `n`.
160    #[inline]
161    fn read_entry_key(&self, n: usize) -> &[u8; 32] {
162        let offset = self.index_offset + n * INDEX_RECORD_SIZE;
163        self.mmap[offset..offset + 32]
164            .try_into()
165            .expect("slice is exactly 32 bytes")
166    }
167
168    /// Read discriminator (u16) at bytes 46-47 of record `n`.
169    #[inline]
170    fn read_discriminator(&self, n: usize) -> u16 {
171        let base = self.index_offset + n * INDEX_RECORD_SIZE;
172        BigEndian::read_u16(&self.mmap[base + 46..base + 48])
173    }
174
175    /// Read data_offset (u64), data_length (u32), flags (u16), discriminator (u16) from record `n`.
176    #[inline]
177    fn read_record_fields(&self, n: usize) -> (u64, u32, IndexFlags, u16) {
178        let base = self.index_offset + n * INDEX_RECORD_SIZE;
179        let data_offset = BigEndian::read_u64(&self.mmap[base + 32..base + 40]);
180        let data_length = BigEndian::read_u32(&self.mmap[base + 40..base + 44]);
181        let flags_raw = BigEndian::read_u16(&self.mmap[base + 44..base + 46]);
182        let flags = IndexFlags::from_bits_truncate(flags_raw);
183        let discriminator = BigEndian::read_u16(&self.mmap[base + 46..base + 48]);
184        (data_offset, data_length, flags, discriminator)
185    }
186
187    /// Get the raw manifest bytes (slice into mmap).
188    pub fn manifest_bytes(&self) -> &[u8] {
189        &self.mmap[self.manifest_start..self.manifest_end]
190    }
191
192    /// Get the raw seal bytes (slice into mmap).
193    pub fn seal_bytes(&self) -> &[u8] {
194        &self.mmap[self.seal_start..self.seal_end]
195    }
196
197    /// Number of index records.
198    pub fn entry_count(&self) -> usize {
199        self.index_count
200    }
201}
202
203#[cfg(test)]
204mod tests {
205    use super::*;
206    use crate::bundle::BundleBuilder;
207    use crate::manifest::*;
208    use sha2::{Digest, Sha256};
209    use uuid::Uuid;
210
211    fn build_test_bundle(n: usize) -> Vec<u8> {
212        let manifest = Manifest {
213            format_version: 1,
214            bundle_id: Uuid::nil(),
215            producer_id: "mmap-test".into(),
216            created_at: 1700000000,
217            bundle_type: BundleType::Full,
218            ca_scopes: vec![CaScope {
219                hash_algorithm: vec![0x01],
220                issuer_name_hash: vec![0xAA; 32],
221                issuer_key_hash: vec![0xBB; 32],
222                epoch: 1,
223                responder_id: ResponderId {
224                    id_type: ResponderIdType::ByKey,
225                    value: vec![0xCC; 20],
226                },
227                responder_chain: None,
228                signature_algorithm: vec![0x02],
229                completeness: Completeness::AuthoritativeComplete,
230            }],
231            window: Window {
232                produced_at: 1700000000,
233                this_update_min: 1700000000,
234                next_update_min: 1700086400,
235                next_update_max: 1700093600,
236            },
237            integrity: Integrity {
238                index_digest: [0; 32],
239                data_digest: [0; 32],
240            },
241            entry_count: 0,
242            continuity: Continuity {
243                prev_manifest_digest: None,
244                base_manifest_digest: None,
245                chain_length: 0,
246            },
247            shard: None,
248            compression: None,
249            extensions: None,
250        };
251
252        let mut builder = BundleBuilder::new(manifest);
253        for i in 0..n {
254            let certid = format!("test-certid-{i:08}");
255            let entry_key = crate::index::compute_entry_key(certid.as_bytes());
256            let response = format!("response-{i:08}").into_bytes();
257            builder.add_entry(entry_key, response);
258        }
259
260        builder.build(|m| Ok(Sha256::digest(m).to_vec())).unwrap()
261    }
262
263    #[test]
264    fn mmap_lookup_matches_heap() {
265        let bytes = build_test_bundle(50);
266
267        let dir = tempfile::tempdir().unwrap();
268        let path = dir.path().join("test.ahu");
269        std::fs::write(&path, &bytes).unwrap();
270
271        let heap_bundle = crate::Bundle::from_bytes(&bytes).unwrap();
272        let mmap_bundle = MmapBundle::open(&path).unwrap();
273
274        assert_eq!(mmap_bundle.entry_count(), heap_bundle.index.len());
275
276        for i in 0..50 {
277            let certid = format!("test-certid-{i:08}");
278            let entry_key = crate::index::compute_entry_key(certid.as_bytes());
279
280            let heap_result = heap_bundle.lookup(&entry_key);
281            let mmap_result = mmap_bundle.lookup(&entry_key);
282
283            assert_eq!(heap_result, mmap_result, "mismatch at entry {i}");
284        }
285    }
286
287    #[test]
288    fn mmap_missing_key_returns_none() {
289        let bytes = build_test_bundle(10);
290
291        let dir = tempfile::tempdir().unwrap();
292        let path = dir.path().join("test.ahu");
293        std::fs::write(&path, &bytes).unwrap();
294
295        let bundle = MmapBundle::open(&path).unwrap();
296        let missing = [0xFF; 32];
297        assert!(bundle.lookup(&missing).is_none());
298    }
299
300    #[test]
301    fn mmap_manifest_and_seal_accessible() {
302        let bytes = build_test_bundle(5);
303
304        let dir = tempfile::tempdir().unwrap();
305        let path = dir.path().join("test.ahu");
306        std::fs::write(&path, &bytes).unwrap();
307
308        let bundle = MmapBundle::open(&path).unwrap();
309        assert!(!bundle.manifest_bytes().is_empty());
310        assert!(!bundle.seal_bytes().is_empty());
311        assert_eq!(bundle.manifest.producer_id, "mmap-test");
312    }
313}
ahu/mmap_bundle.rs

ahu/
mmap_bundle.rs
