lib/aws-eventstream/encoder.rb in aws-eventstream-1.0.1 vs lib/aws-eventstream/encoder.rb in aws-eventstream-1.0.2

@@ -37,10 +37,16 @@
 
       # bytes of total overhead in a message, including prelude
       # and 4 bytes total message crc checksum
       OVERHEAD_LENGTH = 16
 
+      # Maximum header length allowed (after encode) 128kb
+      MAX_HEADERS_LENGTH = 131072
+
+      # Maximum payload length allowed (after encode) 16mb
+      MAX_PAYLOAD_LENGTH = 16777216
+
       # Encodes Aws::EventStream::Message to output IO when
       #   provided, else return the encoded binary string
       #
       # @param [Aws::EventStream::Message] message
       #
@@ -60,17 +66,24 @@
         else
           encoded
         end
       end
 
-      private
-
+      # Encodes an Aws::EventStream::Message
+      #   into Aws::EventStream::BytesBuffer
+      #
+      # @param [Aws::EventStream::Message] msg
+      #
+      # @return [Aws::EventStream::BytesBuffer]
       def encode_message(message)
         # create context buffer with encode headers
         ctx_buffer = encode_headers(message)
         headers_len = ctx_buffer.bytesize
         # encode payload
+        if message.payload.length > MAX_PAYLOAD_LENGTH
+          raise Aws::EventStream::Errors::EventPayloadLengthExceedError.new
+        end
         ctx_buffer << message.payload.read
         total_len = ctx_buffer.bytesize + OVERHEAD_LENGTH
 
         # create message buffer with prelude section
         buffer = prelude(total_len, headers_len)
@@ -83,10 +96,16 @@
         # write buffered message to io
         buffer.rewind
         buffer
       end
 
+      # Encodes headers part of an Aws::EventStream::Message
+      #   into Aws::EventStream::BytesBuffer
+      #
+      # @param [Aws::EventStream::Message] msg
+      #
+      # @return [Aws::EventStream::BytesBuffer]
       def encode_headers(msg)
         buffer = BytesBuffer.new('')
         msg.headers.each do |k, v|
           # header key
           buffer << pack_uint8(k.bytesize)
@@ -99,11 +118,16 @@
           next if !!pattern == pattern
           buffer << pack_uint16(v.value.bytesize) unless val_len
           pattern ? buffer << [v.value].pack(pattern) :
             buffer << v.value
         end
+        if buffer.bytesize > MAX_HEADERS_LENGTH
+          raise Aws::EventStream::Errors::EventHeadersLengthExceedError.new
+        end
         buffer
       end
+
+      private
 
       def prelude(total_len, headers_len)
         BytesBuffer.new(pack_uint32([
           total_len,
           headers_len,