Bound source side peer state

lib/codecs/src/decoding/framing/chunked_gelf.rs

@@ -20,10 +20,33 @@ const GELF_MAGIC: &[u8] = &[0x1e, 0x0f];
 const GELF_MAX_TOTAL_CHUNKS: u8 = 128;
 const DEFAULT_TIMEOUT_SECS: f64 = 5.0;
 
+/// The number of incomplete messages tracked at once, by default.
+///
+/// Each pending message costs roughly 10 KB regardless of how many payload bytes have
+/// actually arrived, because `MessageState` holds a fixed `[Bytes; 128]` slot array. A
+/// 12-byte header with no payload is enough to allocate one, so leaving this unbounded
+/// lets an unauthenticated peer amplify its traffic by ~870x.
+const DEFAULT_PENDING_MESSAGES_LIMIT: usize = 5_000;
+
+/// The maximum reassembled payload of a single chunked message, by default.
+///
+/// Comfortably above any real GELF-over-UDP message: the protocol caps a message at 128
+/// chunks, and clients pick an MTU-sized chunk (Graylog's own go-gelf uses 1420 bytes,
+/// for ~182 KB total), so this only rejects payloads that no conforming sender produces.
+const DEFAULT_MAX_LENGTH_BYTES: usize = 1024 * 1024;
+
 const fn default_timeout_secs() -> f64 {
     DEFAULT_TIMEOUT_SECS
 }
 
+const fn default_pending_messages_limit() -> Option<usize> {
+    Some(DEFAULT_PENDING_MESSAGES_LIMIT)
+}
+
+const fn default_max_length() -> Option<usize> {
+    Some(DEFAULT_MAX_LENGTH_BYTES)
+}
+
 /// Config used to build a `ChunkedGelfDecoder`.
 #[configurable_component]
 #[derive(Debug, Clone, Default, PartialEq, Eq)]
