hermes_core/structures/vector/
mod.rs1pub mod index;
29pub mod ivf;
30pub mod quantization;
31
32#[cfg(target_pointer_width = "64")]
40pub(crate) const MAX_DENSE_ANN_DECODE_BYTES: usize = 4 * 1024 * 1024 * 1024;
41#[cfg(not(target_pointer_width = "64"))]
42pub(crate) const MAX_DENSE_ANN_DECODE_BYTES: usize = 512 * 1024 * 1024;
43
44const ANN_DECODE_EXPANSION_FACTOR: usize = 8;
45const ANN_DECODE_FIXED_HEADROOM: usize = 16 * 1024 * 1024;
46
47fn decode_ann_with_limit<T: serde::de::DeserializeOwned, const LIMIT: usize>(
48 data: &[u8],
49) -> std::io::Result<(T, usize)> {
50 bincode::serde::decode_from_slice(data, bincode::config::standard().with_limit::<LIMIT>())
51 .map_err(|error| std::io::Error::new(std::io::ErrorKind::InvalidData, error))
52}
53
54#[cfg(target_pointer_width = "64")]
55fn decode_ann_with_relative_limit<T: serde::de::DeserializeOwned>(
56 data: &[u8],
57 budget: usize,
58) -> std::io::Result<(T, usize)> {
59 const MIB: usize = 1024 * 1024;
60 const GIB: usize = 1024 * MIB;
61 if budget <= 32 * MIB {
62 decode_ann_with_limit::<T, { 32 * MIB }>(data)
63 } else if budget <= 128 * MIB {
64 decode_ann_with_limit::<T, { 128 * MIB }>(data)
65 } else if budget <= 512 * MIB {
66 decode_ann_with_limit::<T, { 512 * MIB }>(data)
67 } else if budget <= GIB {
68 decode_ann_with_limit::<T, GIB>(data)
69 } else if budget <= 2 * GIB {
70 decode_ann_with_limit::<T, { 2 * GIB }>(data)
71 } else {
72 decode_ann_with_limit::<T, { 4 * GIB }>(data)
73 }
74}
75
76#[cfg(not(target_pointer_width = "64"))]
77fn decode_ann_with_relative_limit<T: serde::de::DeserializeOwned>(
78 data: &[u8],
79 budget: usize,
80) -> std::io::Result<(T, usize)> {
81 const MIB: usize = 1024 * 1024;
82 if budget <= 32 * MIB {
83 decode_ann_with_limit::<T, { 32 * MIB }>(data)
84 } else if budget <= 128 * MIB {
85 decode_ann_with_limit::<T, { 128 * MIB }>(data)
86 } else {
87 decode_ann_with_limit::<T, { 512 * MIB }>(data)
88 }
89}
90
91pub(crate) fn decode_ann_bincode_exact<T: serde::de::DeserializeOwned>(
94 data: &[u8],
95 description: &str,
96) -> std::io::Result<T> {
97 if data.len() > MAX_DENSE_ANN_DECODE_BYTES {
98 return Err(std::io::Error::new(
99 std::io::ErrorKind::InvalidData,
100 format!(
101 "{description} payload is {} bytes, exceeding the {}-byte decode limit",
102 data.len(),
103 MAX_DENSE_ANN_DECODE_BYTES
104 ),
105 ));
106 }
107 let budget = data
108 .len()
109 .checked_mul(ANN_DECODE_EXPANSION_FACTOR)
110 .and_then(|bytes| bytes.checked_add(ANN_DECODE_FIXED_HEADROOM))
111 .unwrap_or(MAX_DENSE_ANN_DECODE_BYTES)
112 .min(MAX_DENSE_ANN_DECODE_BYTES);
113 let (value, consumed) = decode_ann_with_relative_limit(data, budget)?;
114 if consumed != data.len() {
115 return Err(std::io::Error::new(
116 std::io::ErrorKind::InvalidData,
117 format!(
118 "{description} payload contains {} trailing bytes",
119 data.len() - consumed
120 ),
121 ));
122 }
123 Ok(value)
124}
125
126pub use ivf::{
128 ClusterData, ClusterStorage, CoarseCentroids, CoarseConfig, MultiAssignment, QuantizedCode,
129 SoarConfig,
130};
131
132pub use quantization::{
134 DistanceTable, PQCodebook, PQConfig, PQVector, QuantizedQuery, QuantizedVector, Quantizer,
135 RaBitQCodebook, RaBitQConfig,
136};
137
138pub use index::{
140 BinaryIvfConfig, BinaryIvfIndex, IVFPQConfig, IVFPQIndex, IVFRaBitQConfig, IVFRaBitQIndex,
141 RaBitQIndex,
142};
143
144#[cfg(test)]
145mod decode_tests {
146 use super::decode_ann_bincode_exact;
147
148 #[test]
149 fn ann_bincode_decode_is_exact_and_rejects_large_claims_from_tiny_payloads() {
150 let encoded =
151 bincode::serde::encode_to_vec(vec![1u8, 2, 3], bincode::config::standard()).unwrap();
152 assert_eq!(
153 decode_ann_bincode_exact::<Vec<u8>>(&encoded, "test").unwrap(),
154 [1, 2, 3]
155 );
156
157 let mut trailing = encoded;
158 trailing.push(0);
159 assert!(decode_ann_bincode_exact::<Vec<u8>>(&trailing, "test").is_err());
160
161 let oversized_claim =
165 bincode::serde::encode_to_vec(64usize * 1024 * 1024, bincode::config::standard())
166 .unwrap();
167 assert!(decode_ann_bincode_exact::<Vec<u8>>(&oversized_claim, "test").is_err());
168 }
169}