ryerraguntla commented on code in PR #3519: URL: https://github.com/apache/iggy/pull/3519#discussion_r3830789610
########## gateways/kafka/src/protocol/bounds_guard.rs: ########## @@ -0,0 +1,721 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Pre-decode bounds guard against unbounded allocation in `kafka_protocol`'s decode path. +//! +//! `kafka_protocol` 0.17 validates every wire-declared array/string/bytes length against zero +//! but never against what remains in the frame (`kafka-protocol-0.17.0/src/protocol/types.rs:988` +//! `Vec::with_capacity(n as usize)` from a wire `Int32`, `:1096` the same from an unsigned +//! varint). A tiny frame that declares a huge count reaches `handle_alloc_error`, which calls +//! `abort()` - not a panic, not catchable, and it takes down every connection on the process, not +//! just the one that sent it. Reproduced: a 20-byte `CreateTopics` v5 frame claiming a huge +//! topics count requests ~481 GB; a 4-byte Metadata v0 frame requests ~143 GiB. No SASL/TLS gates +//! any of `SUPPORTED_RANGES`, so this is reachable by anyone who can `connect()`. +//! +//! This module walks the same field shape `kafka_protocol`'s real decode walks for each of the +//! six accepted message types, but only to validate every length-prefixed field (array count, +//! string length, bytes length, tagged-field size) against what could still fit in the bytes +//! remaining in the frame - it never materializes a value or allocates a collection. Call the +//! matching `validate_*_shape` function before handing the body to `kafka_protocol`. +//! +//! Kept independent of `kafka_protocol`'s own type decoders on purpose: if this walker's field +//! order ever drifts from the real schema, the worst outcome is a false-positive rejection of a +//! legitimate request (caught immediately by the wire-fixture round-trip tests), never a false +//! negative that lets an oversized count back through to the crate. + +use bytes::{Buf, Bytes}; + +use crate::error::{KafkaProtocolError, Result}; + +/// Same bound the pre-migration hand-rolled codec used: large enough for any real request, small +/// enough that even `MAX_REQUEST_ELEMENTS` identically-sized nested arrays cannot approach a +/// meaningful fraction of memory. +const MAX_COLLECTION_LEN: usize = 65_536; + +/// Cumulative cap on the total array elements this guard will walk across one request. A single +/// array's count is capped by [`MAX_COLLECTION_LEN`], but not the product across nested arrays +/// (`topics` x `partitions`) - this budget bounds the sum across the whole frame. +const MAX_REQUEST_ELEMENTS: usize = MAX_COLLECTION_LEN; + +struct ShapeCursor { + bytes: Bytes, + element_budget: usize, +} + +impl ShapeCursor { + const fn new(body: Bytes) -> Self { + Self { + bytes: body, + element_budget: MAX_REQUEST_ELEMENTS, + } + } + + fn remaining(&self) -> usize { + self.bytes.remaining() + } + + fn ensure(&self, needed: usize) -> Result<()> { + let remaining = self.bytes.remaining(); + if remaining < needed { + return Err(KafkaProtocolError::BufferUnderflow { needed, remaining }); + } + Ok(()) + } + + fn read_i8(&mut self) -> Result<i8> { + self.ensure(1)?; + Ok(self.bytes.get_i8()) + } + + fn read_bool(&mut self) -> Result<bool> { + Ok(self.read_i8()? != 0) + } + + fn read_i16(&mut self) -> Result<i16> { + self.ensure(2)?; + Ok(self.bytes.get_i16()) + } + + fn read_i32(&mut self) -> Result<i32> { + self.ensure(4)?; + Ok(self.bytes.get_i32()) + } + + fn read_i64(&mut self) -> Result<i64> { + self.ensure(8)?; + Ok(self.bytes.get_i64()) + } + + fn read_varint(&mut self) -> Result<u64> { Review Comment: Done! -- 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]
