mz_avro/
util.rs

1// Copyright 2018 Flavien Raynaud.
2// Copyright Materialize, Inc. and contributors. All rights reserved.
3//
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License in the LICENSE file at the
7// root of this repository, or online at
8//
9//     http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing, software
12// distributed under the License is distributed on an "AS IS" BASIS,
13// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14// See the License for the specific language governing permissions and
15// limitations under the License.
16//
17// This file is derived from the avro-rs project, available at
18// https://github.com/flavray/avro-rs. It was incorporated
19// directly into Materialize on March 3, 2020.
20//
21// The original source code is subject to the terms of the MIT license, a copy
22// of which can be found in the LICENSE file at the root of this repository.
23
24use std::io::Read;
25
26use serde_json::{Map, Value};
27
28use crate::error::{DecodeError, Error as AvroError};
29
30/// Maximum number of bytes that can be allocated when decoding
31/// Avro-encoded values. This is a protection against ill-formed
32/// data, whose length field might be interpreted as enormous.
33pub const MAX_ALLOCATION_BYTES: usize = 512 * 1024 * 1024;
34
35#[derive(Debug, Clone, Copy, Eq, PartialEq)]
36pub enum TsUnit {
37    Millis,
38    Micros,
39}
40
41pub trait MapHelper {
42    fn string(&self, key: &str) -> Option<String>;
43
44    fn name(&self) -> Option<String> {
45        self.string("name")
46    }
47
48    fn doc(&self) -> Option<String> {
49        self.string("doc")
50    }
51}
52
53impl MapHelper for Map<String, Value> {
54    fn string(&self, key: &str) -> Option<String> {
55        self.get(key)
56            .and_then(|v| v.as_str())
57            .map(|v| v.to_string())
58    }
59}
60
61pub fn read_long<R: Read>(reader: &mut R) -> Result<i64, AvroError> {
62    zag_i64(reader)
63}
64
65pub fn zig_i32(n: i32, buffer: &mut Vec<u8>) {
66    zig_i64(n as i64, buffer)
67}
68
69pub fn zig_i64(n: i64, buffer: &mut Vec<u8>) {
70    encode_variable(((n << 1) ^ (n >> 63)) as u64, buffer)
71}
72
73pub fn zag_i32<R: Read>(reader: &mut R) -> Result<i32, AvroError> {
74    let i = zag_i64(reader)?;
75    if i < i64::from(i32::MIN) || i > i64::from(i32::MAX) {
76        Err(AvroError::Decode(DecodeError::I32OutOfRange(i)))
77    } else {
78        Ok(i as i32)
79    }
80}
81
82pub fn zag_i64<R: Read>(reader: &mut R) -> Result<i64, AvroError> {
83    let z = decode_variable(reader)?;
84    Ok(if z & 0x1 == 0 {
85        (z >> 1) as i64
86    } else {
87        !(z >> 1) as i64
88    })
89}
90
91fn encode_variable(mut z: u64, buffer: &mut Vec<u8>) {
92    loop {
93        if z <= 0x7F {
94            buffer.push((z & 0x7F) as u8);
95            break;
96        } else {
97            buffer.push((0x80 | (z & 0x7F)) as u8);
98            z >>= 7;
99        }
100    }
101}
102
103fn decode_variable<R: Read>(reader: &mut R) -> Result<u64, AvroError> {
104    let mut i = 0u64;
105    let mut buf = [0u8; 1];
106
107    let mut j = 0;
108    loop {
109        if j > 9 {
110            // if j * 7 > 64
111            return Err(AvroError::Decode(DecodeError::IntDecodeOverflow));
112        }
113        reader.read_exact(&mut buf[..])?;
114        i |= (u64::from(buf[0] & 0x7F)) << (j * 7);
115        if (buf[0] >> 7) == 0 {
116            break;
117        } else {
118            j += 1;
119        }
120    }
121
122    Ok(i)
123}
124
125pub fn safe_len(len: usize) -> Result<usize, AvroError> {
126    if len <= MAX_ALLOCATION_BYTES {
127        Ok(len)
128    } else {
129        Err(AvroError::Allocation {
130            attempted: len,
131            allowed: MAX_ALLOCATION_BYTES,
132        })
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139    use mz_ore::assert_err;
140
141    #[mz_ore::test]
142    fn test_zigzag() {
143        let mut a = Vec::new();
144        let mut b = Vec::new();
145        zig_i32(42i32, &mut a);
146        zig_i64(42i64, &mut b);
147        assert_eq!(a, b);
148    }
149
150    #[mz_ore::test]
151    fn test_zig_i64() {
152        let mut s = Vec::new();
153        zig_i64(2_147_483_647_i64, &mut s);
154        assert_eq!(s, [254, 255, 255, 255, 15]);
155
156        s.clear();
157        zig_i64(2_147_483_648_i64, &mut s);
158        assert_eq!(s, [128, 128, 128, 128, 16]);
159
160        s.clear();
161        zig_i64(-2_147_483_648_i64, &mut s);
162        assert_eq!(s, [255, 255, 255, 255, 15]);
163
164        s.clear();
165        zig_i64(-2_147_483_649_i64, &mut s);
166        assert_eq!(s, [129, 128, 128, 128, 16]);
167
168        s.clear();
169        zig_i64(i64::MAX, &mut s);
170        assert_eq!(s, [254, 255, 255, 255, 255, 255, 255, 255, 255, 1]);
171
172        s.clear();
173        zig_i64(i64::MIN, &mut s);
174        assert_eq!(s, [255, 255, 255, 255, 255, 255, 255, 255, 255, 1]);
175    }
176
177    #[mz_ore::test]
178    fn test_zig_i32() {
179        let mut s = Vec::new();
180        zig_i32(1_073_741_823_i32, &mut s);
181        assert_eq!(s, [254, 255, 255, 255, 7]);
182
183        s.clear();
184        zig_i32(-1_073_741_824_i32, &mut s);
185        assert_eq!(s, [255, 255, 255, 255, 7]);
186
187        s.clear();
188        zig_i32(1_073_741_824_i32, &mut s);
189        assert_eq!(s, [128, 128, 128, 128, 8]);
190
191        s.clear();
192        zig_i32(-1_073_741_825_i32, &mut s);
193        assert_eq!(s, [129, 128, 128, 128, 8]);
194
195        s.clear();
196        zig_i32(2_147_483_647_i32, &mut s);
197        assert_eq!(s, [254, 255, 255, 255, 15]);
198
199        s.clear();
200        zig_i32(-2_147_483_648_i32, &mut s);
201        assert_eq!(s, [255, 255, 255, 255, 15]);
202    }
203
204    #[mz_ore::test]
205    fn test_overflow() {
206        let causes_left_shift_overflow: &[u8] = &[0xe1, 0xe1, 0xe1, 0xe1, 0xe1];
207        assert_err!(decode_variable(&mut &causes_left_shift_overflow[..]));
208    }
209
210    #[mz_ore::test]
211    fn test_safe_len() {
212        assert_eq!(42usize, safe_len(42usize).unwrap());
213        assert_err!(safe_len(1024 * 1024 * 1024));
214    }
215}
mz_avro/util.rs

mz_avro/
util.rs
