✨ Add configurable socket_read_limit

Sets the limit for how many bytes will be read from the socket at a
time.

When both socket_read_limit and max_response_size are non-nil, reads are
limited to the smaller of socket_read_limit and the remaining bytes for
that response.  When reading an IMAP `literal`, reads are also limited
to the remaining bytes in that `literal`

Please note that this only affects Net::IMAP's reads from the socket.
The underlying IO or Socket buffers and system calls will not be
directly affected by this limit.

This is complicated by the possibility that `gets(CRLF, limit)` could
stop in between a CR and a LF.

Author: nevans

lib/net/imap.rb

@@ -234,9 +234,9 @@ module Net
   #
   # Use paginated or limited versions of commands whenever possible.
   #
-  # Use Config#max_response_size to impose a limit on incoming server responses
-  # as they are being read.  <em>This is especially important for untrusted
-  # servers.</em>
+  # Use Config#max_response_size and Config#socket_read_limit to impose limits
+  # on incoming server responses as they are being read.  <em>This is especially
+  # important for untrusted servers.</em>
   #
   # Use #add_response_handler to handle responses after each one is received.
   # Use the +response_handlers+ argument to ::new to assign response handlers
@@ -876,11 +876,18 @@ class << self
     # The maximum allowed server response size, in bytes.
     # Delegates to {config.max_response_size}[rdoc-ref:Config#max_response_size].
 
+    ##
+    # :attr_accessor: socket_read_limit
+    # The limit for each socket read, in bytes.
+    # Delegates to {config.socket_read_limit}[rdoc-ref:Config#socket_read_limit].
+
     # :stopdoc:
     def open_timeout;           config.open_timeout            end
     def idle_response_timeout;  config.idle_response_timeout   end
     def max_response_size;      config.max_response_size       end
     def max_response_size=(val) config.max_response_size = val end
+    def socket_read_limit;      config.socket_read_limit       end
+    def socket_read_limit=(val) config.socket_read_limit = val end
     # :startdoc:
 
     # The hostname this client connected to

lib/net/imap/config.rb

@@ -302,6 +302,31 @@ def self.[](config)
       # * +0.5+: 512 MiB
       attr_accessor :max_response_size, type: Integer?
 
+      # Sets the limit for how many bytes will be read from the socket at a
+      # time.
+      #
+      # When both socket_read_limit and max_response_size are non-nil, reads are
+      # limited to the smaller of #socket_read_limit and the remaining bytes for
+      # that response.  When reading an IMAP +literal+, reads are also limited
+      # to the remaining bytes in that +literal+
+      #
+      # Please note that this only affects Net::IMAP's reads from its connection
+      # object, which is buffered.  The underlying buffers and IO system calls
+      # will not be directly affected by this read limit.
+      #
+      # Related: #max_response_size
+      #
+      # ==== Versioned Defaults
+      #
+      # Net::IMAP#socket_read_limit _and_ Net::IMAP#socket_read_limit= <em>were
+      # added in +v0.2.5+, +v0.3.9+, +v0.4.20+, and +v0.5.7+.</em>
+      #
+      # <em>Config option added in +v0.4.20+ and +v0.5.7+.</em>
+      #
+      # * original: +nil+ <em>(no limit)</em>
+      # * +0.6+: 16KiB
+      attr_accessor :socket_read_limit, type: Integer?
+
       # Controls the behavior of Net::IMAP#responses when called without any
       # arguments (+type+ or +block+).
       #
@@ -484,6 +509,7 @@ def defaults_hash
         responses_without_block: :frozen_dup,
         parser_use_deprecated_uidplus_data: false,
         parser_max_deprecated_uidplus_data_size: 0,
+        socket_read_limit: nil,
       ).freeze
 
       @global = default.new
@@ -523,6 +549,7 @@ def defaults_hash
       ).freeze
 
       version_defaults[0.7r] = Config[0.6r].dup.update(
+        socket_read_limit: 16 << 10, # 16 KiB
       ).freeze
 
       version_defaults[0.8r] = Config[0.7r].dup.update(

lib/net/imap/errors.rb

@@ -50,6 +50,19 @@ def response_size_msg
       end
     end
 
+    # Error raised when Config#socket_read_limit is less than or equal to zero.
+    class SocketReadLimitError < ResponseReadError
+      attr_reader :socket_read_limit
+
+      def initialize(msg = nil, *args, socket_read_limit: nil, **kwargs)
+        @socket_read_limit = socket_read_limit
+        msg ||= [
+          "Socket read limit", socket_read_limit, "<= 0",
+        ].compact.join(" ")
+        super(msg, *args, **kwargs)
+      end
+    end
+
     # Error raised when a response from the server is non-parsable.
     class ResponseParseError < Error
     end

lib/net/imap/response_reader.rb

@@ -88,9 +88,13 @@ def max_response_remaining!
         )
       end
 
-      # TODO: configurable socket_read_limit (currently hardcoded to 4KiB)
+      # Raises SocketReadLimitError if socket_read_limit is zero or negative.
       def socket_read_limit!
-        4096
+        read_limit = client.socket_read_limit
+        if read_limit && !read_limit.positive?
+          raise SocketReadLimitError.new(socket_read_limit: read_limit)
+        end
+        read_limit
       end
 
     end

test/net/imap/test_errors.rb

@@ -37,4 +37,18 @@ class IMAPErrorsTest < Net::IMAP::TestCase
                  "exceeds max_response_size (1200B)", err.message)
   end
 