@@ -57,22 +80,29 @@ pub struct ChunkedGelfDecoderOptions {
     pub timeout_secs: f64,
 
     /// The maximum number of pending incomplete messages. If this limit is reached, the decoder starts
-    /// dropping chunks of new messages, ensuring the memory usage of the decoder's state is bounded.
-    /// If this option is not set, the decoder does not limit the number of pending messages and the memory usage
-    /// of its messages buffer can grow unbounded. This matches Graylog Server's behavior.
-    #[serde(default, skip_serializing_if = "vector_core::serde::is_default")]
+    /// dropping chunks of *new* messages; chunks of messages already being reassembled are still
+    /// accepted, so reaching the limit does not corrupt in-flight messages.
+    ///
+    /// Defaults to 5000. Setting this to `null` restores the unbounded behavior, which matches
+    /// Graylog Server but lets any peer that can reach the socket exhaust memory: a pending message
+    /// costs roughly 10 KB no matter how few payload bytes have arrived, so a 12-byte header is
+    /// enough to allocate one. Do not disable this on an unauthenticated listener.
+    #[serde(default = "default_pending_messages_limit")]
+    #[derivative(Default(value = "default_pending_messages_limit()"))]
     pub pending_messages_limit: Option<usize>,
 
     /// The maximum length of a single GELF message, in bytes. Messages longer than this length will
-    /// be dropped. If this option is not set, the decoder does not limit the length of messages and
-    /// the per-message memory is unbounded.
+    /// be dropped, and the chunk that would exceed it is never buffered.
     ///
     /// Note that a message can be composed of multiple chunks and this limit is applied to the whole
     /// message, not to individual chunks.
     ///
     /// This limit takes only into account the message's payload and the GELF header bytes are excluded from the calculation.
     /// The message's payload is the concatenation of all the chunks' payloads.
-    #[serde(default, skip_serializing_if = "vector_core::serde::is_default")]
+    ///
+    /// Defaults to 1048576 (1 MiB). Setting this to `null` leaves the per-message memory unbounded.
+    #[serde(default = "default_max_length")]
+    #[derivative(Default(value = "default_max_length()"))]
     pub max_length: Option<usize>,
 
     /// Decompression configuration for GELF messages.
@@ -377,9 +407,14 @@ impl ChunkedGelfDecoder {
 
         let mut state_lock = self.state.lock().expect("poisoned lock");
 
+        // The limit bounds how many *distinct* messages are tracked, so it must only reject
+        // message ids that are not already being reassembled. Checking it before the lookup
+        // would drop follow-up chunks of in-flight messages whenever the map is full - i.e.
+        // precisely under the load the limit exists to survive - silently corrupting
+        // legitimate traffic instead of only shedding new attacker-chosen ids.
         if let Some(pending_messages_limit) = self.pending_messages_limit {
             ensure!(
-                state_lock.len() < pending_messages_limit,
+                state_lock.len() < pending_messages_limit || state_lock.contains_key(&message_id),
                 PendingMessagesLimitReachedSnafu {
                     message_id,
                     sequence_number,
@@ -428,10 +463,13 @@ impl ChunkedGelfDecoder {
             return Ok(None);
         }
 
-        message_state.add_chunk(sequence_number, chunk);
-
+        // Checked before the chunk is stored rather than after. This is defensive only: `chunk` is
+        // a refcounted slice of the datagram, so storing it copies nothing and the entry is dropped
+        // on the error path either way - the two orders retain the same memory today. It matters if
+        // `add_chunk` ever starts copying. `length` is the size the message would reach, which is
+        // the same value the previous store-then-measure order reported.
         if let Some(max_length) = self.max_length {
-            let length = message_state.current_length();
+            let length = message_state.current_length() + chunk.remaining();
             if length > max_length {
                 state_lock.remove(&message_id);
                 return Err(ChunkedGelfDecoderError::MaxLengthExceed {
@@ -443,6 +481,8 @@ impl ChunkedGelfDecoder {
             }
         }
 
+        message_state.add_chunk(sequence_number, chunk);
+
         if let Some(message) = message_state.retrieve_message() {
             state_lock.remove(&message_id);
             Ok(Some(message))
@@ -486,8 +526,8 @@ impl Default for ChunkedGelfDecoder {
     fn default() -> Self {
         Self::new(
             DEFAULT_TIMEOUT_SECS,
-            None,
-            None,
+            default_pending_messages_limit(),
+            default_max_length(),
             ChunkedGelfDecompressionConfig::Auto,
         )
     }
@@ -1278,4 +1318,175 @@ mod tests {
 
         assert_eq!(detected_compression, ChunkedGelfDecompression::None);
     }
+
+    // OBE-11235: the reassembly map is keyed solely on an attacker-chosen `message_id` and was
+    // unbounded by default. A 12-byte header with no payload allocates a ~10 KB `MessageState`,
+    // so an unauthenticated peer could amplify its traffic ~870x until the process was OOM-killed.
+
+    /// A header-only chunk: enough to allocate reassembly state, no payload.
+    fn header_only_chunk(message_id: u64) -> BytesMut {
+        create_chunk(message_id, 0, GELF_MAX_TOTAL_CHUNKS, &"")
+    }
+
+    #[tokio::test]
+    async fn defaults_bound_the_reassembly_state() {
+        let options = ChunkedGelfDecoderOptions::default();
+
+        assert_eq!(
+            options.pending_messages_limit,
+            Some(DEFAULT_PENDING_MESSAGES_LIMIT)
+        );
+        assert_eq!(options.max_length, Some(DEFAULT_MAX_LENGTH_BYTES));
+
+        // The socket source reaches the decoder through `Default`, not through the options
+        // struct, so both paths have to carry the limits.
+        let decoder = ChunkedGelfDecoder::default();
+        assert_eq!(
+            decoder.pending_messages_limit,
+            Some(DEFAULT_PENDING_MESSAGES_LIMIT)
+        );
+        assert_eq!(decoder.max_length, Some(DEFAULT_MAX_LENGTH_BYTES));
+    }
+
+    #[tokio::test]
+    async fn pending_map_stops_growing_at_the_default_limit() {
+        let mut decoder = ChunkedGelfDecoder::default();
+        let overshoot = 250;
+
+        let mut rejected = 0;
+        for message_id in 0..(DEFAULT_PENDING_MESSAGES_LIMIT + overshoot) as u64 {
+            if decoder
+                .decode_eof(&mut header_only_chunk(message_id))
+                .is_err()
+            {
+                rejected += 1;
+            }
+        }
+
+        assert_eq!(
+            decoder.state.lock().unwrap().len(),
+            DEFAULT_PENDING_MESSAGES_LIMIT,
+            "the map must stop growing at the limit"
+        );
+        assert_eq!(
+            rejected, overshoot,
+            "every id beyond the limit must be rejected"
+        );
+    }
+
+    #[rstest]
+    #[tokio::test]
+    async fn a_full_pending_map_still_accepts_chunks_of_tracked_messages(
+        two_chunks_message: ([BytesMut; 2], String),
+    ) {
+        // The regression the limit itself used to introduce: because the check ran before the
+        // map lookup, filling the map rejected follow-up chunks of messages already being
+        // reassembled - dropping legitimate traffic exactly when the limit started to bite.
+        let (mut chunks, expected_message) = two_chunks_message;
+        let limit = 4;
+        let mut decoder = ChunkedGelfDecoder {
+            pending_messages_limit: Some(limit),
+            ..Default::default()
+        };
+
+        // Start reassembling a real message, then fill the rest of the map with other ids.
+        assert!(decoder.decode_eof(&mut chunks[0]).unwrap().is_none());
+        let tracked_id = *decoder.state.lock().unwrap().keys().next().unwrap();
+        for message_id in 0..limit as u64 {
+            if message_id != tracked_id {
+                let _ = decoder.decode_eof(&mut header_only_chunk(message_id));
+            }
+        }
+        assert_eq!(decoder.state.lock().unwrap().len(), limit, "map is full");
+
+        // A brand-new id is shed...
+        let error = decoder
+            .decode_eof(&mut header_only_chunk(u64::MAX))
+            .unwrap_err();
+        assert!(matches!(
+            downcast_framing_error(&error),
+            ChunkedGelfDecoderError::PendingMessagesLimitReached { .. }
+        ));
+
+        // ...but the in-flight message still completes.
+        let frame = decoder
+            .decode_eof(&mut chunks[1])
+            .expect("a chunk of an already-tracked message must not be rejected")
+            .expect("the message should be complete");
+        assert_eq!(frame, expected_message);
+    }
+
+    #[tokio::test]
+    async fn oversized_message_is_rejected_and_leaves_no_state() {
+        let message_id = 1u64;
+        let max_length = 32;
+        let mut decoder = ChunkedGelfDecoder {
+            max_length: Some(max_length),
+            ..Default::default()
+        };
+
+        let payload = "a".repeat(max_length + 1);
+        let error = decoder
+            .decode_eof(&mut create_chunk(message_id, 0, 2, &payload))
+            .unwrap_err();
+
+        assert!(matches!(
+            downcast_framing_error(&error),
+            ChunkedGelfDecoderError::MaxLengthExceed { .. }
+        ));
+        assert!(
+            decoder.state.lock().unwrap().is_empty(),
+            "the over-budget message must not be left buffered"
+        );
+    }
+
+    #[tokio::test]
+    async fn message_over_the_default_max_length_is_rejected() {
+        let mut decoder = ChunkedGelfDecoder::default();
+
+        // Two chunks, so the message never completes and the length check is what stops it.
+        let payload = "a".repeat(DEFAULT_MAX_LENGTH_BYTES + 1);
+        let error = decoder
+            .decode_eof(&mut create_chunk(1, 0, 2, &payload))
+            .unwrap_err();
+
+        assert!(matches!(
+            downcast_framing_error(&error),
+            ChunkedGelfDecoderError::MaxLengthExceed { .. }
+        ));
+    }
+
+    #[tokio::test]
+    async fn message_exactly_at_max_length_is_accepted() {
+        let message_id = 1u64;
+        let max_length = 32;
+        let mut decoder = ChunkedGelfDecoder {
+            max_length: Some(max_length),
+            ..Default::default()
+        };
+
+        let payload = "a".repeat(max_length);
+        let frame = decoder
+            .decode_eof(&mut create_chunk(message_id, 0, 1, &payload))
+            .expect("a message exactly at the limit must be accepted")
+            .expect("the single-chunk message should be complete");
+
+        assert_eq!(frame, payload);
+    }
+
+    #[rstest]
+    #[tokio::test]
+    async fn ordinary_chunked_message_decodes_under_the_default_limits(
+        three_chunks_message: ([BytesMut; 3], String),
+    ) {
+        let (mut chunks, expected_message) = three_chunks_message;
+        let mut decoder = ChunkedGelfDecoder::default();
+
+        assert!(decoder.decode_eof(&mut chunks[0]).unwrap().is_none());
+        assert!(decoder.decode_eof(&mut chunks[1]).unwrap().is_none());
+        let frame = decoder.decode_eof(&mut chunks[2]).unwrap().unwrap();
+
+        assert_eq!(frame, expected_message);
+        assert!(decoder.state.lock().unwrap().is_empty());
+    }
 }

lib/vector-core/src/event/merge_state.rs

@@ -1,29 +1,50 @@
-use super::LogEvent;
+use super::{EstimatedJsonEncodedSizeOf, LogEvent};
 
 /// Encapsulates the inductive events merging algorithm.
 ///
-/// In the future, this might be extended by various counters (the number of
-/// events that contributed to the current merge event for instance, or the
-/// event size) to support circuit breaker logic.
+/// Tracks the size of everything merged so far so callers can apply circuit-breaker logic:
+/// a merge that is never terminated returns nothing to the pipeline, so bounded-channel
+/// backpressure cannot engage and the accumulator is otherwise unbounded.
 #[derive(Debug)]
 pub struct LogEventMergeState {
     /// Intermediate event we merge into.
     intermediate_merged_event: LogEvent,
+    /// Running total of the sizes of every event folded in so far.
+    merged_bytes: usize,
 }
 
 impl LogEventMergeState {
     /// Initialize the algorithm with a first (partial) event.
     pub fn new(first_partial_event: LogEvent) -> Self {
+        let merged_bytes = first_partial_event.estimated_json_encoded_size_of().get();
         Self {
             intermediate_merged_event: first_partial_event,
+            merged_bytes,
         }
     }
 
     /// Merge the incoming (partial) event in.
     pub fn merge_in_next_event(&mut self, incoming: LogEvent, fields: &[impl AsRef<str>]) {
+        // Measured on the incoming event rather than the accumulator, so this stays O(incoming)
+        // and adds no term to the merge's existing cost.
+        self.merged_bytes = self
+            .merged_bytes
+            .saturating_add(incoming.estimated_json_encoded_size_of().get());
         self.intermediate_merged_event.merge(incoming, fields);
     }
 
+    /// The total size of every event folded in so far.
+    pub const fn merged_bytes(&self) -> usize {
+        self.merged_bytes
+    }
+
+    /// Take the event accumulated so far, abandoning the merge.
+    ///
+    /// Used to flush an over-budget merge downstream instead of growing it further.
+    pub fn into_merged_event(self) -> LogEvent {
+        self.intermediate_merged_event
+    }
+
     /// Merge the final (non-partial) event in and return the resulting (merged)
     /// event.
     pub fn merge_in_final_event(
@@ -61,4 +82,39 @@ mod test {
             b"hello world"
         );
     }
+
+    #[test]
+    fn merged_bytes_grows_with_every_event_folded_in() {
+        let fields = vec!["message".to_string()];
+
+        let mut state = LogEventMergeState::new(log_event_with_message("hel"));
+        let after_first = state.merged_bytes();
+        assert!(after_first > 0, "the initial event must be accounted for");
+
+        state.merge_in_next_event(log_event_with_message("lo "), &fields);
+        let after_second = state.merged_bytes();
+        assert!(
+            after_second > after_first,
+            "folding an event in must grow the running total"
+        );
+
+        // Growth tracks payload size, which is what a caller budgets against.
+        state.merge_in_next_event(log_event_with_message(&"x".repeat(1024)), &fields);
+        assert!(state.merged_bytes() >= after_second + 1024);
+    }
+
+    #[test]
+    fn into_merged_event_returns_what_was_accumulated() {
+        let fields = vec!["message".to_string()];
+
+        let mut state = LogEventMergeState::new(log_event_with_message("hel"));
+        state.merge_in_next_event(log_event_with_message("lo"), &fields);
+
+        let event = state.into_merged_event();
+
+        assert_eq!(
+            event.get("message").unwrap().coerce_to_bytes().as_ref(),
+            b"hello"
+        );
+    }
 }