Fix up buffer treatment of stream parsing

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