✨ Limit max response size to 512MiB (hard-coded)

_Please note:_ this only limits the size per response.  It does _not_
limit how many unhandled responses may be stored on the responses hash.

Author: nevans

lib/net/imap/errors.rb

@@ -11,6 +11,39 @@ class Error < StandardError
     class DataFormatError < Error
     end
 
+    # Error raised when the socket cannot be read, due to a Config limit.
+    class ResponseReadError < Error
+    end
+
+    # Error raised when a response is larger than IMAP#max_response_size.
+    class ResponseTooLargeError < ResponseReadError
+      attr_reader :bytes_read, :literal_size
+      attr_reader :max_response_size
+
+      def initialize(msg = nil, *args,
+                     bytes_read:        nil,
+                     literal_size:      nil,
+                     max_response_size: nil,
+                     **kwargs)
+        @bytes_read        = bytes_read
+        @literal_size      = literal_size
+        @max_response_size = max_response_size
+        msg ||= [
+          "Response size", response_size_msg, "exceeds max_response_size",
+          max_response_size && "(#{max_response_size}B)",
+        ].compact.join(" ")
+        super(msg, *args, **kwargs)
+      end
+
+      private
+
+      def response_size_msg
+        if bytes_read && literal_size
+          "(#{bytes_read}B read + #{literal_size}B literal)"
+        end
+      end
+    end
+
     # Error raised when a response from the server is non-parsable.
     class ResponseParseError < Error
     end

lib/net/imap/response_reader.rb

@@ -28,19 +28,46 @@ def read_response_buffer
 
       attr_reader :buff, :literal_size
 
+      def bytes_read          = buff.bytesize
+      def empty?              = buff.empty?
+      def done?               = line_done? && !get_literal_size
+      def line_done?          = buff.end_with?(CRLF)
       def get_literal_size; /\{(\d+)\}\r\n\z/n =~ buff && $1.to_i end
 
       def read_line
-        buff << (@sock.gets(CRLF) or throw :eof)
+        buff << (@sock.gets(CRLF, read_limit) or throw :eof)
+        max_response_remaining! unless line_done?
       end
 
       def read_literal
+        # check before allocating memory for literal
+        max_response_remaining!
         literal = String.new(capacity: literal_size)
-        buff << (@sock.read(literal_size, literal) or throw :eof)
+        buff << (@sock.read(read_limit(literal_size), literal) or throw :eof)
       ensure
         @literal_size = nil
       end
 
+      def read_limit(limit = nil)
+        [limit, max_response_remaining!].compact.min
+      end
+
+      def max_response_size      = 512 << 20 # TODO: Config#max_response_size
+      def max_response_remaining = max_response_size &.- bytes_read
+      def response_too_large?    = max_response_size &.< min_response_size
+      def min_response_size      = bytes_read + min_response_remaining
+
+      def min_response_remaining
+        empty? ? 3 : done? ? 0 : (literal_size || 0) + 2
+      end
+
+      def max_response_remaining!
+        return max_response_remaining unless response_too_large?
+        raise ResponseTooLargeError.new(
+          max_response_size:, bytes_read:, literal_size:,
+        )
+      end
+
     end
   end
 end

test/net/imap/test_errors.rb

@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+require "net/imap"
+require "test/unit"
+
+class IMAPErrorsTest < Test::Unit::TestCase
+
+  test "ResponseTooLargeError" do
+    err = Net::IMAP::ResponseTooLargeError.new
+    assert_nil err.bytes_read
+    assert_nil err.literal_size
+    assert_nil err.max_response_size
+
+    err = Net::IMAP::ResponseTooLargeError.new("manually set message")
+    assert_equal "manually set message", err.message
+    assert_nil err.bytes_read
+    assert_nil err.literal_size
+    assert_nil err.max_response_size
+
+    err = Net::IMAP::ResponseTooLargeError.new(max_response_size: 1024)
+    assert_equal "Response size exceeds max_response_size (1024B)", err.message
+    assert_nil err.bytes_read
+    assert_nil err.literal_size
+    assert_equal 1024, err.max_response_size
+
+    err = Net::IMAP::ResponseTooLargeError.new(bytes_read:        1200,
+                                               max_response_size: 1200)
+    assert_equal 1200, err.bytes_read
+    assert_equal "Response size exceeds max_response_size (1200B)", err.message
+
+    err = Net::IMAP::ResponseTooLargeError.new(bytes_read:        800,
+                                               literal_size:      1000,
+                                               max_response_size: 1200)
+    assert_equal  800, err.bytes_read
+    assert_equal 1000, err.literal_size
+    assert_equal("Response size (800B read + 1000B literal) " \
+                 "exceeds max_response_size (1200B)", err.message)
+  end
+
+end

test/net/imap/test_response_reader.rb

@@ -49,4 +49,27 @@ def literal(str) "{#{str.bytesize}}\r\n#{str}" end
     assert_equal "",           rcvr.read_response_buffer.to_str
   end
 
+  class LimitedResponseReader < Net::IMAP::ResponseReader
+    attr_reader :max_response_size
+    def initialize(*args, max_response_size:)
+      super(*args)
+      @max_response_size = max_response_size
+    end
+  end
+
+  test "#read_response_buffer with max_response_size" do
+    client = FakeClient.new
+    max_response_size = 10
+    under = "+ 3456\r\n"
+    exact = "+ 345678\r\n"
+    over  = "+ 3456789\r\n"
+    io = StringIO.new([under, exact, over].join)
+    rcvr = LimitedResponseReader.new(client, io, max_response_size:)
+    assert_equal under, rcvr.read_response_buffer.to_str
+    assert_equal exact, rcvr.read_response_buffer.to_str
+    assert_raise Net::IMAP::ResponseTooLargeError do
+      rcvr.read_response_buffer
+    end
+  end
+
 end