avro_rs/
util.rs

1use crate::{AvroResult, Error};
2use serde_json::{Map, Value};
3use std::{convert::TryFrom, i64, io::Read, sync::Once};
4
5/// Maximum number of bytes that can be allocated when decoding
6/// Avro-encoded values. This is a protection against ill-formed
7/// data, whose length field might be interpreted as enourmous.
8/// See max_allocation_bytes to change this limit.
9pub static mut MAX_ALLOCATION_BYTES: usize = 512 * 1024 * 1024;
10static MAX_ALLOCATION_BYTES_ONCE: Once = Once::new();
11
12pub trait MapHelper {
13    fn string(&self, key: &str) -> Option<String>;
14
15    fn name(&self) -> Option<String> {
16        self.string("name")
17    }
18
19    fn doc(&self) -> Option<String> {
20        self.string("doc")
21    }
22}
23
24impl MapHelper for Map<String, Value> {
25    fn string(&self, key: &str) -> Option<String> {
26        self.get(key)
27            .and_then(|v| v.as_str())
28            .map(|v| v.to_string())
29    }
30}
31
32pub fn read_long<R: Read>(reader: &mut R) -> AvroResult<i64> {
33    zag_i64(reader)
34}
35
36pub fn zig_i32(n: i32, buffer: &mut Vec<u8>) {
37    zig_i64(n as i64, buffer)
38}
39
40pub fn zig_i64(n: i64, buffer: &mut Vec<u8>) {
41    encode_variable(((n << 1) ^ (n >> 63)) as u64, buffer)
42}
43
44pub fn zag_i32<R: Read>(reader: &mut R) -> AvroResult<i32> {
45    let i = zag_i64(reader)?;
46    i32::try_from(i).map_err(|e| Error::ZagI32(e, i))
47}
48
49pub fn zag_i64<R: Read>(reader: &mut R) -> AvroResult<i64> {
50    let z = decode_variable(reader)?;
51    Ok(if z & 0x1 == 0 {
52        (z >> 1) as i64
53    } else {
54        !(z >> 1) as i64
55    })
56}
57
58fn encode_variable(mut z: u64, buffer: &mut Vec<u8>) {
59    loop {
60        if z <= 0x7F {
61            buffer.push((z & 0x7F) as u8);
62            break;
63        } else {
64            buffer.push((0x80 | (z & 0x7F)) as u8);
65            z >>= 7;
66        }
67    }
68}
69
70fn decode_variable<R: Read>(reader: &mut R) -> AvroResult<u64> {
71    let mut i = 0u64;
72    let mut buf = [0u8; 1];
73
74    let mut j = 0;
75    loop {
76        if j > 9 {
77            // if j * 7 > 64
78            return Err(Error::IntegerOverflow);
79        }
80        reader
81            .read_exact(&mut buf[..])
82            .map_err(Error::ReadVariableIntegerBytes)?;
83        i |= (u64::from(buf[0] & 0x7F)) << (j * 7);
84        if (buf[0] >> 7) == 0 {
85            break;
86        } else {
87            j += 1;
88        }
89    }
90
91    Ok(i)
92}
93
94/// Set a new maximum number of bytes that can be allocated when decoding data.
95/// Once called, the limit cannot be changed.
96///
97/// **NOTE** This function must be called before decoding **any** data. The
98/// library leverages [`std::sync::Once`](https://doc.rust-lang.org/std/sync/struct.Once.html)
99/// to set the limit either when calling this method, or when decoding for
100/// the first time.
101pub fn max_allocation_bytes(num_bytes: usize) -> usize {
102    unsafe {
103        MAX_ALLOCATION_BYTES_ONCE.call_once(|| {
104            MAX_ALLOCATION_BYTES = num_bytes;
105        });
106        MAX_ALLOCATION_BYTES
107    }
108}
109
110pub fn safe_len(len: usize) -> AvroResult<usize> {
111    let max_bytes = max_allocation_bytes(512 * 1024 * 1024);
112
113    if len <= max_bytes {
114        Ok(len)
115    } else {
116        Err(Error::MemoryAllocation {
117            desired: len,
118            maximum: max_bytes,
119        })
120    }
121}
122
123#[cfg(test)]
124mod tests {
125    use super::*;
126
127    #[test]
128    fn test_zigzag() {
129        let mut a = Vec::new();
130        let mut b = Vec::new();
131        zig_i32(42i32, &mut a);
132        zig_i64(42i64, &mut b);
133        assert_eq!(a, b);
134    }
135
136    #[test]
137    fn test_zig_i64() {
138        let mut s = Vec::new();
139        zig_i64(std::i32::MAX as i64, &mut s);
140        assert_eq!(s, [254, 255, 255, 255, 15]);
141
142        s.clear();
143        zig_i64(std::i32::MAX as i64 + 1, &mut s);
144        assert_eq!(s, [128, 128, 128, 128, 16]);
145
146        s.clear();
147        zig_i64(std::i32::MIN as i64, &mut s);
148        assert_eq!(s, [255, 255, 255, 255, 15]);
149
150        s.clear();
151        zig_i64(std::i32::MIN as i64 - 1, &mut s);
152        assert_eq!(s, [129, 128, 128, 128, 16]);
153
154        s.clear();
155        zig_i64(i64::MAX, &mut s);
156        assert_eq!(s, [254, 255, 255, 255, 255, 255, 255, 255, 255, 1]);
157
158        s.clear();
159        zig_i64(i64::MIN, &mut s);
160        assert_eq!(s, [255, 255, 255, 255, 255, 255, 255, 255, 255, 1]);
161    }
162
163    #[test]
164    fn test_zig_i32() {
165        let mut s = Vec::new();
166        zig_i32(std::i32::MAX / 2, &mut s);
167        assert_eq!(s, [254, 255, 255, 255, 7]);
168
169        s.clear();
170        zig_i32(std::i32::MIN / 2, &mut s);
171        assert_eq!(s, [255, 255, 255, 255, 7]);
172
173        s.clear();
174        zig_i32(-(std::i32::MIN / 2), &mut s);
175        assert_eq!(s, [128, 128, 128, 128, 8]);
176
177        s.clear();
178        zig_i32(std::i32::MIN / 2 - 1, &mut s);
179        assert_eq!(s, [129, 128, 128, 128, 8]);
180
181        s.clear();
182        zig_i32(std::i32::MAX, &mut s);
183        assert_eq!(s, [254, 255, 255, 255, 15]);
184
185        s.clear();
186        zig_i32(std::i32::MIN, &mut s);
187        assert_eq!(s, [255, 255, 255, 255, 15]);
188    }
189
190    #[test]
191    fn test_overflow() {
192        let causes_left_shift_overflow: &[u8] = &[0xe1, 0xe1, 0xe1, 0xe1, 0xe1];
193        assert!(decode_variable(&mut &causes_left_shift_overflow[..]).is_err());
194    }
195
196    #[test]
197    fn test_safe_len() {
198        assert_eq!(42usize, safe_len(42usize).unwrap());
199        assert!(safe_len(1024 * 1024 * 1024).is_err());
200    }
201}
avro_rs/util.rs

avro_rs/
util.rs
