apache_avro/
util.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Utility functions, like configuring various global settings.
19
20use crate::{AvroResult, error::Details, schema::Documentation};
21use log::warn;
22use serde_json::{Map, Value};
23use std::{
24    io::{Read, Write},
25    sync::OnceLock,
26};
27
28/// Maximum number of bytes that can be allocated when decoding Avro-encoded values.
29///
30/// This is a protection against ill-formed data, whose length field might be interpreted as enormous.
31///
32/// See [`max_allocation_bytes`] to change this limit.
33pub const DEFAULT_MAX_ALLOCATION_BYTES: usize = 512 * 1024 * 1024;
34static MAX_ALLOCATION_BYTES: OnceLock<usize> = OnceLock::new();
35
36/// Maximum recursion depth when decoding or resolving Avro-encoded values.
37///
38/// This protects against stack exhaustion aborts from deeply nested data.
39///
40/// See [`max_decode_recursion_depth`] to change this limit.
41pub const DEFAULT_MAX_DECODE_RECURSION_DEPTH: usize = 32;
42static MAX_DECODE_RECURSION_DEPTH: OnceLock<usize> = OnceLock::new();
43
44/// Whether to set serialization & deserialization traits as `human_readable` or not.
45///
46/// See [`set_serde_human_readable`] to change this value.
47pub const DEFAULT_SERDE_HUMAN_READABLE: bool = false;
48/// Whether the serializer and deserializer should indicate to types that the format is human-readable.
49// crate-visible for testing
50pub(crate) static SERDE_HUMAN_READABLE: OnceLock<bool> = OnceLock::new();
51
52pub(crate) trait MapHelper {
53    fn string(&mut self, key: &'static str) -> AvroResult<Option<String>>;
54
55    fn str(&self, key: &'static str) -> AvroResult<Option<&str>>;
56
57    fn name(&mut self) -> AvroResult<String> {
58        self.string("name")?
59            .ok_or_else(|| Details::GetNameField.into())
60    }
61
62    fn name_ref(&self) -> AvroResult<Option<&str>> {
63        self.str("name")
64    }
65
66    fn doc(&mut self) -> AvroResult<Documentation> {
67        self.string("doc")
68    }
69
70    fn aliases(&mut self) -> AvroResult<Option<Vec<String>>>;
71}
72
73impl MapHelper for Map<String, Value> {
74    fn string(&mut self, key: &'static str) -> AvroResult<Option<String>> {
75        match self.remove(key) {
76            Some(Value::String(s)) => Ok(Some(s)),
77            Some(value) => Err(Details::GetStringInvalidType(key, value.description()).into()),
78            None => Ok(None),
79        }
80    }
81
82    fn str(&self, key: &'static str) -> AvroResult<Option<&str>> {
83        match self.get(key) {
84            Some(Value::String(s)) => Ok(Some(s)),
85            Some(value) => Err(Details::GetStringInvalidType(key, value.description()).into()),
86            None => Ok(None),
87        }
88    }
89
90    fn aliases(&mut self) -> AvroResult<Option<Vec<String>>> {
91        match self.remove("aliases") {
92            Some(Value::Array(array)) => array
93                .into_iter()
94                .map(|v| match v {
95                    Value::String(s) => Ok(s),
96                    _ => Err(Details::GetAliasesFieldArrayInvalidType(v.description()).into()),
97                })
98                .collect::<Result<Vec<_>, _>>()
99                .map(Some),
100            Some(value) => Err(Details::GetAliasesFieldInvalidType(value.description()).into()),
101            None => Ok(None),
102        }
103    }
104}
105
106pub(crate) trait JsonValueDescriber {
107    fn description(&self) -> &'static str;
108}
109impl JsonValueDescriber for Value {
110    fn description(&self) -> &'static str {
111        match self {
112            Value::Null => "null",
113            Value::Bool(_) => "bool",
114            Value::Number(_) => "number",
115            Value::String(_) => "string",
116            Value::Array(_) => "array",
117            Value::Object(_) => "object",
118        }
119    }
120}
121
122/// Decode a long from the reader and convert it to a usize.
123pub(crate) fn read_usize<R: Read>(reader: &mut R) -> AvroResult<usize> {
124    let long = zag_i64(reader)?;
125    usize::try_from(long).map_err(|e| Details::ConvertI64ToUsize(e, long).into())
126}
127
128/// Write the number as a zigzagged varint to the writer.
129pub(crate) fn zig_i32<W: Write>(n: i32, buffer: W) -> AvroResult<usize> {
130    zig_i64(n as i64, buffer)
131}
132
133/// Write the number as a zigzagged varint to the writer.
134pub(crate) fn zig_i64<W: Write>(n: i64, writer: W) -> AvroResult<usize> {
135    let zigzagged = ((n << 1) ^ (n >> 63)) as u64;
136    encode_variable(zigzagged, writer)
137}
138
139/// Decode a zigzagged varint from the reader.
140pub(crate) fn zag_i32<R: Read>(reader: &mut R) -> AvroResult<i32> {
141    let i = zag_i64(reader)?;
142    i32::try_from(i).map_err(|e| Details::ZagI32(e, i).into())
143}
144
145/// Decode a zigzagged varint from the reader.
146pub(crate) fn zag_i64<R: Read>(reader: &mut R) -> AvroResult<i64> {
147    let z = decode_variable(reader)?;
148    Ok(if z & 0x1 == 0 {
149        (z >> 1) as i64
150    } else {
151        !(z >> 1) as i64
152    })
153}
154
155/// Write the number as a varint to the writer.
156///
157/// Note: this function does not do zigzag encoding, for that see [`zig_i32`] and [`zig_i64`].
158fn encode_variable<W: Write>(mut zigzagged: u64, mut writer: W) -> AvroResult<usize> {
159    // Ensure the number is little endian for the varint encoding (no-op on LE systems)
160    zigzagged = zigzagged.to_le();
161    // Encode the number as a varint
162    let mut buffer = [0u8; 10];
163    let mut i: usize = 0;
164    loop {
165        if zigzagged <= 0x7F {
166            buffer[i] = (zigzagged & 0x7F) as u8;
167            i += 1;
168            break;
169        } else {
170            buffer[i] = (0x80 | (zigzagged & 0x7F)) as u8;
171            i += 1;
172            zigzagged >>= 7;
173        }
174    }
175    writer
176        .write_all(&buffer[..i])
177        .map_err(Details::WriteBytes)?;
178    Ok(i)
179}
180
181/// Read a varint from the reader.
182///
183/// Note: this function does not do zigzag decoding, for that see [`zag_i32`] and [`zag_i64`].
184fn decode_variable<R: Read>(reader: &mut R) -> AvroResult<u64> {
185    let mut i = 0u64;
186    let mut buf = [0u8; 1];
187
188    let mut j = 0;
189    loop {
190        if j > 9 {
191            // if j * 7 > 64
192            return Err(Details::IntegerOverflow.into());
193        }
194        reader
195            .read_exact(&mut buf[..])
196            .map_err(Details::ReadVariableIntegerBytes)?;
197        i |= (u64::from(buf[0] & 0x7F)) << (j * 7);
198        if (buf[0] >> 7) == 0 {
199            break;
200        } else {
201            j += 1;
202        }
203    }
204
205    Ok(u64::from_le(i))
206}
207
208/// Set the maximum number of bytes that can be allocated when decoding data.
209///
210/// This function only changes the setting once. On subsequent calls the value will stay the same
211/// as the first time it is called. It is automatically called on first allocation and defaults to
212/// [`DEFAULT_MAX_ALLOCATION_BYTES`].
213///
214/// # Returns
215/// The configured maximum, which might be different from what the function was called with if the
216/// value was already set before.
217pub fn max_allocation_bytes(num_bytes: usize) -> usize {
218    *MAX_ALLOCATION_BYTES.get_or_init(|| num_bytes)
219}
220
221/// Set the maximum recursion depth used when decoding or resolving data.
222///
223/// This function only changes the setting once. On subsequent calls the value will stay the same
224/// as the first time it is called. It is automatically called on first decode and defaults to
225/// [`DEFAULT_MAX_DECODE_RECURSION_DEPTH`].
226///
227/// # Returns
228/// The configured maximum, which might be different from what the function was called with if the
229/// value was already set before. In that case a warning is logged.
230pub fn max_decode_recursion_depth(depth: usize) -> usize {
231    let configured = *MAX_DECODE_RECURSION_DEPTH.get_or_init(|| depth);
232    if configured != depth {
233        warn!(
234            "max_decode_recursion_depth({depth}) has no effect: the limit was already set to {configured}"
235        );
236    }
237    configured
238}
239
240/// The effective decode recursion limit, initializing the default if unset.
241pub(crate) fn decode_recursion_limit() -> usize {
242    *MAX_DECODE_RECURSION_DEPTH.get_or_init(|| DEFAULT_MAX_DECODE_RECURSION_DEPTH)
243}
244
245pub(crate) fn safe_len(len: usize) -> AvroResult<usize> {
246    let max_bytes = max_allocation_bytes(DEFAULT_MAX_ALLOCATION_BYTES);
247
248    if len <= max_bytes {
249        Ok(len)
250    } else {
251        Err(Details::MemoryAllocation {
252            desired: Some(len),
253            maximum: max_bytes,
254        }
255        .into())
256    }
257}
258
259/// Bound the cumulative number of elements a collection (array or map) may hold.
260///
261/// This is equivalent to `safe_len(total_items * size_of::<T>)`
262pub(crate) fn safe_collection_len<T>(total_items: usize) -> AvroResult<()> {
263    let max_bytes = max_allocation_bytes(DEFAULT_MAX_ALLOCATION_BYTES);
264    // Use checked_mul (not saturating_mul): saturating to usize::MAX could pass
265    // the check below when max_bytes is configured to usize::MAX, letting the
266    // subsequent reserve() hit a capacity-overflow panic instead of erroring.
267    let desired = total_items
268        .checked_mul(size_of::<T>())
269        .ok_or(Details::IntegerOverflow)?;
270
271    if desired <= max_bytes {
272        Ok(())
273    } else {
274        Err(Details::MemoryAllocation {
275            desired: Some(desired),
276            maximum: max_bytes,
277        }
278        .into())
279    }
280}
281
282/// Set whether the serializer and deserializer should indicate to types that the format is human-readable.
283///
284/// This function only changes the setting once. On subsequent calls the value will stay the same
285/// as the first time it is called. It is automatically called on first allocation and defaults to
286/// [`DEFAULT_SERDE_HUMAN_READABLE`].
287///
288/// *NOTE*: Changing this setting can change the output of [`from_value`](crate::from_value) and the
289/// accepted input of [`to_value`](crate::to_value).
290///
291/// # Returns
292/// The configured human-readable value, which might be different from what the function was called
293/// with if the value was already set before.
294pub fn set_serde_human_readable(human_readable: bool) -> bool {
295    *SERDE_HUMAN_READABLE.get_or_init(|| human_readable)
296}
297
298pub(crate) fn is_human_readable() -> bool {
299    *SERDE_HUMAN_READABLE.get_or_init(|| DEFAULT_SERDE_HUMAN_READABLE)
300}
301
302#[cfg(test)]
303mod tests {
304    use super::*;
305    use apache_avro_test_helper::TestResult;
306    use pretty_assertions::assert_eq;
307
308    #[test]
309    fn test_zigzag() {
310        let mut a = Vec::new();
311        let mut b = Vec::new();
312        zig_i32(42i32, &mut a).unwrap();
313        zig_i64(42i64, &mut b).unwrap();
314        assert_eq!(a, b);
315    }
316
317    #[test]
318    fn test_zig_i64() {
319        let mut s = Vec::new();
320
321        zig_i64(0, &mut s).unwrap();
322        assert_eq!(s, [0]);
323
324        s.clear();
325        zig_i64(-1, &mut s).unwrap();
326        assert_eq!(s, [1]);
327
328        s.clear();
329        zig_i64(1, &mut s).unwrap();
330        assert_eq!(s, [2]);
331
332        s.clear();
333        zig_i64(-64, &mut s).unwrap();
334        assert_eq!(s, [127]);
335
336        s.clear();
337        zig_i64(64, &mut s).unwrap();
338        assert_eq!(s, [128, 1]);
339
340        s.clear();
341        zig_i64(i32::MAX as i64, &mut s).unwrap();
342        assert_eq!(s, [254, 255, 255, 255, 15]);
343
344        s.clear();
345        zig_i64(i32::MAX as i64 + 1, &mut s).unwrap();
346        assert_eq!(s, [128, 128, 128, 128, 16]);
347
348        s.clear();
349        zig_i64(i32::MIN as i64, &mut s).unwrap();
350        assert_eq!(s, [255, 255, 255, 255, 15]);
351
352        s.clear();
353        zig_i64(i32::MIN as i64 - 1, &mut s).unwrap();
354        assert_eq!(s, [129, 128, 128, 128, 16]);
355
356        s.clear();
357        zig_i64(i64::MAX, &mut s).unwrap();
358        assert_eq!(s, [254, 255, 255, 255, 255, 255, 255, 255, 255, 1]);
359
360        s.clear();
361        zig_i64(i64::MIN, &mut s).unwrap();
362        assert_eq!(s, [255, 255, 255, 255, 255, 255, 255, 255, 255, 1]);
363    }
364
365    #[test]
366    fn test_zig_i32() {
367        let mut s = Vec::new();
368        zig_i32(i32::MAX / 2, &mut s).unwrap();
369        assert_eq!(s, [254, 255, 255, 255, 7]);
370
371        s.clear();
372        zig_i32(i32::MIN / 2, &mut s).unwrap();
373        assert_eq!(s, [255, 255, 255, 255, 7]);
374
375        s.clear();
376        zig_i32(-(i32::MIN / 2), &mut s).unwrap();
377        assert_eq!(s, [128, 128, 128, 128, 8]);
378
379        s.clear();
380        zig_i32(i32::MIN / 2 - 1, &mut s).unwrap();
381        assert_eq!(s, [129, 128, 128, 128, 8]);
382
383        s.clear();
384        zig_i32(i32::MAX, &mut s).unwrap();
385        assert_eq!(s, [254, 255, 255, 255, 15]);
386
387        s.clear();
388        zig_i32(i32::MIN, &mut s).unwrap();
389        assert_eq!(s, [255, 255, 255, 255, 15]);
390    }
391
392    #[test]
393    fn test_overflow() {
394        let causes_left_shift_overflow: &[u8] = &[0xe1; 10];
395        assert!(matches!(
396            decode_variable(&mut &*causes_left_shift_overflow)
397                .unwrap_err()
398                .details(),
399            Details::IntegerOverflow
400        ));
401    }
402
403    #[test]
404    fn test_safe_len() -> TestResult {
405        assert_eq!(42usize, safe_len(42usize)?);
406        assert!(safe_len(1024 * 1024 * 1024).is_err());
407
408        Ok(())
409    }
410}
apache_avro/util.rs

apache_avro/
util.rs
