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 serde_json::{Map, Value};
22use std::{
23    io::{Read, Write},
24    sync::OnceLock,
25};
26
27/// Maximum number of bytes that can be allocated when decoding Avro-encoded values.
28///
29/// This is a protection against ill-formed data, whose length field might be interpreted as enormous.
30///
31/// See [`max_allocation_bytes`] to change this limit.
32pub const DEFAULT_MAX_ALLOCATION_BYTES: usize = 512 * 1024 * 1024;
33static MAX_ALLOCATION_BYTES: OnceLock<usize> = OnceLock::new();
34
35/// Whether to set serialization & deserialization traits as `human_readable` or not.
36///
37/// See [`set_serde_human_readable`] to change this value.
38pub const DEFAULT_SERDE_HUMAN_READABLE: bool = false;
39/// Whether the serializer and deserializer should indicate to types that the format is human-readable.
40// crate-visible for testing
41pub(crate) static SERDE_HUMAN_READABLE: OnceLock<bool> = OnceLock::new();
42
43pub(crate) trait MapHelper {
44    fn string(&self, key: &str) -> Option<&str>;
45
46    fn name(&self) -> Option<&str> {
47        self.string("name")
48    }
49
50    fn doc(&self) -> Documentation {
51        self.string("doc").map(Into::into)
52    }
53
54    fn aliases(&self) -> Option<Vec<String>>;
55}
56
57impl MapHelper for Map<String, Value> {
58    fn string(&self, key: &str) -> Option<&str> {
59        self.get(key).and_then(|v| v.as_str())
60    }
61
62    fn aliases(&self) -> Option<Vec<String>> {
63        // FIXME no warning when aliases aren't a json array of json strings
64        self.get("aliases")
65            .and_then(|aliases| aliases.as_array())
66            .and_then(|aliases| {
67                aliases
68                    .iter()
69                    .map(|alias| alias.as_str())
70                    .map(|alias| alias.map(|a| a.to_string()))
71                    .collect::<Option<_>>()
72            })
73    }
74}
75
76/// Decode a long from the reader and convert it to a usize.
77pub(crate) fn read_usize<R: Read>(reader: &mut R) -> AvroResult<usize> {
78    let long = zag_i64(reader)?;
79    usize::try_from(long).map_err(|e| Details::ConvertI64ToUsize(e, long).into())
80}
81
82/// Write the number as a zigzagged varint to the writer.
83pub(crate) fn zig_i32<W: Write>(n: i32, buffer: W) -> AvroResult<usize> {
84    zig_i64(n as i64, buffer)
85}
86
87/// Write the number as a zigzagged varint to the writer.
88pub(crate) fn zig_i64<W: Write>(n: i64, writer: W) -> AvroResult<usize> {
89    let zigzagged = ((n << 1) ^ (n >> 63)) as u64;
90    encode_variable(zigzagged, writer)
91}
92
93/// Decode a zigzagged varint from the reader.
94pub(crate) fn zag_i32<R: Read>(reader: &mut R) -> AvroResult<i32> {
95    let i = zag_i64(reader)?;
96    i32::try_from(i).map_err(|e| Details::ZagI32(e, i).into())
97}
98
99/// Decode a zigzagged varint from the reader.
100pub(crate) fn zag_i64<R: Read>(reader: &mut R) -> AvroResult<i64> {
101    let z = decode_variable(reader)?;
102    Ok(if z & 0x1 == 0 {
103        (z >> 1) as i64
104    } else {
105        !(z >> 1) as i64
106    })
107}
108
109/// Write the number as a varint to the writer.
110///
111/// Note: this function does not do zigzag encoding, for that see [`zig_i32`] and [`zig_i64`].
112fn encode_variable<W: Write>(mut zigzagged: u64, mut writer: W) -> AvroResult<usize> {
113    // Ensure the number is little endian for the varint encoding (no-op on LE systems)
114    zigzagged = zigzagged.to_le();
115    // Encode the number as a varint
116    let mut buffer = [0u8; 10];
117    let mut i: usize = 0;
118    loop {
119        if zigzagged <= 0x7F {
120            buffer[i] = (zigzagged & 0x7F) as u8;
121            i += 1;
122            break;
123        } else {
124            buffer[i] = (0x80 | (zigzagged & 0x7F)) as u8;
125            i += 1;
126            zigzagged >>= 7;
127        }
128    }
129    writer
130        .write_all(&buffer[..i])
131        .map_err(Details::WriteBytes)?;
132    Ok(i)
133}
134
135/// Read a varint from the reader.
136///
137/// Note: this function does not do zigzag decoding, for that see [`zag_i32`] and [`zag_i64`].
138fn decode_variable<R: Read>(reader: &mut R) -> AvroResult<u64> {
139    let mut i = 0u64;
140    let mut buf = [0u8; 1];
141
142    let mut j = 0;
143    loop {
144        if j > 9 {
145            // if j * 7 > 64
146            return Err(Details::IntegerOverflow.into());
147        }
148        reader
149            .read_exact(&mut buf[..])
150            .map_err(Details::ReadVariableIntegerBytes)?;
151        i |= (u64::from(buf[0] & 0x7F)) << (j * 7);
152        if (buf[0] >> 7) == 0 {
153            break;
154        } else {
155            j += 1;
156        }
157    }
158
159    Ok(u64::from_le(i))
160}
161
162/// Set the maximum number of bytes that can be allocated when decoding data.
163///
164/// This function only changes the setting once. On subsequent calls the value will stay the same
165/// as the first time it is called. It is automatically called on first allocation and defaults to
166/// [`DEFAULT_MAX_ALLOCATION_BYTES`].
167///
168/// # Returns
169/// The configured maximum, which might be different from what the function was called with if the
170/// value was already set before.
171pub fn max_allocation_bytes(num_bytes: usize) -> usize {
172    *MAX_ALLOCATION_BYTES.get_or_init(|| num_bytes)
173}
174
175pub(crate) fn safe_len(len: usize) -> AvroResult<usize> {
176    let max_bytes = max_allocation_bytes(DEFAULT_MAX_ALLOCATION_BYTES);
177
178    if len <= max_bytes {
179        Ok(len)
180    } else {
181        Err(Details::MemoryAllocation {
182            desired: Some(len),
183            maximum: max_bytes,
184        }
185        .into())
186    }
187}
188
189/// Bound the cumulative number of elements a collection (array or map) may hold.
190///
191/// This is equivalent to `safe_len(total_items * size_of::<T>)`
192pub(crate) fn safe_collection_len<T>(total_items: usize) -> AvroResult<()> {
193    let max_bytes = max_allocation_bytes(DEFAULT_MAX_ALLOCATION_BYTES);
194    // Use checked_mul (not saturating_mul): saturating to usize::MAX could pass
195    // the check below when max_bytes is configured to usize::MAX, letting the
196    // subsequent reserve() hit a capacity-overflow panic instead of erroring.
197    let desired = total_items
198        .checked_mul(size_of::<T>())
199        .ok_or(Details::IntegerOverflow)?;
200
201    if desired <= max_bytes {
202        Ok(())
203    } else {
204        Err(Details::MemoryAllocation {
205            desired: Some(desired),
206            maximum: max_bytes,
207        }
208        .into())
209    }
210}
211
212/// Set whether the serializer and deserializer should indicate to types that the format is human-readable.
213///
214/// This function only changes the setting once. On subsequent calls the value will stay the same
215/// as the first time it is called. It is automatically called on first allocation and defaults to
216/// [`DEFAULT_SERDE_HUMAN_READABLE`].
217///
218/// *NOTE*: Changing this setting can change the output of [`from_value`](crate::from_value) and the
219/// accepted input of [`to_value`](crate::to_value).
220///
221/// # Returns
222/// The configured human-readable value, which might be different from what the function was called
223/// with if the value was already set before.
224pub fn set_serde_human_readable(human_readable: bool) -> bool {
225    *SERDE_HUMAN_READABLE.get_or_init(|| human_readable)
226}
227
228pub(crate) fn is_human_readable() -> bool {
229    *SERDE_HUMAN_READABLE.get_or_init(|| DEFAULT_SERDE_HUMAN_READABLE)
230}
231
232#[cfg(test)]
233mod tests {
234    use super::*;
235    use apache_avro_test_helper::TestResult;
236    use pretty_assertions::assert_eq;
237
238    #[test]
239    fn test_zigzag() {
240        let mut a = Vec::new();
241        let mut b = Vec::new();
242        zig_i32(42i32, &mut a).unwrap();
243        zig_i64(42i64, &mut b).unwrap();
244        assert_eq!(a, b);
245    }
246
247    #[test]
248    fn test_zig_i64() {
249        let mut s = Vec::new();
250
251        zig_i64(0, &mut s).unwrap();
252        assert_eq!(s, [0]);
253
254        s.clear();
255        zig_i64(-1, &mut s).unwrap();
256        assert_eq!(s, [1]);
257
258        s.clear();
259        zig_i64(1, &mut s).unwrap();
260        assert_eq!(s, [2]);
261
262        s.clear();
263        zig_i64(-64, &mut s).unwrap();
264        assert_eq!(s, [127]);
265
266        s.clear();
267        zig_i64(64, &mut s).unwrap();
268        assert_eq!(s, [128, 1]);
269
270        s.clear();
271        zig_i64(i32::MAX as i64, &mut s).unwrap();
272        assert_eq!(s, [254, 255, 255, 255, 15]);
273
274        s.clear();
275        zig_i64(i32::MAX as i64 + 1, &mut s).unwrap();
276        assert_eq!(s, [128, 128, 128, 128, 16]);
277
278        s.clear();
279        zig_i64(i32::MIN as i64, &mut s).unwrap();
280        assert_eq!(s, [255, 255, 255, 255, 15]);
281
282        s.clear();
283        zig_i64(i32::MIN as i64 - 1, &mut s).unwrap();
284        assert_eq!(s, [129, 128, 128, 128, 16]);
285
286        s.clear();
287        zig_i64(i64::MAX, &mut s).unwrap();
288        assert_eq!(s, [254, 255, 255, 255, 255, 255, 255, 255, 255, 1]);
289
290        s.clear();
291        zig_i64(i64::MIN, &mut s).unwrap();
292        assert_eq!(s, [255, 255, 255, 255, 255, 255, 255, 255, 255, 1]);
293    }
294
295    #[test]
296    fn test_zig_i32() {
297        let mut s = Vec::new();
298        zig_i32(i32::MAX / 2, &mut s).unwrap();
299        assert_eq!(s, [254, 255, 255, 255, 7]);
300
301        s.clear();
302        zig_i32(i32::MIN / 2, &mut s).unwrap();
303        assert_eq!(s, [255, 255, 255, 255, 7]);
304
305        s.clear();
306        zig_i32(-(i32::MIN / 2), &mut s).unwrap();
307        assert_eq!(s, [128, 128, 128, 128, 8]);
308
309        s.clear();
310        zig_i32(i32::MIN / 2 - 1, &mut s).unwrap();
311        assert_eq!(s, [129, 128, 128, 128, 8]);
312
313        s.clear();
314        zig_i32(i32::MAX, &mut s).unwrap();
315        assert_eq!(s, [254, 255, 255, 255, 15]);
316
317        s.clear();
318        zig_i32(i32::MIN, &mut s).unwrap();
319        assert_eq!(s, [255, 255, 255, 255, 15]);
320    }
321
322    #[test]
323    fn test_overflow() {
324        let causes_left_shift_overflow: &[u8] = &[0xe1; 10];
325        assert!(matches!(
326            decode_variable(&mut &*causes_left_shift_overflow)
327                .unwrap_err()
328                .details(),
329            Details::IntegerOverflow
330        ));
331    }
332
333    #[test]
334    fn test_safe_len() -> TestResult {
335        assert_eq!(42usize, safe_len(42usize)?);
336        assert!(safe_len(1024 * 1024 * 1024).is_err());
337
338        Ok(())
339    }
340}
apache_avro/util.rs

apache_avro/
util.rs
