Add back validation for the eventstream payload and header (#3208)

* Add back an eventstream payload and header limit but with 24MB

* Increase the limit

* Fix limit in error message

---------

Co-authored-by: Richard Wang <[email protected]>

Author: mullermp

gems/aws-eventstream/CHANGELOG.md

@@ -1,6 +1,8 @@
 Unreleased Changes
 ------------------
 
+* Issue - Add back event stream max payload size and headers length checks with a 24MB limit.
+
 1.3.1 (2025-02-13)
 ------------------
 

gems/aws-eventstream/lib/aws-eventstream/encoder.rb

@@ -41,6 +41,12 @@ class Encoder
       # and 4 bytes total message crc checksum
       OVERHEAD_LENGTH = 16
 
+      # Maximum header length allowed (after encode) 128kb
+      MAX_HEADERS_LENGTH = 1024 * 128
+
+      # Maximum payload length allowed (after encode) 24mb
+      MAX_PAYLOAD_LENGTH = 1024 * 1024 * 24
+
       # Encodes Aws::EventStream::Message to output IO when
       #   provided, else return the encoded binary string
       #
@@ -75,6 +81,9 @@ def encode_message(message)
         encoded_header = encode_headers(message)
         header_length = encoded_header.bytesize
         # encode payload
+        if message.payload.length > MAX_PAYLOAD_LENGTH
+          raise Aws::EventStream::Errors::EventPayloadLengthExceedError.new
+        end
         encoded_payload = message.payload.read
         total_length = header_length + encoded_payload.bytesize + OVERHEAD_LENGTH
 
@@ -115,7 +124,10 @@ def encode_headers(message)
             pattern ? [value.value].pack(pattern) : value.value,
           ].pack('a*a*a*')
         end
-        header_entries.join
+        header_entries.join.tap do |encoded_header|
+          break encoded_header if encoded_header.bytesize <= MAX_HEADERS_LENGTH
+          raise Aws::EventStream::Errors::EventHeadersLengthExceedError.new
+        end
       end
 
       private

gems/aws-eventstream/lib/aws-eventstream/errors.rb

@@ -32,6 +32,18 @@ def initialize(*args)
         end
       end
 
+      class EventPayloadLengthExceedError < RuntimeError
+        def initialize(*args)
+          super("Payload length of a message should be under 24mb.")
+        end
+      end
+
+      class EventHeadersLengthExceedError < RuntimeError
+        def initialize(*args)
+          super("Encoded headers length of a message should be under 128kb.")
+        end
+      end
+
     end
   end
 end

gems/aws-eventstream/spec/encoder_spec.rb

@@ -25,25 +25,33 @@ module EventStream
             expect(test_io.string.freeze).to eq(expectation)
           end
         end
+      end
+
+      describe '#encode error' do
 
-        it 'encodes large payloads' do
+        it 'raises an error when payload exceeds' do
+          payload = double('payload', :length => 24 * 1024 * 1024 + 1)
           message = Aws::EventStream::Message.new(
             headers: {},
-            payload: StringIO.new('.' * 24 * 1024 * 1024)
+            payload: payload
           )
-          Encoder.new.encode(message)
+          expect {
+            Encoder.new.encode(message)
+          }.to raise_error(Aws::EventStream::Errors::EventPayloadLengthExceedError)
         end
 
-        it 'encodes long headers' do
+        it 'raises an error when encoded headers exceeds' do
           headers = {}
           headers['foo'] = Aws::EventStream::HeaderValue.new(
-            value: '*' * 131_073, type: 'string'
+            value: '*' * 131073, type: 'string'
           )
           message = Aws::EventStream::Message.new(
             headers: headers,
             payload: StringIO.new
           )
-          Encoder.new.encode(message)
+          expect {
+            Encoder.new.encode(message)
+          }.to raise_error(Aws::EventStream::Errors::EventHeadersLengthExceedError)
 
         end