+  test "SocketReadLimitError" do
+    err = Net::IMAP::SocketReadLimitError.new("manually set message")
+    assert_nil err.socket_read_limit
+    assert_equal "manually set message", err.message
+
+    err = Net::IMAP::SocketReadLimitError.new
+    assert_nil err.socket_read_limit
+    assert_equal "Socket read limit <= 0", err.message
+
+    err = Net::IMAP::SocketReadLimitError.new(socket_read_limit: -1)
+    assert_equal(-1, err.socket_read_limit)
+    assert_equal "Socket read limit -1 <= 0", err.message
+  end
+
 end

test/net/imap/test_imap_socket_read_limit.rb

@@ -0,0 +1,103 @@
+# frozen_string_literal: true
+
+require "net/imap"
+require "test/unit"
+require_relative "fake_server"
+
+class IMAPSocketReadLimitTest < Test::Unit::TestCase
+  include Net::IMAP::FakeServer::TestHelper
+
+  def setup
+    Net::IMAP.config.reset
+    @do_not_reverse_lookup = Socket.do_not_reverse_lookup
+    Socket.do_not_reverse_lookup = true
+    @threads = []
+  end
+
+  def teardown
+    if !@threads.empty?
+      assert_join_threads(@threads)
+    end
+  ensure
+    Socket.do_not_reverse_lookup = @do_not_reverse_lookup
+  end
+
+  def setup_various_responses_on_noop(server)
+    server.on("NOOP") do |resp|
+      # 10 byte response (prefix: "* ", suffix: "\r\n")
+      resp.untagged("OK 678")
+      # 20 byte response (prefix: "* ", suffix: "\r\n")
+      resp.untagged("OK 6789012345678")
+      # 21 byte response (prefix: "* ", suffix: "\r\n")
+      resp.untagged("OK 67890123456789")
+      # 22 byte response (prefix: "* ", suffix: "\r\n")
+      resp.untagged("OK 678901234567890")
+      # very large literal
+      resp.untagged("1 FETCH (BODY[] {12345}\r\n" + "a" * 12_345 + ")")
+      resp.done_ok
+    end
+  end
+
+  def setup_illegal_CRs_response_on_noop(server)
+    # response with many illegal CR chars (not followed by LF)
+    server.on("NOOP") do |resp|
+      resp.untagged("OK #{"\r" * 50} Oops!")
+      resp.done_ok
+    end
+  end
+
+  data " 1b",    1
+  data " 2b",    2
+  data " 9b",    9
+  data "10b",   10
+  data "20b",   20
+  data "21b",   21
+  data "22b",   22
+  data " 1KiB",  1 << 10
+  data "16KiB", 16 << 10
+  data "16MiB", 16 << 20
+  data "nil",   nil
+  test "#config.socket_read_limit" do |limit|
+    Net::IMAP.config.max_response_size = nil
+    Net::IMAP.config.socket_read_limit = limit
+    with_fake_server do |server, imap|
+      setup_various_responses_on_noop(server)
+      responses = []
+      imap.add_response_handler do
+        responses << _1.data if _1 in Net::IMAP::UntaggedResponse
+      end
+
+      imap.noop
+      assert_equal 5, responses.count
+      assert_equal "678",             responses[0].text
+      assert_equal "6789012345678",   responses[1].text
+      assert_equal "67890123456789",  responses[2].text
+      assert_equal "678901234567890", responses[3].text
+      assert_equal "a" * 12_345,      responses[4].message
+    end
+
+    with_fake_server(ignore_io_error: true) do |server, imap|
+      setup_illegal_CRs_response_on_noop(server)
+      # ResponseParseError means it was successfully sent to the parser
+      assert_raise(Net::IMAP::ResponseParseError) do
+        imap.noop
+      end
+    end
+  end
+
+  test "#config.socket_read_limit <= zero" do
+    with_fake_server(ignore_io_error: true) do |server, imap|
+      imap.config.socket_read_limit = 0
+      assert_raise(Net::IMAP::SocketReadLimitError) do
+        imap.noop
+      end
+    end
+    with_fake_server(ignore_io_error: true) do |server, imap|
+      imap.config.socket_read_limit = -1
+      assert_raise(Net::IMAP::SocketReadLimitError) do
+        imap.noop
+      end
+    end
+  end
+
+end

test/net/imap/test_response_reader.rb

@@ -9,6 +9,7 @@ class ResponseReaderTest < Net::IMAP::TestCase
   class FakeClient
     def config = @config ||= Net::IMAP.config.new
     def max_response_size = config.max_response_size
+    def socket_read_limit = config.socket_read_limit
   end
 
   def literal(str) = "{#{str.bytesize}}\r\n#{str}"
@@ -82,4 +83,25 @@ def literal(str) = "{#{str.bytesize}}\r\n#{str}"
     end
   end
 
+  test "#read_response_buffer with socket_read_limit" do
+    client = FakeClient.new
+    client.config.socket_read_limit = 10
+    aaaaaaaaa    = "a" * (20 << 10)
+    many_crs     = "\r" * 1000
+    many_crlfs   = "\r\n" * 100
+    mega_response = [
+      "* fake",
+      aaaaaaaaa,
+      literal(aaaaaaaaa),
+      literal(many_crlfs),
+      literal(literal(literal(many_crlfs))),
+      literal(many_crs),
+      many_crs,
+    ].join(" ")
+    io = StringIO.new(mega_response)
+    rcvr = Net::IMAP::ResponseReader.new(client, io)
+    assert_equal mega_response, rcvr.read_response_buffer.to_str
+    assert_equal "",           rcvr.read_response_buffer.to_str
+  end
+
 end