Add limit option to Net::BufferedIO#readuntil

lib/net/protocol.rb

@@ -79,6 +79,28 @@ class ProtoRetriableError    < ProtocolError; end
   ProtocRetryError = ProtoRetriableError
   # :startdoc:
 
+  ##
+  # ReadLimitExceeded, a subclass of ProtocolError, is raised if the
+  # terminator is not found within the byte limit given to
+  # Net::BufferedIO#readuntil.
+  #
+  # The limit is the largest result readuntil may return, counting the
+  # terminator itself, so a limit of 4 accepts "abc\n" and rejects
+  # "abcd\n". Unlike the limit of IO#gets it never truncates a result to
+  # fit. Either the whole thing comes back or this is raised, except
+  # under ignore_eof, which still returns what was buffered when the
+  # stream ended. The count is in bytes while the IO hands back binary
+  # strings, which every real one does.
+  #
+  # It bounds one call, not a connection. The unconsumed buffer can
+  # still run one BUFSIZE past the limit, and a peer sending endless
+  # short lines is not bounded at all.
+  #
+  # Nothing is consumed when this is raised, so the usual response is to
+  # close the connection rather than read on under a wider limit.
+
+  class ReadLimitExceeded < ProtocolError; end
+
   ##
   # OpenTimeout, a subclass of Timeout::Error, is raised if a connection cannot
   # be created within the open_timeout.
@@ -205,21 +227,31 @@ def read_all(dest = ''.b)
       dest
     end
 
-    def readuntil(terminator, ignore_eof = false)
+    def readuntil(terminator, ignore_eof = false, limit: nil)
+      unless limit.nil? || (Integer === limit && limit > 0)
+        # Integer === calls nothing on limit, and only an Integer is
+        # echoed back, so validation never runs the caller's code.
+        got = Integer === limit ? limit : "a non-Integer"
+        raise ArgumentError, "limit must be a positive Integer, got #{got}"
+      end
       offset = @rbuf_offset
       begin
         until idx = @rbuf.index(terminator, offset)
-          # Rewind by terminator.bytesize - 1 so that a terminator split
-          # across reads is not missed, however many reads it spans.
-          # @rbuf_offset is the floor for two reasons. A negative offset
-          # makes String#index search relative to the end of the buffer,
-          # skipping a match near its start. An offset below @rbuf_offset
-          # matches a terminator beginning inside bytes already returned
-          # to the caller, yielding a slice that does not end with one.
+          if limit && rbuf_size > limit
+            raise ReadLimitExceeded, "exceeded the #{limit} byte read limit"
+          end
+          # Rewind so a terminator split across reads is still found. The
+          # floor guards two things. String#index reads a negative offset
+          # as counting from the end, and an offset below @rbuf_offset
+          # matches inside bytes already returned.
           offset = [@rbuf.bytesize - terminator.bytesize + 1, @rbuf_offset].max
           rbuf_fill
         end
-        return rbuf_consume(idx + terminator.bytesize - @rbuf_offset)
+        len = idx + terminator.bytesize - @rbuf_offset
+        if limit && len > limit
+          raise ReadLimitExceeded, "exceeded the #{limit} byte read limit"
+        end
+        return rbuf_consume(len)
       rescue EOFError
         raise unless ignore_eof
         return rbuf_consume

test/net/protocol/test_protocol.rb

@@ -65,6 +65,115 @@ def test_readuntil
     end
   end
 
