1use arrow_schema::{DataType, Field, Fields, Schema};
12use iris_abi::Node;
13
14use crate::check::{MAX_DEPTH, check};
15use crate::error::{Invariant, Result};
16use crate::indirect::{check_dictionary, check_views};
17
18#[derive(Clone, Debug)]
20pub struct Case {
21 pub name: &'static str,
23 pub why: &'static str,
25 pub expected: Option<Invariant>,
27 pub subject: Subject,
29}
30
31#[derive(Clone, Debug)]
33#[non_exhaustive]
34pub enum Subject {
35 Batch {
37 schema: Schema,
39 rows: u64,
41 nodes: Vec<Node>,
43 buffers: Vec<Vec<u8>>,
45 },
46 Dictionary {
48 keys: Vec<u8>,
50 key_type: DataType,
52 len: u64,
54 dictionary_len: u64,
56 },
57 Views {
59 views: Vec<u8>,
61 data: Vec<Vec<u8>>,
63 len: u64,
65 },
66}
67
68impl Subject {
69 pub fn run(&self) -> Result<()> {
75 match self {
76 Self::Batch {
77 schema,
78 rows,
79 nodes,
80 buffers,
81 } => check(schema, *rows, nodes, buffers),
82 Self::Dictionary {
83 keys,
84 key_type,
85 len,
86 dictionary_len,
87 } => check_dictionary(keys, key_type, *len, *dictionary_len, "keys"),
88 Self::Views { views, data, len } => check_views(views, data, *len, "views"),
89 }
90 }
91}
92
93#[must_use]
95pub fn cases() -> Vec<Case> {
96 vec![
97 sound_integers(),
98 sound_strings_with_nulls(),
99 sound_nested_struct(),
100 offset_one_past_the_end(),
101 null_count_off_by_one(),
102 dictionary_index_equal_to_the_dictionary_length(),
103 view_buffer_index_equal_to_the_buffer_count(),
104 length_times_width_that_overflows(),
105 length_that_wraps_the_count_of_offsets(),
106 child_one_row_short_of_its_parent(),
107 schema_nesting_without_a_bound(),
108 ]
109}
110
111fn i64s(values: &[i64]) -> Vec<u8> {
112 values.iter().flat_map(|v| v.to_le_bytes()).collect()
113}
114
115fn i32s(values: &[i32]) -> Vec<u8> {
116 values.iter().flat_map(|v| v.to_le_bytes()).collect()
117}
118
119fn sound_integers() -> Case {
120 Case {
121 name: "three integers",
122 why: "a checker that refuses everything passes every corpus, so the corpus has to include \
123 batches that are fine",
124 expected: None,
125 subject: Subject::Batch {
126 schema: Schema::new(vec![Field::new("a", DataType::Int64, false)]),
127 rows: 3,
128 nodes: vec![Node {
129 length: 3,
130 null_count: 0,
131 }],
132 buffers: vec![Vec::new(), i64s(&[1, 2, 3])],
133 },
134 }
135}
136
137fn sound_strings_with_nulls() -> Case {
138 Case {
139 name: "two strings, one of them null",
140 why: "the validity path and the offsets path are the two that get the most attention here, \
141 so both need a case that passes",
142 expected: None,
143 subject: Subject::Batch {
144 schema: Schema::new(vec![Field::new("s", DataType::Utf8, true)]),
145 rows: 2,
146 nodes: vec![Node {
147 length: 2,
148 null_count: 1,
149 }],
150 buffers: vec![vec![0b0000_0001], i32s(&[0, 2, 2]), b"ho".to_vec()],
152 },
153 }
154}
155
156fn sound_nested_struct() -> Case {
157 let children = Fields::from(vec![
158 Field::new("x", DataType::Int64, false),
159 Field::new("y", DataType::Int64, false),
160 ]);
161 Case {
162 name: "a struct of two integers",
163 why: "nesting is where the buffer counting is easiest to get wrong in either direction",
164 expected: None,
165 subject: Subject::Batch {
166 schema: Schema::new(vec![Field::new("p", DataType::Struct(children), false)]),
167 rows: 2,
168 nodes: vec![
169 Node {
170 length: 2,
171 null_count: 0,
172 },
173 Node {
174 length: 2,
175 null_count: 0,
176 },
177 Node {
178 length: 2,
179 null_count: 0,
180 },
181 ],
182 buffers: vec![
183 Vec::new(),
184 Vec::new(),
185 i64s(&[1, 2]),
186 Vec::new(),
187 i64s(&[3, 4]),
188 ],
189 },
190 }
191}
192
193fn offset_one_past_the_end() -> Case {
194 Case {
195 name: "an offset one past the end of its buffer",
196 why: "the classic off by one. The offsets are ordered and the buffer is nearly long \
197 enough, so nothing about the array looks wrong until something reads the last value",
198 expected: Some(Invariant::OffsetRange),
199 subject: Subject::Batch {
200 schema: Schema::new(vec![Field::new("s", DataType::Utf8, false)]),
201 rows: 2,
202 nodes: vec![Node {
203 length: 2,
204 null_count: 0,
205 }],
206 buffers: vec![Vec::new(), i32s(&[0, 2, 6]), b"hoyea".to_vec()],
207 },
208 }
209}
210
211fn null_count_off_by_one() -> Case {
212 Case {
213 name: "a null count off by one",
214 why: "the one number in a batch that nothing else would catch. An array that lies about \
215 its nulls produces wrong answers rather than an error",
216 expected: Some(Invariant::NullCount),
217 subject: Subject::Batch {
218 schema: Schema::new(vec![Field::new("a", DataType::Int64, true)]),
219 rows: 3,
220 nodes: vec![Node {
221 length: 3,
222 null_count: 2,
223 }],
224 buffers: vec![vec![0b0000_0111], i64s(&[7, 8, 9])],
226 },
227 }
228}
229
230fn dictionary_index_equal_to_the_dictionary_length() -> Case {
231 Case {
232 name: "a dictionary index equal to the dictionary length",
233 why: "in range for the arithmetic and one past the end of the data, which is what an off \
234 by one in a decoder produces",
235 expected: Some(Invariant::DictionaryIndex),
236 subject: Subject::Dictionary {
237 keys: i32s(&[0, 1, 3]),
238 key_type: DataType::Int32,
239 len: 3,
240 dictionary_len: 3,
241 },
242 }
243}
244
245fn view_buffer_index_equal_to_the_buffer_count() -> Case {
246 let mut views = Vec::with_capacity(16);
247 views.extend_from_slice(&18u32.to_le_bytes());
248 views.extend_from_slice(&[0u8; 4]);
249 views.extend_from_slice(&1u32.to_le_bytes());
250 views.extend_from_slice(&0u32.to_le_bytes());
251
252 Case {
253 name: "a view buffer index equal to the buffer count",
254 why: "the same off by one as the dictionary key, in the one array layout where the number \
255 of buffers is not fixed by the schema",
256 expected: Some(Invariant::ViewBuffer),
257 subject: Subject::Views {
258 views,
259 data: vec![b"hello there friend".to_vec()],
260 len: 1,
261 },
262 }
263}
264
265fn length_times_width_that_overflows() -> Case {
266 Case {
267 name: "a length times an element width that overflows",
268 why: "the arithmetic a checker does is itself an attack surface. A length that wraps when \
269 multiplied by a width turns a bounds check into a permission slip",
270 expected: Some(Invariant::Size),
271 subject: Subject::Batch {
272 schema: Schema::new(vec![Field::new("a", DataType::Int64, false)]),
273 rows: u64::MAX,
274 nodes: vec![Node {
275 length: u64::MAX,
276 null_count: 0,
277 }],
278 buffers: vec![Vec::new(), i64s(&[1])],
279 },
280 }
281}
282
283fn length_that_wraps_the_count_of_offsets() -> Case {
284 Case {
285 name: "a length that wraps the count of offsets",
286 why: "found by the fuzzer rather than by anybody thinking about it. There is one more \
287 offset than there are slots, and adding that one to the largest length there is \
288 wrapped the count to zero, so the buffer needed no bytes and the loop over the \
289 offsets ran no times. The array came back sound with nothing in it",
290 expected: Some(Invariant::Size),
291 subject: Subject::Batch {
292 schema: Schema::new(vec![Field::new("s", DataType::Utf8, false)]),
293 rows: u64::MAX,
294 nodes: vec![Node {
295 length: u64::MAX,
296 null_count: 0,
297 }],
298 buffers: vec![Vec::new(), Vec::new(), Vec::new()],
299 },
300 }
301}
302
303fn child_one_row_short_of_its_parent() -> Case {
304 let children = Fields::from(vec![Field::new("x", DataType::Int64, false)]);
305 Case {
306 name: "a child array one row short of its parent",
307 why: "a struct's fields are read by the parent's length, so a short child is read past its \
308 end on the last row and nowhere else",
309 expected: Some(Invariant::ChildLength),
310 subject: Subject::Batch {
311 schema: Schema::new(vec![Field::new("p", DataType::Struct(children), false)]),
312 rows: 3,
313 nodes: vec![
314 Node {
315 length: 3,
316 null_count: 0,
317 },
318 Node {
319 length: 2,
320 null_count: 0,
321 },
322 ],
323 buffers: vec![Vec::new(), Vec::new(), i64s(&[1, 2])],
324 },
325 }
326}
327
328fn schema_nesting_without_a_bound() -> Case {
329 let mut data_type = DataType::Int64;
330 for _ in 0..MAX_DEPTH + 10 {
331 data_type = DataType::List(std::sync::Arc::new(Field::new("item", data_type, false)));
332 }
333
334 Case {
335 name: "a schema nested deeper than anything will walk",
336 why: "everything downstream of the guard walks a schema recursively, so an unbounded \
337 schema is a stack overflow rather than an error, and a stack overflow is not \
338 something a host can turn into a failed query",
339 expected: Some(Invariant::Depth),
340 subject: Subject::Batch {
341 schema: Schema::new(vec![Field::new("deep", data_type, false)]),
342 rows: 0,
343 nodes: Vec::new(),
344 buffers: Vec::<Vec<u8>>::new(),
345 },
346 }
347}