open_geocode/
container.rs

1//! Single-file Pack container.
2//!
3//! A Pack is one file: a short header, named sections, and a table of contents
4//! at the end. Each section carries its own format version so one part of the
5//! Pack can evolve without touching the others. One file is one thing to
6//! upload, copy or swap, and readers memory-map it once and hand out borrowed
7//! slices per section. Every section carries an XXH3 checksum, so a copy
8//! damaged in transit is caught by [`Container::verify`] before it is served.
9//!
10//! ```text
11//! header   magic "OGPACK01" | container version u32 | reserved u32
12//! sections each starts on a 64-byte boundary
13//! toc      varint count, then per section: name, version, offset, length, xxh3
14//! trailer  toc offset u64 | toc length u64 | magic "OGPACKTC"
15//! ```
16
17use std::{
18    collections::BTreeMap,
19    fs::File,
20    io::{self, BufWriter, Seek, Write},
21    ops::Deref,
22    path::{Path, PathBuf},
23    sync::Arc,
24};
25
26use anyhow::{Context, Result, bail};
27use memmap2::{Mmap, MmapOptions};
28pub use tantivy::directory::OwnedBytes as Bytes;
29
30use rayon::prelude::*;
31use xxhash_rust::xxh3::{Xxh3, xxh3_64};
32
33use crate::util::codec::{get_string, get_u32, get_u64, put_str, put_u64, read_u64_le};
34
35const MAGIC: &[u8; 8] = b"OGPACK01";
36const TRAILER_MAGIC: &[u8; 8] = b"OGPACKTC";
37const CONTAINER_VERSION: u32 = 2;
38const HEADER_BYTES: u64 = 16;
39const TRAILER_BYTES: usize = 24;
40const SECTION_ALIGN: u64 = 64;
41const WRITE_BUFFER_BYTES: usize = 4 << 20;
42
43#[derive(Debug, Clone, PartialEq, Eq)]
44pub struct SectionInfo {
45    pub version: u32,
46    pub offset: u64,
47    pub len: u64,
48    /// XXH3-64 of the section bytes.
49    pub checksum: u64,
50}
51
52pub struct ContainerWriter {
53    path: PathBuf,
54    file: BufWriter<File>,
55    offset: u64,
56    sections: BTreeMap<String, SectionInfo>,
57    open: Option<OpenSection>,
58}
59
60struct OpenSection {
61    name: String,
62    version: u32,
63    start: u64,
64    hasher: Xxh3,
65}
66
67impl ContainerWriter {
68    pub fn create(path: impl AsRef<Path>) -> Result<Self> {
69        let path = path.as_ref().to_path_buf();
70        let file =
71            File::create(&path).with_context(|| format!("failed to create {}", path.display()))?;
72        let mut file = BufWriter::with_capacity(WRITE_BUFFER_BYTES, file);
73        file.write_all(MAGIC)?;
74        file.write_all(&CONTAINER_VERSION.to_le_bytes())?;
75        file.write_all(&0u32.to_le_bytes())?;
76        Ok(Self {
77            path,
78            file,
79            offset: HEADER_BYTES,
80            sections: BTreeMap::new(),
81            open: None,
82        })
83    }
84
85    /// Start a section. Bytes written until [`Self::end`] belong to it.
86    pub fn begin(&mut self, name: &str, version: u32) -> Result<()> {
87        if self.open.is_some() {
88            bail!("section {name} started while another section is open");
89        }
90        if self.sections.contains_key(name) {
91            bail!("duplicate Pack section {name}");
92        }
93        let padding = (SECTION_ALIGN - self.offset % SECTION_ALIGN) % SECTION_ALIGN;
94        self.file.write_all(&vec![0u8; padding as usize])?;
95        self.offset += padding;
96        self.open = Some(OpenSection {
97            name: name.to_string(),
98            version,
99            start: self.offset,
100            hasher: Xxh3::new(),
101        });
102        Ok(())
103    }
104
105    pub fn end(&mut self) -> Result<()> {
106        let section = self.open.take().context("no open Pack section")?;
107        self.sections.insert(
108            section.name,
109            SectionInfo {
110                version: section.version,
111                offset: section.start,
112                len: self.offset - section.start,
113                checksum: section.hasher.digest(),
114            },
115        );
116        Ok(())
117    }
118
119    pub fn add(&mut self, name: &str, version: u32, bytes: &[u8]) -> Result<()> {
120        self.begin(name, version)?;
121        self.write_all(bytes)?;
122        self.end()
123    }
124
125    pub fn add_file(&mut self, name: &str, version: u32, path: &Path) -> Result<()> {
126        self.begin(name, version)?;
127        let mut source =
128            File::open(path).with_context(|| format!("failed to open {}", path.display()))?;
129        io::copy(&mut source, self)
130            .with_context(|| format!("failed to copy {} into the Pack", path.display()))?;
131        self.end()
132    }
133
134    /// Write the table of contents and sync the file. Returns the Pack size.
135    pub fn finish(mut self) -> Result<u64> {
136        if self.open.is_some() {
137            bail!("Pack finished with an open section");
138        }
139        let mut toc = Vec::new();
140        put_u64(&mut toc, self.sections.len() as u64);
141        for (name, info) in &self.sections {
142            put_str(&mut toc, name);
143            put_u64(&mut toc, u64::from(info.version));
144            put_u64(&mut toc, info.offset);
145            put_u64(&mut toc, info.len);
146            put_u64(&mut toc, info.checksum);
147        }
148        let toc_offset = self.offset;
149        self.file.write_all(&toc)?;
150        self.file.write_all(&toc_offset.to_le_bytes())?;
151        self.file.write_all(&(toc.len() as u64).to_le_bytes())?;
152        self.file.write_all(TRAILER_MAGIC)?;
153        self.file.flush()?;
154        let file = self.file.get_mut();
155        file.sync_all()
156            .with_context(|| format!("failed to sync {}", self.path.display()))?;
157        Ok(file.stream_position()?)
158    }
159}
160
161impl Write for ContainerWriter {
162    fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
163        let Some(section) = &mut self.open else {
164            return Err(io::Error::other("Pack bytes written outside a section"));
165        };
166        let written = self.file.write(bytes)?;
167        section.hasher.update(&bytes[..written]);
168        self.offset += written as u64;
169        Ok(written)
170    }
171
172    fn flush(&mut self) -> io::Result<()> {
173        self.file.flush()
174    }
175}
176
177/// Read-only view of a Pack file.
178#[derive(Clone)]
179pub struct Container {
180    bytes: SharedMmap,
181    sections: Arc<BTreeMap<String, SectionInfo>>,
182}
183
184impl std::fmt::Debug for Container {
185    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
186        formatter
187            .debug_struct("Container")
188            .field("bytes", &self.bytes.len())
189            .field("sections", &self.sections.len())
190            .finish()
191    }
192}
193
194impl Container {
195    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
196        let path = path.as_ref();
197        let file =
198            File::open(path).with_context(|| format!("failed to open {}", path.display()))?;
199        // SAFETY: Pack files are immutable once published and mapped read-only.
200        let mmap = unsafe { MmapOptions::new().map(&file) }
201            .with_context(|| format!("failed to mmap {}", path.display()))?;
202        let sections =
203            parse_toc(&mmap).with_context(|| format!("invalid Pack {}", path.display()))?;
204        Ok(Self {
205            bytes: SharedMmap(Arc::new(mmap)),
206            sections: Arc::new(sections),
207        })
208    }
209
210    /// Size of the Pack file in bytes.
211    pub fn file_size(&self) -> u64 {
212        self.bytes.len() as u64
213    }
214
215    pub fn sections(&self) -> &BTreeMap<String, SectionInfo> {
216        &self.sections
217    }
218
219    /// Check every section against its checksum. Reads the whole file, so it
220    /// runs when a Pack is published or copied, not on every open.
221    pub fn verify(&self) -> Result<()> {
222        self.sections.par_iter().try_for_each(|(name, info)| {
223            let bytes = self.bytes_of(info);
224            if xxh3_64(&bytes) != info.checksum {
225                bail!("Pack section {name} is corrupt: its checksum does not match");
226            }
227            Ok(())
228        })
229    }
230
231    /// Bytes of a section whose version must match `version`. The handle keeps
232    /// the mapping alive, so it can be stored next to the `Container`.
233    pub fn section(&self, name: &str, version: u32) -> Result<Bytes> {
234        let info = self.info(name)?;
235        if info.version != version {
236            bail!(
237                "Pack section {name} has version {}, this build reads version {version}; rebuild the Pack",
238                info.version
239            );
240        }
241        Ok(self.bytes_of(info))
242    }
243
244    /// Bytes of a section whose format is owned by another library (Tantivy).
245    pub fn raw_section(&self, name: &str) -> Result<Bytes> {
246        Ok(self.bytes_of(self.info(name)?))
247    }
248
249    fn info(&self, name: &str) -> Result<&SectionInfo> {
250        self.sections
251            .get(name)
252            .with_context(|| format!("Pack is missing section {name}"))
253    }
254
255    fn bytes_of(&self, info: &SectionInfo) -> Bytes {
256        // Bounds were validated when the table of contents was parsed.
257        Bytes::new(self.bytes.clone())
258            .slice(info.offset as usize..(info.offset + info.len) as usize)
259    }
260}
261
262fn parse_toc(bytes: &[u8]) -> Result<BTreeMap<String, SectionInfo>> {
263    if bytes.len() < HEADER_BYTES as usize + TRAILER_BYTES || &bytes[0..8] != MAGIC {
264        bail!("not an open-geocode Pack file");
265    }
266    let version = u32::from_le_bytes(bytes[8..12].try_into().expect("header slice"));
267    if version != CONTAINER_VERSION {
268        bail!("Pack container version {version} is unsupported; rebuild the Pack");
269    }
270    let trailer = bytes.len() - TRAILER_BYTES;
271    if &bytes[trailer + 16..] != TRAILER_MAGIC {
272        bail!("Pack file is truncated or still being written");
273    }
274    let toc_offset = read_u64_le(bytes, trailer).expect("trailer slice") as usize;
275    let toc_len = read_u64_le(bytes, trailer + 8).expect("trailer slice") as usize;
276    let toc_end = toc_offset
277        .checked_add(toc_len)
278        .filter(|end| *end == trailer)
279        .context("Pack table of contents is out of bounds")?;
280    let mut input = &bytes[toc_offset..toc_end];
281    let count = get_u64(&mut input)?;
282    let mut sections = BTreeMap::new();
283    for _ in 0..count {
284        let name = get_string(&mut input)?;
285        let info = SectionInfo {
286            version: get_u32(&mut input)?,
287            offset: get_u64(&mut input)?,
288            len: get_u64(&mut input)?,
289            checksum: get_u64(&mut input)?,
290        };
291        let end = info
292            .offset
293            .checked_add(info.len)
294            .context("Pack section bounds overflow")?;
295        if info.offset < HEADER_BYTES || end > toc_offset as u64 {
296            bail!("Pack section {name} is out of bounds");
297        }
298        sections.insert(name, info);
299    }
300    if !input.is_empty() {
301        bail!("Pack table of contents has trailing bytes");
302    }
303    Ok(sections)
304}
305
306#[derive(Clone)]
307struct SharedMmap(Arc<Mmap>);
308
309impl Deref for SharedMmap {
310    type Target = [u8];
311
312    fn deref(&self) -> &[u8] {
313        &self.0
314    }
315}
316
317// SAFETY: the bytes live in the mapping owned by the `Arc`, so their address
318// does not change when the handle moves or is cloned.
319unsafe impl stable_deref_trait::StableDeref for SharedMmap {}
320
321#[cfg(test)]
322mod tests {
323    use super::*;
324
325    fn temp_file(name: &str) -> PathBuf {
326        std::env::temp_dir().join(format!(
327            "open-geocode-container-{name}-{}",
328            uuid::Uuid::new_v4()
329        ))
330    }
331
332    #[test]
333    fn round_trips_aligned_sections() {
334        let path = temp_file("round-trip");
335        let source = temp_file("source");
336        std::fs::write(&source, b"from a file").expect("source");
337        let mut writer = ContainerWriter::create(&path).expect("writer");
338        writer.add("alpha", 3, b"abc").expect("alpha");
339        writer.begin("beta", 1).expect("begin");
340        writer.write_all(b"streamed ").expect("write");
341        writer.write_all(b"bytes").expect("write");
342        writer.end().expect("end");
343        writer.add_file("gamma", 1, &source).expect("gamma");
344        writer.add("empty", 1, b"").expect("empty");
345        let size = writer.finish().expect("finish");
346
347        let container = Container::open(&path).expect("open");
348        assert_eq!(container.file_size(), size);
349        assert_eq!(
350            container.section("alpha", 3).expect("alpha").as_slice(),
351            b"abc"
352        );
353        assert_eq!(
354            container.section("beta", 1).expect("beta").as_slice(),
355            b"streamed bytes"
356        );
357        assert_eq!(
358            container.section("gamma", 1).expect("gamma").as_slice(),
359            b"from a file"
360        );
361        assert_eq!(
362            container.section("empty", 1).expect("empty").as_slice(),
363            b""
364        );
365        assert_eq!(container.sections()["beta"].offset % SECTION_ALIGN, 0);
366        assert_eq!(
367            container.raw_section("gamma").expect("raw").as_slice(),
368            b"from a file"
369        );
370        let error = container.section("alpha", 2).unwrap_err().to_string();
371        assert!(error.contains("rebuild the Pack"), "{error}");
372        assert!(container.section("missing", 1).is_err());
373    }
374
375    #[test]
376    fn verify_catches_a_flipped_byte_inside_a_section() {
377        let path = temp_file("checksum");
378        let mut writer = ContainerWriter::create(&path).expect("writer");
379        writer.add("alpha", 1, &[7u8; 4096]).expect("alpha");
380        writer.add("beta", 1, b"untouched").expect("beta");
381        writer.finish().expect("finish");
382        Container::open(&path)
383            .expect("open")
384            .verify()
385            .expect("intact");
386
387        let offset = Container::open(&path).expect("open").sections()["alpha"].offset as usize;
388        let mut bytes = std::fs::read(&path).expect("read");
389        bytes[offset + 2000] ^= 1;
390        std::fs::write(&path, &bytes).expect("corrupt");
391        // Structure is intact, so the file still opens; the checksum catches it.
392        let container = Container::open(&path).expect("open");
393        let error = container.verify().unwrap_err().to_string();
394        assert!(error.contains("alpha"), "{error}");
395    }
396
397    #[test]
398    fn rejects_truncated_files() {
399        let path = temp_file("truncated");
400        let mut writer = ContainerWriter::create(&path).expect("writer");
401        writer.add("alpha", 1, b"abcdef").expect("alpha");
402        writer.finish().expect("finish");
403        let bytes = std::fs::read(&path).expect("read");
404        std::fs::write(&path, &bytes[..bytes.len() - 3]).expect("truncate");
405        assert!(Container::open(&path).is_err());
406    }
407}
open_geocode/container.rs

open_geocode/
container.rs
