Kriskras99 commented on code in PR #639:
URL: https://github.com/apache/avro-rs/pull/639#discussion_r3860320375
##########
avro/src/decode.rs:
##########
@@ -68,17 +71,64 @@ fn decode_seq_len<R: Read>(reader: &mut R) ->
AvroResult<usize> {
)
}
+/// Per-datum decoding state.
+///
+/// Tracks the cumulative number of bytes allocated on behalf of a single
+/// datum, so that nested collections cannot multiply the allocation budget:
+/// every allocation performed while decoding one datum is debited from a
+/// shared budget of [`max_allocation_bytes`] bytes, instead of each
+/// collection only being checked in isolation.
+pub(crate) struct DecodeContext {
+ /// Bytes still available for allocations while decoding the current datum.
+ remaining_budget: usize,
+}
+
+impl DecodeContext {
+ /// Create a fresh context. Call once per datum.
+ pub(crate) fn new() -> Self {
+ Self {
+ remaining_budget:
max_allocation_bytes(DEFAULT_MAX_ALLOCATION_BYTES),
+ }
+ }
+
+ /// Debit `bytes` from the per-datum allocation budget, erroring when the
+ /// cumulative allocations for this datum would exceed it.
+ fn debit(&mut self, bytes: usize) -> AvroResult<()> {
Review Comment:
I think renaming this to `debit_bytes` would make it a lot clearer to future
readers of the code
##########
avro/src/decode.rs:
##########
@@ -68,17 +71,64 @@ fn decode_seq_len<R: Read>(reader: &mut R) ->
AvroResult<usize> {
)
}
+/// Per-datum decoding state.
+///
+/// Tracks the cumulative number of bytes allocated on behalf of a single
+/// datum, so that nested collections cannot multiply the allocation budget:
+/// every allocation performed while decoding one datum is debited from a
+/// shared budget of [`max_allocation_bytes`] bytes, instead of each
+/// collection only being checked in isolation.
+pub(crate) struct DecodeContext {
+ /// Bytes still available for allocations while decoding the current datum.
+ remaining_budget: usize,
+}
+
+impl DecodeContext {
+ /// Create a fresh context. Call once per datum.
Review Comment:
```suggestion
/// Create a new context.
///
/// This should only be done when a new datum is being decoded, never
during the decoding.
```
##########
avro/src/decode.rs:
##########
@@ -68,17 +71,64 @@ fn decode_seq_len<R: Read>(reader: &mut R) ->
AvroResult<usize> {
)
}
+/// Per-datum decoding state.
+///
+/// Tracks the cumulative number of bytes allocated on behalf of a single
+/// datum, so that nested collections cannot multiply the allocation budget:
+/// every allocation performed while decoding one datum is debited from a
+/// shared budget of [`max_allocation_bytes`] bytes, instead of each
+/// collection only being checked in isolation.
+pub(crate) struct DecodeContext {
+ /// Bytes still available for allocations while decoding the current datum.
+ remaining_budget: usize,
+}
+
+impl DecodeContext {
+ /// Create a fresh context. Call once per datum.
+ pub(crate) fn new() -> Self {
+ Self {
+ remaining_budget:
max_allocation_bytes(DEFAULT_MAX_ALLOCATION_BYTES),
+ }
+ }
+
+ /// Debit `bytes` from the per-datum allocation budget, erroring when the
+ /// cumulative allocations for this datum would exceed it.
Review Comment:
```suggestion
/// Debit `bytes` from the per-datum allocation budget
///
/// # Errors
/// `Details::MemoryAllocation` if the maximum budget is exceeded.
```
##########
avro/src/decode.rs:
##########
@@ -68,17 +71,64 @@ fn decode_seq_len<R: Read>(reader: &mut R) ->
AvroResult<usize> {
)
}
+/// Per-datum decoding state.
+///
+/// Tracks the cumulative number of bytes allocated on behalf of a single
+/// datum, so that nested collections cannot multiply the allocation budget:
+/// every allocation performed while decoding one datum is debited from a
+/// shared budget of [`max_allocation_bytes`] bytes, instead of each
+/// collection only being checked in isolation.
+pub(crate) struct DecodeContext {
+ /// Bytes still available for allocations while decoding the current datum.
+ remaining_budget: usize,
+}
+
+impl DecodeContext {
+ /// Create a fresh context. Call once per datum.
+ pub(crate) fn new() -> Self {
+ Self {
+ remaining_budget:
max_allocation_bytes(DEFAULT_MAX_ALLOCATION_BYTES),
+ }
+ }
+
+ /// Debit `bytes` from the per-datum allocation budget, erroring when the
+ /// cumulative allocations for this datum would exceed it.
+ fn debit(&mut self, bytes: usize) -> AvroResult<()> {
+ match self.remaining_budget.checked_sub(bytes) {
+ Some(remaining) => {
+ self.remaining_budget = remaining;
+ Ok(())
+ }
+ None => Err(Details::MemoryAllocation {
+ desired: Some(bytes),
+ maximum: max_allocation_bytes(DEFAULT_MAX_ALLOCATION_BYTES),
+ }
+ .into()),
+ }
+ }
+
+ /// Debit the budget for `items` collection elements of type `T`.
+ fn debit_items<T>(&mut self, items: usize) -> AvroResult<()> {
+ let bytes = items
Review Comment:
```suggestion
/// Debit the amount of bytes for `n` items of `T`.
///
/// # Errors
/// `Details::MemoryAllocation` if the maximum budget is exceeded.
fn debit_items<T>(&mut self, n: usize) -> AvroResult<()> {
let bytes = n
```
##########
avro/src/decode.rs:
##########
@@ -226,6 +280,10 @@ pub(crate) fn decode_internal<R: Read, S: Borrow<Schema>>(
}
}
Schema::Fixed(FixedSchema { size, .. }) => {
+ // The size is schema-declared, not wire-declared, but the schema
+ // itself may be attacker-supplied (e.g. an OCF header), so it must
+ // be debited from the allocation budget like any other length.
Review Comment:
```suggestion
```
Doesn't add any extra information that is not in the documentation for
`DecodeContext`
##########
avro/src/decode.rs:
##########
@@ -68,17 +71,64 @@ fn decode_seq_len<R: Read>(reader: &mut R) ->
AvroResult<usize> {
)
}
+/// Per-datum decoding state.
+///
+/// Tracks the cumulative number of bytes allocated on behalf of a single
+/// datum, so that nested collections cannot multiply the allocation budget:
+/// every allocation performed while decoding one datum is debited from a
+/// shared budget of [`max_allocation_bytes`] bytes, instead of each
+/// collection only being checked in isolation.
+pub(crate) struct DecodeContext {
+ /// Bytes still available for allocations while decoding the current datum.
+ remaining_budget: usize,
+}
+
+impl DecodeContext {
+ /// Create a fresh context. Call once per datum.
+ pub(crate) fn new() -> Self {
+ Self {
+ remaining_budget:
max_allocation_bytes(DEFAULT_MAX_ALLOCATION_BYTES),
+ }
+ }
+
+ /// Debit `bytes` from the per-datum allocation budget, erroring when the
+ /// cumulative allocations for this datum would exceed it.
+ fn debit(&mut self, bytes: usize) -> AvroResult<()> {
+ match self.remaining_budget.checked_sub(bytes) {
+ Some(remaining) => {
+ self.remaining_budget = remaining;
+ Ok(())
+ }
+ None => Err(Details::MemoryAllocation {
+ desired: Some(bytes),
+ maximum: max_allocation_bytes(DEFAULT_MAX_ALLOCATION_BYTES),
+ }
+ .into()),
+ }
+ }
+
+ /// Debit the budget for `items` collection elements of type `T`.
+ fn debit_items<T>(&mut self, items: usize) -> AvroResult<()> {
+ let bytes = items
+ .checked_mul(size_of::<T>())
+ .ok_or(Details::IntegerOverflow)?;
+ self.debit(bytes)
+ }
+}
+
/// Decode a `Value` from avro format given its `Schema`.
pub fn decode<R: Read>(schema: &Schema, reader: &mut R) -> AvroResult<Value> {
let rs = ResolvedSchema::try_from(schema)?;
- decode_internal(schema, rs.get_names(), None, reader)
+ let mut ctx = DecodeContext::new();
+ decode_internal(schema, rs.get_names(), None, reader, &mut ctx)
Review Comment:
```suggestion
decode_internal(schema, rs.get_names(), None, reader, &mut
DecodeContext::new())
```
Doing it like this prevents the accidental re-use of the context
##########
avro/src/decode.rs:
##########
@@ -247,6 +305,10 @@ pub(crate) fn decode_internal<R: Read, S: Borrow<Schema>>(
.checked_add(len)
.ok_or(Details::IntegerOverflow)?;
safe_collection_len::<Value>(total)?;
+ // Elements can be arbitrarily cheap on the wire (e.g. null),
+ // so also debit the per-datum budget: nested collections must
+ // not multiply the allocation budget.
Review Comment:
```suggestion
```
Doesn't add any extra information that is not in the documentation for
`DecodeContext`.
Also, it's not only nested collections right? It's for the entire datum
##########
avro/src/decode.rs:
##########
@@ -279,13 +342,21 @@ pub(crate) fn decode_internal<R: Read, S: Borrow<Schema>>(
.checked_add(len)
.ok_or(Details::IntegerOverflow)?;
safe_collection_len::<(String, Value)>(total)?;
+ // See the Array arm: nested collections share one budget.
Review Comment:
```suggestion
```
Doesn't add any extra information that is not in the documentation for
`DecodeContext`
##########
avro/src/decode.rs:
##########
@@ -318,9 +389,13 @@ pub(crate) fn decode_internal<R: Read, S: Borrow<Schema>>(
},
Schema::Record(RecordSchema { name, fields, .. }) => {
let fully_qualified_name =
name.fully_qualified_name(enclosing_namespace);
+ // Records can consume zero wire bytes (e.g. all-null fields), so
+ // debit the budget for the field vector and the cloned names.
Review Comment:
```suggestion
```
Doesn't add any extra information that is not in the documentation for
`DecodeContext`
##########
avro/src/reader/block.rs:
##########
@@ -195,11 +195,13 @@ impl<'r, R: Read> Block<'r, R> {
let mut block_bytes = &self.buf[self.buf_idx..];
let b_original = block_bytes.len();
+ let mut ctx = DecodeContext::new();
let item = decode_internal(
&self.writer_schema,
&self.names_refs,
None,
&mut block_bytes,
+ &mut ctx,
Review Comment:
```suggestion
let item = decode_internal(
&self.writer_schema,
&self.names_refs,
None,
&mut block_bytes,
&mut DecodeContext::new(),
```
Prevent accidental re-use
##########
avro/src/reader/datum.rs:
##########
@@ -130,7 +130,14 @@ impl<'s, S: generic_datum_reader_builder::State>
GenericDatumReaderBuilder<'s, S
impl<'s> GenericDatumReader<'s> {
/// Read a Avro datum from the reader.
pub fn read_value<R: Read>(&self, reader: &mut R) -> AvroResult<Value> {
- let value = decode_internal(self.writer, self.resolved.get_names(),
None, reader)?;
+ let mut ctx = DecodeContext::new();
+ let value = decode_internal(
+ self.writer,
+ self.resolved.get_names(),
+ None,
+ reader,
+ &mut ctx,
Review Comment:
```suggestion
let value = decode_internal(
self.writer,
self.resolved.get_names(),
None,
reader,
&mut DecodeContext::new(),
```
Prevent accidental re-use
##########
avro/src/decode.rs:
##########
@@ -339,6 +415,9 @@ pub(crate) fn decode_internal<R: Read, S: Borrow<Schema>>(
let index = usize::try_from(raw_index)
.map_err(|e| Details::ConvertI32ToUsize(e, raw_index))?;
if (0..symbols.len()).contains(&index) {
+ // Cloning the symbol allocates without consuming wire
+ // bytes, so it counts against the per-datum budget.
Review Comment:
This is a very useful comment!
##########
avro/src/decode.rs:
##########
@@ -472,6 +552,41 @@ mod tests {
Ok(())
}
+ #[test]
+ fn avro_rs_639_test_nested_collections_share_one_allocation_budget() ->
TestResult {
+ use crate::util::{DEFAULT_MAX_ALLOCATION_BYTES, max_allocation_bytes};
+
+ // Each inner array<null> block passes the per-collection check on its
+ // own, but the shared per-datum budget must reject the cumulative
+ // total: elements of type null cost zero wire bytes, so without a
+ // cumulative budget a handful of ~10-byte inner arrays would pin an
+ // unbounded multiple of the allocation limit in memory at once.
+ let budget = max_allocation_bytes(DEFAULT_MAX_ALLOCATION_BYTES);
+ let inner_count = (budget / 2) / size_of::<Value>() + 1;
+
+ let inner_arrays_count = 2;
+ let mut payload = Vec::new();
+ // Outer array: a single block declaring two inner arrays.
+ crate::util::zig_i64(inner_arrays_count, &mut payload)?;
+ for _ in 0..inner_arrays_count {
+ payload.extend(create_block(inner_count as i64));
+ }
+ // Outer array terminator.
+ payload.push(0x00);
+
+ let result = decode(
+ &Schema::array(Schema::array(Schema::Null).build()).build(),
+ &mut payload.as_slice(),
+ );
+
+ assert!(
+ result.is_err(),
+ "nested collections must share one allocation budget, got
{result:?}"
+ );
Review Comment:
Please check the error is actually `Details::MemoryAllocation`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