+  def test_readuntil_limit
+    io = Net::BufferedIO.new(StringIO.new("123\n45678\n".dup))
+    assert_equal "123\n", io.readuntil("\n", limit: 4)
+    assert_raise(Net::ReadLimitExceeded) { io.readuntil("\n", limit: 4) }
+  end
+
+  # The limit measures the result, not the position of the terminator in
+  # the buffer, so bytes consumed by an earlier read must not count
+  # against it.
+  def test_readuntil_limit_ignores_already_consumed_bytes
+    io = Net::BufferedIO.new(StringIO.new("123\n45678\n".dup))
+    assert_equal "123\n", io.readuntil("\n", limit: 4)
+    assert_equal "45678\n", io.readuntil("\n", limit: 6)
+  end
+
+  def test_readuntil_limit_is_a_protocol_error
+    assert_operator Net::ReadLimitExceeded, :<, Net::ProtocolError
+  end
+
+  # Which of the two checks fires is decided by how the peer split its
+  # writes, so both have to report the same thing.
+  def test_readuntil_limit_message_does_not_depend_on_chunking
+    whole = Net::BufferedIO.new(StringIO.new("45678\n".dup))
+    split = Net::BufferedIO.new(FakeReadPartialIO.new(["45678", "\n"]))
+    messages = [whole, split].map do |io|
+      assert_raise(Net::ReadLimitExceeded) { io.readuntil("\n", limit: 4) }.message
+    end
+    assert_equal messages.first, messages.last
+    assert_match(/\b4\b/, messages.first)
+    assert_match(/limit/, messages.first)
+  end
+
+  def test_readuntil_limit_after_a_long_earlier_read
+    io = Net::BufferedIO.new(FakeReadPartialIO.new(["aaaaaaaaaa\nbc", "\n"]))
+    assert_equal "aaaaaaaaaa\n", io.readuntil("\n", limit: 11)
+    assert_equal "bc\n", io.readuntil("\n", limit: 3)
+  end
+
+  def test_readuntil_limit_rejects_values_that_are_not_a_positive_integer
+    { 0 => "0", -1 => "-1", false => "a non-Integer",
+      4.5 => "a non-Integer", "4" => "a non-Integer" }.each do |limit, expected|
+      io = Net::BufferedIO.new(StringIO.new("123\n".dup))
+      e = assert_raise(ArgumentError, "limit: #{limit.inspect}") do
+        io.readuntil("\n", limit: limit)
+      end
+      assert_equal "limit must be a positive Integer, got #{expected}", e.message
+    end
+
+    io = Net::BufferedIO.new(StringIO.new("123\n".dup))
+    assert_equal "123\n", io.readuntil("\n", limit: nil)
+  end
+
+  def test_readuntil_limit_counts_the_terminator
+    io = Net::BufferedIO.new(StringIO.new("1234\n".dup))
+    assert_raise(Net::ReadLimitExceeded) { io.readuntil("\n", limit: 4) }
+
+    io = Net::BufferedIO.new(StringIO.new("1234\n".dup))
+    assert_equal "1234\n", io.readuntil("\n", limit: 5)
+  end
+
+  def test_readuntil_limit_consumes_nothing_when_it_raises
+    io = Net::BufferedIO.new(StringIO.new("45678\nrest\n".dup))
+    assert_raise(Net::ReadLimitExceeded) { io.readuntil("\n", limit: 4) }
+    assert_equal "45678\n", io.readuntil("\n", limit: 6)
+    assert_equal "rest\n", io.readuntil("\n")
+  end
+
+  def test_readuntil_limit_ignore_eof
+    io = Net::BufferedIO.new(StringIO.new("abc".dup))
+    assert_equal "abc", io.readuntil("\n", true, limit: 10)
+  end
+
+  # The EOF path returns the buffer without consulting the limit, so
+  # only the loop's earlier check keeps it inside.
+  def test_readuntil_limit_bounds_what_ignore_eof_returns_at_eof
+    io = Net::BufferedIO.new(FakeReadPartialIO.new(["abcde"]))
+    assert_equal "abcde", io.readuntil("\n", true, limit: 5)
+
+    io = Net::BufferedIO.new(FakeReadPartialIO.new(["abcdef"]))
+    assert_raise(Net::ReadLimitExceeded) { io.readuntil("\n", true, limit: 5) }
+  end
+
+  def test_readuntil_limit_applies_with_ignore_eof
+    io = Net::BufferedIO.new(StringIO.new("abcdefghij".dup))
+    assert_raise(Net::ReadLimitExceeded) { io.readuntil("\n", true, limit: 5) }
+  end
+
+  # Never yields the terminator. Capping the reads makes a regression in
+  # the limit check fail instead of running the CI host out of memory.
+  class EndlessIO
+    MAX_READS = 2
+
+    def initialize
+      @reads = 0
+    end
+
+    def read_nonblock(size, buf = nil, exception: false)
+      @reads += 1
+      raise "readuntil ignored its limit: #{@reads} reads" if @reads > MAX_READS
+      s = ("a" * size).b
+      buf ? buf.replace(s) : s
+    end
+  end
+
+  def test_readuntil_limit_endless_stream
+    io = Net::BufferedIO.new(EndlessIO.new)
+    assert_raise(Net::ReadLimitExceeded) { io.readuntil("\n", limit: 1024) }
+  end
+
   def test_write0_multibyte
     mockio = create_mockio(max: 1)
     io = Net::BufferedIO.new(mockio)
@@ -130,9 +239,8 @@ def test_write0_timeout_multi2
 
   class FakeReadPartialIO
     def initialize(chunks)
-      # Binary, like the bytes a real IO hands back. String#b also copies,
-      # which matters because rbuf_fill clears a string read_nonblock
-      # returns without having been handed it as the buffer.
+      # Binary, like a real IO. String#b also copies, which matters
+      # because rbuf_fill clears a string it was not handed as the buffer.
       @chunks = chunks.map(&:b)
     end
 
@@ -162,6 +270,14 @@ def test_shareable_buffer_leak # https://github.com/ruby/net-protocol/pull/19
     assert_equal expected_chunks, actual_chunks
   end
 
+  def test_readuntil_limit_with_a_terminator_spanning_chunks
+    io = Net::BufferedIO.new(FakeReadPartialIO.new(["abc\r", "\ndef\r\n"]))
+    assert_equal "abc\r\n", io.readuntil("\r\n", limit: 5)
+
+    io = Net::BufferedIO.new(FakeReadPartialIO.new(["abc\r", "\ndef\r\n"]))
+    assert_raise(Net::ReadLimitExceeded) { io.readuntil("\r\n", limit: 4) }
+  end
+
   def test_readuntil_terminator_spanning_chunks # https://github.com/ruby/net-protocol/pull/66
     fake_io = FakeReadPartialIO.new(["abc\r", "\ndef\r\n"])
     io = Net::BufferedIO.new(fake_io)
@@ -178,17 +294,13 @@ def test_readuntil_terminator_spanning_more_than_two_chunks # https://github.com
   def test_readuntil_clamps_a_negative_rewind # https://github.com/ruby/net-protocol/pull/66
     fake_io = FakeReadPartialIO.new(["ab\n"])
     io = Net::BufferedIO.new(fake_io)
-    # Any buffer shorter than the terminator drives the rewind below zero,
-    # and String#index reads a negative offset as counting from the end.
     assert_equal "ab", io.readuntil("ab")
   end
 
   def test_readuntil_does_not_rewind_into_consumed_bytes # https://github.com/ruby/net-protocol/pull/66
     fake_io = FakeReadPartialIO.new(["ab\r\n\r", "\nc"])
     io = Net::BufferedIO.new(fake_io)
     assert_equal "ab\r", io.readuntil("\r")
-    # The terminator is longer than what is left unconsumed, so the rewind
-    # would reach back into the bytes readuntil already returned.
     assert_raise(EOFError) { io.readuntil("\r\n\r\n") }
   end