Fix up buffer treatment of stream parsing

A couple of things to fix, notably that we relied on the stream
returning strings and relied on it not lying about how much was
being written. Now we no longer trust either and we verify
instead.

Co-authored-by: Sutou Kouhei <kou@clear-code.com>

Author: kddnewton

ext/prism/extension.c

@@ -1067,26 +1067,64 @@ parse_stream_eof(void *stream) {
 }
 
 /**
- * An implementation of fgets that is suitable for use with Ruby IO objects.
+ * The largest number of bytes a single character can occupy in any encoding
+ * that Ruby supports (CESU-8). `IO#gets(limit)` treats its argument as a soft
+ * limit: to avoid splitting a multi-byte character it may return up to
+ * `MAX_ENC_LEN - 1` more bytes than requested. Reserving `MAX_ENC_LEN` bytes of
+ * headroom therefore leaves room for both that overshoot and the terminating
+ * NUL byte.
+ *
+ * The value can be re-derived from the Ruby source tree with:
+ *
+ *   grep -Erh "max (enc|byte) length" enc | grep -Eo '[0-9]+' | sort | tail -n1
+ */
+#define MAX_ENC_LEN 6
+
+/**
+ * An implementation of fgets that is suitable for use with Ruby IO objects. As
+ * required by pm_source_stream_fgets_t, this never writes more than `size`
+ * bytes into `string` (including the terminating NUL byte).
  */
 static char *
 parse_stream_fgets(char *string, int size, void *stream) {
-    RUBY_ASSERT(size > 0);
+    RUBY_ASSERT(size > MAX_ENC_LEN);
+
+    /*
+     * Request fewer bytes than the buffer can hold. `gets` may return more
+     * bytes than requested when the limit falls in the middle of a multi-byte
+     * character, and the reserved headroom guarantees the result still fits.
+     */
+    VALUE line = rb_funcall((VALUE) stream, rb_intern("gets"), 1, INT2FIX(size - MAX_ENC_LEN));
 
-    VALUE line = rb_funcall((VALUE) stream, rb_intern("gets"), 1, INT2FIX(size - 1));
+    /*
+     * A well-behaved stream returns a String, or nil at EOF. Coerce anything
+     * else to nil so that we never treat a non-String as a byte buffer.
+     */
+    line = rb_check_string_type(line);
     if (NIL_P(line)) {
         return NULL;
     }
 
     const char *cstr = RSTRING_PTR(line);
     long length = RSTRING_LEN(line);
 
-    memcpy(string, cstr, length);
+    /*
+     * Defensively clamp the copy. A misbehaving `gets` may ignore the limit
+     * entirely and return an arbitrarily long string; we must never write past
+     * the caller's buffer. One byte is reserved for the NUL terminator.
+     */
+    if (length > (long) (size - 1)) {
+        length = (long) (size - 1);
+    }
+
+    memcpy(string, cstr, (size_t) length);
     string[length] = '\0';
 
     return string;
 }
 
+#undef MAX_ENC_LEN
+
 /**
  * :markup: markdown
  * call-seq:

include/prism/source.h

@@ -24,6 +24,13 @@ typedef struct pm_source_t pm_source_t;
  * This function is used to retrieve a line of input from a stream. It closely
  * mirrors that of fgets so that fgets can be used as the default
  * implementation.
+ *
+ * As with fgets, implementations MUST write at most `size` bytes into
+ * `string`, including the terminating NUL byte (i.e. at most `size - 1` bytes
+ * of data followed by a '\0'). The caller relies on this bound for memory
+ * safety, so an implementation that reads from a source which may return more
+ * data than requested (for example a multi-byte-aware `gets`) is responsible
+ * for clamping the amount it copies. Returns `string`, or NULL on EOF.
  */
 typedef char * (pm_source_stream_fgets_t)(char *string, int size, void *stream);
 

lib/prism/ffi.rb

@@ -294,11 +294,23 @@ def parse_file(filepath, **options)
     # Mirror the Prism.parse_stream API by using the serialization API.
     def parse_stream(stream, **options)
       LibRubyParser::PrismBuffer.with do |buffer|
