nodedb_vector/quantize/
pq_decode.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Allocation-free validation for persisted product-quantizer codebooks.
4
5use std::mem::size_of;
6
7use crate::error::VectorError;
8
9fn invalid_payload(detail: &str) -> VectorError {
10    VectorError::DeserializationFailed(detail.to_string())
11}
12
13struct Cursor<'a> {
14    bytes: &'a [u8],
15    position: usize,
16}
17
18impl<'a> Cursor<'a> {
19    fn new(bytes: &'a [u8]) -> Self {
20        Self { bytes, position: 0 }
21    }
22
23    fn byte(&mut self) -> Result<u8, VectorError> {
24        let value = self
25            .bytes
26            .get(self.position)
27            .copied()
28            .ok_or_else(|| invalid_payload("truncated PQ payload"))?;
29        self.position += 1;
30        Ok(value)
31    }
32
33    fn exact<const N: usize>(&mut self) -> Result<[u8; N], VectorError> {
34        let end = self
35            .position
36            .checked_add(N)
37            .ok_or_else(|| invalid_payload("PQ payload offset overflow"))?;
38        let value = self
39            .bytes
40            .get(self.position..end)
41            .and_then(|slice| slice.try_into().ok())
42            .ok_or_else(|| invalid_payload("truncated PQ payload"))?;
43        self.position = end;
44        Ok(value)
45    }
46
47    fn array_len(&mut self) -> Result<usize, VectorError> {
48        match self.byte()? {
49            marker @ 0x90..=0x9f => Ok(usize::from(marker & 0x0f)),
50            0xdc => Ok(usize::from(u16::from_be_bytes(self.exact()?))),
51            0xdd => usize::try_from(u32::from_be_bytes(self.exact()?))
52                .map_err(|_| invalid_payload("PQ array length exceeds platform size")),
53            _ => Err(invalid_payload("expected PQ array")),
54        }
55    }
56
57    fn usize(&mut self) -> Result<usize, VectorError> {
58        let value = match self.byte()? {
59            marker @ 0x00..=0x7f => u64::from(marker),
60            0xcc => u64::from(self.byte()?),
61            0xcd => u64::from(u16::from_be_bytes(self.exact()?)),
62            0xce => u64::from(u32::from_be_bytes(self.exact()?)),
63            0xcf => u64::from_be_bytes(self.exact()?),
64            _ => return Err(invalid_payload("expected PQ unsigned integer")),
65        };
66        usize::try_from(value).map_err(|_| invalid_payload("PQ integer exceeds platform size"))
67    }
68
69    fn f32(&mut self) -> Result<(), VectorError> {
70        if self.byte()? != 0xca {
71            return Err(invalid_payload("expected PQ float component"));
72        }
73        let _: [u8; 4] = self.exact()?;
74        Ok(())
75    }
76}
77
78/// Prove every allocation-bearing MessagePack cardinality before the generic
79/// decoder is allowed to reserve its nested vectors.
80pub(crate) fn preflight_pq_payload(
81    payload: &[u8],
82    max_dimension: usize,
83    max_allocation_bytes: usize,
84) -> Result<(), VectorError> {
85    let mut reader = Cursor::new(payload);
86    if reader.array_len()? != 5 {
87        return Err(invalid_payload("invalid PQ field count"));
88    }
89    let metadata = (
90        reader.usize()?,
91        reader.usize()?,
92        reader.usize()?,
93        reader.usize()?,
94    );
95    let observed = read_codebook_shape(&mut reader, max_dimension, max_allocation_bytes)?;
96    if observed != (metadata.1, metadata.2, metadata.3) {
97        return Err(invalid_payload("PQ codebook shape does not match metadata"));
98    }
99    if metadata.0 > max_dimension
100        || metadata
101            .1
102            .checked_mul(metadata.3)
103            .is_none_or(|dimension| dimension != metadata.0)
104    {
105        return Err(invalid_payload("invalid PQ dimensions"));
106    }
107    Ok(())
108}
109
110fn read_codebook_shape(
111    reader: &mut Cursor<'_>,
112    max_dimension: usize,
113    max_allocation_bytes: usize,
114) -> Result<(usize, usize, usize), VectorError> {
115    let outer = reader.array_len()?;
116    if outer == 0 || outer > max_dimension {
117        return Err(invalid_payload("invalid PQ codebook count"));
118    }
119    let mut expected_centroids = None;
120    let mut expected_components = None;
121    let mut total_bytes = outer
122        .checked_mul(size_of::<Vec<Vec<f32>>>())
123        .ok_or_else(|| invalid_payload("PQ codebook allocation overflow"))?;
124
125    for _ in 0..outer {
126        let centroid_count = reader.array_len()?;
127        if centroid_count == 0 || centroid_count > usize::from(u8::MAX) + 1 {
128            return Err(invalid_payload("invalid PQ centroid count"));
129        }
130        if expected_centroids
131            .replace(centroid_count)
132            .is_some_and(|value| value != centroid_count)
133        {
134            return Err(invalid_payload("inconsistent PQ centroid count"));
135        }
136        total_bytes = total_bytes
137            .checked_add(
138                centroid_count
139                    .checked_mul(size_of::<Vec<f32>>())
140                    .ok_or_else(|| invalid_payload("PQ centroid allocation overflow"))?,
141            )
142            .ok_or_else(|| invalid_payload("PQ codebook allocation overflow"))?;
143
144        for _ in 0..centroid_count {
145            let components = reader.array_len()?;
146            if components == 0 || components > max_dimension {
147                return Err(invalid_payload("invalid PQ centroid dimension"));
148            }
149            if expected_components
150                .replace(components)
151                .is_some_and(|value| value != components)
152            {
153                return Err(invalid_payload("inconsistent PQ centroid dimension"));
154            }
155            total_bytes = total_bytes
156                .checked_add(
157                    components
158                        .checked_mul(size_of::<f32>())
159                        .ok_or_else(|| invalid_payload("PQ component allocation overflow"))?,
160                )
161                .ok_or_else(|| invalid_payload("PQ codebook allocation overflow"))?;
162            if total_bytes > max_allocation_bytes {
163                return Err(invalid_payload("PQ codebook allocation exceeds limit"));
164            }
165            for _ in 0..components {
166                reader.f32()?;
167            }
168        }
169    }
170
171    Ok((
172        outer,
173        expected_centroids.ok_or_else(|| invalid_payload("missing PQ centroids"))?,
174        expected_components.ok_or_else(|| invalid_payload("missing PQ components"))?,
175    ))
176}
nodedb_vector/quantize/pq_decode.rs

nodedb_vector/quantize/
pq_decode.rs
