urna_runtime/bm25/
codec.rs

1//! On-disk encoding/decoding for the BM25 section (`0x08`). Term order:
2//! alphabetical (deterministic across builds).
3//!
4//! Payload version 2 bitpacks the integer-heavy parts with the shared
5//! `intpack` codec: the doc lengths, the delta-gapped (sorted) doc ids,
6//! and the term frequencies each become one `intpack` column. The decoded
7//! index is identical to the built one, so scores are unchanged. Version 1
8//! (flat `u32` postings) is still accepted on read. The section is optional
9//! and excluded from content_hash.
10
11use std::collections::HashMap;
12
13use urna_format::encoding::{pack_u64s, unpack_u64s};
14use urna_format::error::UrnaError;
15
16use super::index::{BM25_PAYLOAD_VERSION, Bm25Index, Posting, TermEntry};
17use crate::error::RuntimeError;
18use urna_format::bytes::{le_f32, le_u32};
19
20impl Bm25Index {
21    /// Encode for storage in section `0x08` (v2).
22    pub fn to_bytes(&self) -> Vec<u8> {
23        let mut out = Vec::new();
24        out.extend_from_slice(&BM25_PAYLOAD_VERSION.to_le_bytes());
25        out.extend_from_slice(&self.k1.to_le_bytes());
26        out.extend_from_slice(&self.b.to_le_bytes());
27        out.extend_from_slice(&self.avgdl.to_le_bytes());
28        out.extend_from_slice(&(self.n_docs as u32).to_le_bytes());
29        out.extend_from_slice(&(self.n_terms as u32).to_le_bytes());
30
31        let dls: Vec<u64> = self.doc_lengths.iter().map(|&x| x as u64).collect();
32        write_blob(&mut out, &pack_u64s(&dls));
33
34        // Sorted by term for determinism.
35        let mut terms: Vec<(&String, &TermEntry)> = self.terms.iter().collect();
36        terms.sort_by(|a, b| a.0.cmp(b.0));
37        let mut doc_gaps: Vec<u64> = Vec::new();
38        let mut tfs: Vec<u64> = Vec::new();
39        for (term, entry) in &terms {
40            let bs = term.as_bytes();
41            out.extend_from_slice(&(bs.len() as u32).to_le_bytes());
42            out.extend_from_slice(bs);
43            out.extend_from_slice(&entry.df.to_le_bytes());
44            let mut prev = 0u32;
45            for p in &entry.postings {
46                doc_gaps.push(p.doc.wrapping_sub(prev) as u64);
47                prev = p.doc;
48                tfs.push(p.tf as u64);
49            }
50        }
51        write_blob(&mut out, &pack_u64s(&doc_gaps));
52        write_blob(&mut out, &pack_u64s(&tfs));
53        out
54    }
55
56    pub fn from_bytes(bytes: &[u8]) -> Result<Self, RuntimeError> {
57        let mut cur = ByteCursor::new(bytes);
58        let version = cur.u32()?;
59        let k1 = cur.f32()?;
60        let b = cur.f32()?;
61        let avgdl = cur.f32()?;
62        let n_docs = cur.u32()? as usize;
63        let n_terms = cur.u32()? as usize;
64        let (doc_lengths, terms) = match version {
65            1 => decode_v1(&mut cur, n_docs, n_terms)?,
66            2 => decode_v2(&mut cur, n_docs, n_terms)?,
67            other => {
68                return Err(RuntimeError::Format(UrnaError::UnsupportedSectionVersion {
69                    section_id: urna_format::layout::SECTION_BM25_INDEX,
70                    version: other,
71                }));
72            }
73        };
74        // semantic validation of an untrusted payload (found by the mutation
75        // fuzz harness): a posting whose doc id is outside `doc_lengths`
76        // indexed out of bounds at query time, and non-finite parameters
77        // would turn every score into NaN.
78        if !(k1.is_finite() && b.is_finite() && avgdl.is_finite()) {
79            return Err(malformed("bm25: non-finite k1 / b / avgdl"));
80        }
81        if doc_lengths.len() != n_docs {
82            return Err(malformed("bm25: doc-length count != n_docs"));
83        }
84        for entry in terms.values() {
85            if entry.postings.iter().any(|p| p.doc as usize >= n_docs) {
86                return Err(malformed("bm25: posting doc id >= n_docs"));
87            }
88        }
89        Ok(Self {
90            k1,
91            b,
92            avgdl,
93            n_docs,
94            n_terms,
95            doc_lengths,
96            terms,
97        })
98    }
99}
100
101fn malformed(reason: impl Into<String>) -> RuntimeError {
102    RuntimeError::Format(UrnaError::MalformedSectionPayload {
103        section_id: urna_format::layout::SECTION_BM25_INDEX,
104        reason: reason.into(),
105    })
106}
107
108fn write_blob(out: &mut Vec<u8>, blob: &[u8]) {
109    out.extend_from_slice(&(blob.len() as u32).to_le_bytes());
110    out.extend_from_slice(blob);
111}
112
113/// hostile-input guard: `n_docs`, `n_terms`, and `df` are raw `u32` read from
114/// an untrusted section and would otherwise be passed straight to
115/// `Vec::/HashMap::with_capacity`, letting a crafted header (e.g. `df =
116/// 0xFFFF_FFFF`) force a multi-GiB reservation and abort the process at
117/// `open()`. Cap the capacity HINT and let the collection grow on demand; the
118/// cursor's per-read EOF check bounds the actual work to the bytes present.
119const MAX_PREALLOC: usize = 1 << 20;
120
121type Decoded = (Vec<u32>, HashMap<String, TermEntry>);
122
123/// v1: flat `u32` doc lengths and `(doc, tf)` postings.
124fn decode_v1(cur: &mut ByteCursor, n_docs: usize, n_terms: usize) -> Result<Decoded, RuntimeError> {
125    let mut doc_lengths = Vec::with_capacity(n_docs.min(MAX_PREALLOC));
126    for _ in 0..n_docs {
127        doc_lengths.push(cur.u32()?);
128    }
129    let mut terms: HashMap<String, TermEntry> = HashMap::with_capacity(n_terms.min(MAX_PREALLOC));
130    for _ in 0..n_terms {
131        let term = cur.term()?;
132        let df = cur.u32()?;
133        let mut postings = Vec::with_capacity((df as usize).min(MAX_PREALLOC));
134        for _ in 0..df {
135            let doc = cur.u32()?;
136            let tf = cur.u32()?;
137            postings.push(Posting { doc, tf });
138        }
139        terms.insert(term, TermEntry { df, postings });
140    }
141    Ok((doc_lengths, terms))
142}
143
144/// v2: `intpack` columns for doc lengths, delta-gapped doc ids, and tfs.
145fn decode_v2(cur: &mut ByteCursor, n_docs: usize, n_terms: usize) -> Result<Decoded, RuntimeError> {
146    let dls = cur.intpack_column()?;
147    if dls.len() != n_docs {
148        return Err(malformed("bm25 v2: doc-length column mismatch"));
149    }
150    let doc_lengths: Vec<u32> = dls.iter().map(|&x| x as u32).collect();
151
152    // term dictionary: (term, df) in alphabetical order.
153    let mut dict: Vec<(String, u32)> = Vec::with_capacity(n_terms.min(MAX_PREALLOC));
154    let mut total: usize = 0;
155    for _ in 0..n_terms {
156        let term = cur.term()?;
157        let df = cur.u32()?;
158        total = total.saturating_add(df as usize);
159        dict.push((term, df));
160    }
161    let gaps = cur.intpack_column()?;
162    let tfs = cur.intpack_column()?;
163    if gaps.len() != total || tfs.len() != total {
164        return Err(malformed("bm25 v2: posting column mismatch"));
165    }
166    let mut terms: HashMap<String, TermEntry> = HashMap::with_capacity(n_terms.min(MAX_PREALLOC));
167    let mut k = 0usize;
168    for (term, df) in dict {
169        let mut postings = Vec::with_capacity((df as usize).min(MAX_PREALLOC));
170        let mut prev = 0u32;
171        for _ in 0..df {
172            let doc = prev.wrapping_add(gaps[k] as u32);
173            prev = doc;
174            postings.push(Posting {
175                doc,
176                tf: tfs[k] as u32,
177            });
178            k += 1;
179        }
180        terms.insert(term, TermEntry { df, postings });
181    }
182    Ok((doc_lengths, terms))
183}
184
185struct ByteCursor<'a> {
186    buf: &'a [u8],
187    pos: usize,
188}
189impl<'a> ByteCursor<'a> {
190    fn new(buf: &'a [u8]) -> Self {
191        Self { buf, pos: 0 }
192    }
193    fn u32(&mut self) -> Result<u32, RuntimeError> {
194        Ok(le_u32(self.bytes(4)?)?)
195    }
196    fn f32(&mut self) -> Result<f32, RuntimeError> {
197        Ok(le_f32(self.bytes(4)?)?)
198    }
199    fn bytes(&mut self, n: usize) -> Result<&'a [u8], RuntimeError> {
200        if n > self.buf.len() - self.pos {
201            return Err(malformed(format!("bm25: unexpected EOF (need {})", n)));
202        }
203        let out = &self.buf[self.pos..self.pos + n];
204        self.pos += n;
205        Ok(out)
206    }
207    fn term(&mut self) -> Result<String, RuntimeError> {
208        let len = self.u32()? as usize;
209        let b = self.bytes(len)?;
210        std::str::from_utf8(b)
211            .map(|s| s.to_string())
212            .map_err(|e| malformed(format!("term utf-8: {}", e)))
213    }
214    fn intpack_column(&mut self) -> Result<Vec<u64>, RuntimeError> {
215        let len = self.u32()? as usize;
216        let blob = self.bytes(len)?;
217        unpack_u64s(blob).map_err(RuntimeError::Format)
218    }
219}
220
221#[cfg(test)]
222mod tests {
223    use super::*;
224
225    fn v1_header(n_docs: u32, n_terms: u32) -> Vec<u8> {
226        let mut out = 1u32.to_le_bytes().to_vec(); // version 1
227        out.extend_from_slice(&1.2f32.to_le_bytes()); // k1
228        out.extend_from_slice(&0.75f32.to_le_bytes()); // b
229        out.extend_from_slice(&1.0f32.to_le_bytes()); // avgdl
230        out.extend_from_slice(&n_docs.to_le_bytes());
231        out.extend_from_slice(&n_terms.to_le_bytes());
232        out
233    }
234
235    // a header claiming ~4.3e9 docs must not reserve ~17 GiB before the read
236    // loop hits EOF; it must return a typed error, never OOM/abort.
237    #[test]
238    fn v1_hostile_n_docs_errors_without_ooming() {
239        let p = v1_header(u32::MAX, 0); // no doc-length bytes follow
240        assert!(Bm25Index::from_bytes(&p).is_err());
241    }
242
243    // a term declaring df = u32::MAX must not reserve ~34 GiB of postings.
244    #[test]
245    fn v1_hostile_df_errors_without_ooming() {
246        let mut p = v1_header(0, 1);
247        p.extend_from_slice(&1u32.to_le_bytes()); // term length 1
248        p.extend_from_slice(b"a"); // term
249        p.extend_from_slice(&u32::MAX.to_le_bytes()); // df = u32::MAX
250        // no posting bytes follow
251        assert!(Bm25Index::from_bytes(&p).is_err());
252    }
253}
urna_runtime/bm25/codec.rs

urna_runtime/bm25/
codec.rs