+        # The largest number of bytes a single character can occupy in any
+        # encoding Ruby supports. IO#gets(limit) may return up to
+        # (max_enc_len - 1) bytes more than requested to avoid splitting a
+        # multi-byte character, so we reserve that much headroom (plus the NUL
+        # terminator) to guarantee the result fits in the caller's buffer. This
+        # mirrors MAX_ENC_LEN in ext/prism/extension.c.
+        max_enc_len = 6
+
         source = +""
         callback = -> (string, size, _) {
-          raise "Expected size to be >= 0, got: #{size}" if size <= 0
+          raise "Expected size to be > #{max_enc_len}, got: #{size}" if size <= max_enc_len
 
-          if !(line = stream.gets(size - 1)).nil?
+          line = String.try_convert(stream.gets(size - max_enc_len))
+          if !line.nil?
+            # A misbehaving `gets` may ignore the limit; never write past the
+            # buffer (one byte is reserved for the NUL terminator).
+            line = line.byteslice(0, size - 1) if line.bytesize > size - 1
             source << line
             string.write_string("#{line}\x00", line.bytesize + 1)
           end

src/source.c

@@ -368,25 +368,50 @@ pm_source_stream_read(pm_source_t *source) {
 #define LINE_SIZE 4096
     char line[LINE_SIZE];
 
-    while (memset(line, '\n', LINE_SIZE), source->stream.fgets(line, LINE_SIZE, source->stream.stream) != NULL) {
+    /*
+     * A chunk read from the stream may legitimately contain embedded NUL bytes,
+     * so the NUL terminator written by fgets cannot be located with strlen.
+     * Instead we prefill the buffer with a non-NUL sentinel before each read;
+     * the terminator is then the only NUL at or after the data, so it can be
+     * found by scanning from the end. The value only needs to differ from '\0'.
+     */
+#define LINE_FILL 0xff
+
+    while (memset(line, LINE_FILL, LINE_SIZE), source->stream.fgets(line, LINE_SIZE, source->stream.stream) != NULL) {
+        /*
+         * Locate the NUL terminator written by fgets (the last NUL in the
+         * buffer) to recover the length of the chunk, then drop the terminator.
+         */
         size_t length = LINE_SIZE;
-        while (length > 0 && line[length - 1] == '\n') length--;
-
-        if (length == LINE_SIZE) {
-            /*
-             * If we read a line that is the maximum size and it doesn't end
-             * with a newline, then we'll just append it to the buffer and
-             * continue reading.
-             */
-            length--;
-            pm_buffer_append_string(buffer, line, length);
-            continue;
+        while (length > 0 && line[length - 1] != '\0') length--;
+        if (length > 0) length--;
+
+        /*
+         * An empty chunk means the stream returned an empty string instead of
+         * signaling EOF (a well-behaved stream never does this). We can't make
+         * progress on it, so stop reading rather than spinning forever.
+         */
+        if (length == 0) {
+            break;
         }
 
-        /* Append the line to the buffer. */
-        length--;
+        /* Append the chunk to the buffer. */
         pm_buffer_append_string(buffer, line, length);
 
+        bool newline = (line[length - 1] == '\n');
+
+        /*
+         * A chunk with no trailing newline is either a line longer than the
+         * read buffer (keep reading the rest of it) or the final line at EOF.
+         * This is the only place the stream's EOF state changes what we do
+         * next, so it is the only place we ask the stream whether it has hit
+         * EOF. A newline-terminated line needs no such check: if it happens to
+         * be the last line, the next fgets returns NULL and ends the loop.
+         */
+        if (!newline && !source->stream.feof(source->stream.stream)) {
+            continue;
+        }
+
         /*
          * Check if the line matches the __END__ marker. If it does, then stop
          * reading and return false. In most circumstances, this means we should
@@ -418,14 +443,16 @@ pm_source_stream_read(pm_source_t *source) {
         }
 
         /*
-         * All data should be read via gets. If the string returned by gets
-         * _doesn't_ end with a newline, then we assume we hit EOF condition.
+         * A chunk that reached here without a trailing newline is the final
+         * line at EOF (the check above would have continued otherwise), so
+         * there is nothing more to read.
          */
-        if (source->stream.feof(source->stream.stream)) {
+        if (!newline) {
             break;
         }
     }
 
+#undef LINE_FILL
 #undef LINE_SIZE
 
     source->stream.eof = true;

test/prism/api/parse_stream_test.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require_relative "../test_helper"
+require "timeout"
 
 module Prism
   class ParseStreamTest < TestCase
@@ -43,6 +44,18 @@ def test___END__
       assert_equal "5 + 6\n", io.read
     end
 
+    # __END__ as the final line with no trailing newline must still be detected
+    # as the terminator. This exercises the EOF check for a chunk that does not
+    # end in a newline, distinct from a line longer than the internal read
+    # buffer (which also has no trailing newline but is not the end of input).
+    def test___END___without_trailing_newline
+      io = StringIO.new("1 + 2\n__END__")
+      result = Prism.parse_stream(io)
+
+      assert result.success?
+      assert_equal 1, result.value.statements.body.length
+    end
+
     def test_false___END___in_string
       io = StringIO.new(<<~RUBY)
         1 + 2
@@ -114,5 +127,70 @@ def test_nul_bytes
       assert result.success?
       assert_equal 3, result.value.statements.body.length
     end
+
+    # `IO#gets(limit)` will not split a multi-byte character, so it can return
+    # more bytes than requested when the limit falls in the middle of one. When
+    # that character straddles the internal read buffer boundary, the old code
+    # wrote past the buffer and dropped the overflowing bytes, corrupting the
+    # content (and overflowing the stack). Sweep offsets around the 4096-byte
+    # boundary with both 3- and 4-byte characters so the straddle is hit
+    # regardless of the exact internal read size.
+    def test_multibyte_on_read_boundary
+      ["あ", "\u{1F600}"].each do |char|
+        (4080..4100).each do |prefix|
+          body = ("a" * prefix) + char
+          result = Prism.parse_stream(StringIO.new("\"#{body}\""))
+
+          assert result.success?, "parse failed at prefix=#{prefix} char=#{char.dump}"
+          assert_equal body, result.value.statements.body[0].content, "content mismatch at prefix=#{prefix} char=#{char.dump}"
+        end
+      end
+    end
+
+    # A misbehaving stream whose `gets` ignores its limit and returns an
+    # arbitrarily long string must not overflow the internal buffer. The old
+    # code copied the full returned length into a fixed 4096-byte buffer.
+    def test_gets_exceeding_limit
+      stream = Object.new
+      def stream.gets(limit = nil)
+        return nil if defined?(@done) && @done
+        @done = true
+        ("x" * 100_000) + "\n"
+      end
+      def stream.eof?; defined?(@done) && @done; end
+
+      result = Prism.parse_stream(stream)
+      assert result.success?
+    end
+
+    # A misbehaving stream whose `gets` returns a non-String must not be read as
+    # a byte buffer. The old code called RSTRING_PTR on whatever was returned.
+    def test_gets_returning_non_string
+      stream = Object.new
+      def stream.gets(limit = nil)
+        return nil if defined?(@done) && @done
+        @done = true
+        1234
+      end
+      def stream.eof?; defined?(@done) && @done; end
+
+      assert_nothing_raised do
+        Prism.parse_stream(stream)
+      end
+    end
+
+    # A misbehaving stream whose `gets` returns an empty string instead of nil
+    # while never reporting EOF must not loop forever. A well-behaved stream
+    # returns nil at EOF, so an empty chunk is treated as the end of input.
+    def test_gets_returning_empty_string
+      stream = Object.new
+      def stream.gets(limit = nil); ""; end
+      def stream.eof?; false; end
+
+      result = Timeout.timeout(10) { Prism.parse_stream(stream) }
+      assert result.success?
+    rescue Timeout::Error
+      flunk "Prism.parse_stream looped forever on a stream returning empty strings"
+    end
   end
 end

test/prism/newline_test.rb

@@ -12,6 +12,7 @@ class NewlineTest < TestCase
       regexp_test.rb
       test_helper.rb
       unescape_test.rb
+      api/parse_stream_test.rb
       encoding/regular_expression_encoding_test.rb
       encoding/string_encoding_test.rb
       result/breadth_first_search_test.rb