arrow_avro/reader/
block.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//! Decoder for [`Block`]
19
20use crate::errors::AvroError;
21use crate::reader::vlq::VLQDecoder;
22
23/// A file data block
24///
25/// <https://avro.apache.org/docs/1.11.1/specification/#object-container-files>
26#[derive(Debug, Default)]
27pub struct Block {
28    /// The number of objects in this block
29    pub count: usize,
30    /// The serialized objects within this block
31    pub data: Vec<u8>,
32    /// The sync marker
33    pub sync: [u8; 16],
34}
35
36/// A decoder for [`Block`]
37#[derive(Debug)]
38pub struct BlockDecoder {
39    state: BlockDecoderState,
40    in_progress: Block,
41    vlq_decoder: VLQDecoder,
42    bytes_remaining: usize,
43}
44
45#[derive(Debug)]
46pub(crate) enum BlockDecoderState {
47    Count,
48    Size,
49    Data,
50    Sync,
51    Finished,
52}
53
54impl Default for BlockDecoder {
55    fn default() -> Self {
56        Self {
57            state: BlockDecoderState::Count,
58            in_progress: Default::default(),
59            vlq_decoder: Default::default(),
60            bytes_remaining: 0,
61        }
62    }
63}
64
65impl BlockDecoder {
66    /// Parse [`Block`] from `buf`, returning the number of bytes read
67    ///
68    /// This method can be called multiple times with consecutive chunks of data, allowing
69    /// integration with chunked IO systems like [`BufRead::fill_buf`]
70    ///
71    /// All errors should be considered fatal, and decoding aborted
72    ///
73    /// Once an entire [`Block`] has been decoded this method will not read any further
74    /// input bytes, until [`Self::flush`] is called. Afterwards [`Self::decode`]
75    /// can then be used again to read the next block, if any
76    ///
77    /// [`BufRead::fill_buf`]: std::io::BufRead::fill_buf
78    pub fn decode(&mut self, mut buf: &[u8]) -> Result<usize, AvroError> {
79        let max_read = buf.len();
80        while !buf.is_empty() {
81            match self.state {
82                BlockDecoderState::Count => {
83                    if let Some(c) = self.vlq_decoder.long(&mut buf)? {
84                        self.in_progress.count = c.try_into().map_err(|_| {
85                            AvroError::ParseError(format!(
86                                "Block count cannot be negative, got {c}"
87                            ))
88                        })?;
89
90                        self.state = BlockDecoderState::Size;
91                    }
92                }
93                BlockDecoderState::Size => {
94                    if let Some(c) = self.vlq_decoder.long(&mut buf)? {
95                        self.bytes_remaining = c.try_into().map_err(|_| {
96                            AvroError::ParseError(format!("Block size cannot be negative, got {c}"))
97                        })?;
98
99                        // Only reserve what the current input backs: the block size is
100                        // input specified so could be an extreme value (e.g. i64::MAX)
101                        // in case of corrupted/malicious input. The rest is reserved
102                        // lazily by `extend_from_slice` below as data arrives.
103                        self.in_progress
104                            .data
105                            .reserve(self.bytes_remaining.min(buf.len()));
106                        self.state = BlockDecoderState::Data;
107                    }
108                }
109                BlockDecoderState::Data => {
110                    let to_read = self.bytes_remaining.min(buf.len());
111                    self.in_progress.data.extend_from_slice(&buf[..to_read]);
112                    buf = &buf[to_read..];
113                    self.bytes_remaining -= to_read;
114                    if self.bytes_remaining == 0 {
115                        self.bytes_remaining = 16;
116                        self.state = BlockDecoderState::Sync;
117                    }
118                }
119                BlockDecoderState::Sync => {
120                    let to_decode = buf.len().min(self.bytes_remaining);
121                    // Fill from the front: the marker may arrive split across decode() calls.
122                    let offset = 16 - self.bytes_remaining;
123                    self.in_progress.sync[offset..offset + to_decode]
124                        .copy_from_slice(&buf[..to_decode]);
125                    self.bytes_remaining -= to_decode;
126                    buf = &buf[to_decode..];
127                    if self.bytes_remaining == 0 {
128                        self.state = BlockDecoderState::Finished;
129                    }
130                }
131                BlockDecoderState::Finished => return Ok(max_read - buf.len()),
132            }
133        }
134        Ok(max_read)
135    }
136
137    /// Flush this decoder returning the parsed [`Block`] if any
138    pub fn flush(&mut self) -> Option<Block> {
139        match self.state {
140            BlockDecoderState::Finished => {
141                self.state = BlockDecoderState::Count;
142                Some(std::mem::take(&mut self.in_progress))
143            }
144            _ => None,
145        }
146    }
147}
148
149#[cfg(feature = "async")]
150impl BlockDecoder {
151    pub(crate) fn state(&self) -> &BlockDecoderState {
152        &self.state
153    }
154
155    pub(crate) fn bytes_remaining(&self) -> usize {
156        self.bytes_remaining
157    }
158}
159
160#[cfg(test)]
161mod tests {
162    use super::*;
163
164    /// Zig-zag encode `value` as an Avro `long` (variable-length integer).
165    fn encode_long(value: i64, out: &mut Vec<u8>) {
166        let mut n = ((value << 1) ^ (value >> 63)) as u64;
167        while n >= 0x80 {
168            out.push((n as u8) | 0x80);
169            n >>= 7;
170        }
171        out.push(n as u8);
172    }
173
174    #[test]
175    fn test_oversized_block_size_bounds_reserve() {
176        // A block advertising `i64::MAX` bytes must not reserve that up front when only
177        // a few payload bytes are present, or a crafted OCF aborts on a huge alloc (#10234).
178        let mut buf = Vec::new();
179        encode_long(1, &mut buf); // object count
180        encode_long(i64::MAX, &mut buf); // attacker-controlled block size
181        buf.extend_from_slice(&[0u8; 8]); // a handful of real bytes
182
183        let mut decoder = BlockDecoder::default();
184        let read = decoder.decode(&buf).unwrap();
185
186        assert_eq!(read, buf.len(), "all available input should be consumed");
187        assert!(
188            decoder.in_progress.data.capacity() <= buf.len(),
189            "capacity {} must stay bounded by available input {}, not the advertised i64::MAX",
190            decoder.in_progress.data.capacity(),
191            buf.len(),
192        );
193    }
194
195    #[test]
196    fn test_negative_block_size_errors() {
197        let mut buf = Vec::new();
198        encode_long(1, &mut buf); // object count
199        encode_long(-1, &mut buf); // invalid (negative) block size
200
201        let mut decoder = BlockDecoder::default();
202        let err = decoder.decode(&buf).unwrap_err();
203        assert!(
204            err.to_string().contains("Block size cannot be negative"),
205            "unexpected error: {err}",
206        );
207    }
208
209    #[test]
210    fn test_well_formed_block_round_trips() {
211        // The capped reserve must not change decoding of a normal block.
212        let payload = [1u8, 2, 3, 4];
213        let sync = [7u8; 16];
214        let mut buf = Vec::new();
215        encode_long(2, &mut buf); // object count
216        encode_long(payload.len() as i64, &mut buf); // block size
217        buf.extend_from_slice(&payload);
218        buf.extend_from_slice(&sync);
219
220        let mut decoder = BlockDecoder::default();
221        assert_eq!(decoder.decode(&buf).unwrap(), buf.len());
222        let block = decoder.flush().expect("a complete block");
223        assert_eq!(block.count, 2);
224        assert_eq!(block.data, payload);
225        assert_eq!(block.sync, sync);
226    }
227
228    #[test]
229    fn test_sync_marker_split_across_decode_calls() {
230        // count=1 (zig-zag 0x02), size=1 (0x02), one data byte, 16-byte sync marker
231        let sync: [u8; 16] = core::array::from_fn(|i| i as u8);
232        let mut block_bytes = vec![0x02, 0x02, 0xAA];
233        block_bytes.extend_from_slice(&sync);
234
235        for chunk_size in 1..block_bytes.len() {
236            let mut decoder = BlockDecoder::default();
237            for chunk in block_bytes.chunks(chunk_size) {
238                decoder.decode(chunk).unwrap();
239            }
240            let block = decoder.flush().expect("complete block");
241            assert_eq!(block.count, 1);
242            assert_eq!(block.data, vec![0xAA]);
243            assert_eq!(block.sync, sync, "chunk_size {chunk_size}");
244        }
245    }
246}
arrow_avro/reader/block.rs

arrow_avro/reader/
block.rs
