Fix a bug that Prism.parse_stream may write too much data to internal buffer

`pm_source_stream_read()` has 4096 bytes internal buffer to read a
line:
https://github.com/ruby/prism/blob/39fb41f06e5f4ccbe1783122fa8a4470e09dfcfe/src/source.c#L368-L369

It uses `gets(4095)` to read the next chunk to the buffer:
https://github.com/ruby/prism/blob/39fb41f06e5f4ccbe1783122fa8a4470e09dfcfe/ext/prism/extension.c#L1076

But `gets(4095)` may read 4095 or more data when the character at 4095
is a multi-byte character. In the case, `parse_stream_fgets()` writes
4097 or more data to the buffer.

We can avoid it by increasing internal buffer size to `4096 + 6` where
`6` is the max character length in all available encodings. We can
reduce it by using `stream.external_encoding` and `rb_enc_mbmaxlen()`
but it may slow down. So I used constant `6` in this change.

This also changes `pm_source_stream_fgets_t`'s return type to
`int` (read bytes) from `char *` (read data) to reduce needless length
check by caller.

Author: kou

ext/prism/extension.c

@@ -1069,22 +1069,21 @@ parse_stream_eof(void *stream) {
 /**
  * An implementation of fgets that is suitable for use with Ruby IO objects.
  */
-static char *
+static int
 parse_stream_fgets(char *string, int size, void *stream) {
     RUBY_ASSERT(size > 0);
 
     VALUE line = rb_funcall((VALUE) stream, rb_intern("gets"), 1, INT2FIX(size - 1));
     if (NIL_P(line)) {
-        return NULL;
+        return 0;
     }
 
     const char *cstr = RSTRING_PTR(line);
     long length = RSTRING_LEN(line);
 
     memcpy(string, cstr, length);
-    string[length] = '\0';
 
-    return string;
+    return length;
 }
 
 /**

include/prism/source.h

@@ -25,7 +25,7 @@ typedef struct pm_source_t pm_source_t;
  * mirrors that of fgets so that fgets can be used as the default
  * implementation.
  */
-typedef char * (pm_source_stream_fgets_t)(char *string, int size, void *stream);
+typedef int (pm_source_stream_fgets_t)(char *string, int size, void *stream);
 
 /**
  * This function is used to check whether a stream is at EOF. It closely mirrors

lib/prism/ffi.rb

@@ -88,7 +88,7 @@ def self.load_exported_functions_from(header, *functions, callbacks)
       raise "Could not find functions #{functions.inspect}" unless functions.empty?
     end
 
-    callback :pm_source_stream_fgets_t, [:pointer, :int, :pointer], :pointer
+    callback :pm_source_stream_fgets_t, [:pointer, :int, :pointer], :int
     callback :pm_source_stream_feof_t, [:pointer], :int
     pm_source_init_result_values = %i[PM_SOURCE_INIT_SUCCESS PM_SOURCE_INIT_ERROR_GENERIC PM_SOURCE_INIT_ERROR_DIRECTORY PM_SOURCE_INIT_ERROR_NON_REGULAR]
     enum :pm_source_init_result_t, pm_source_init_result_values
@@ -298,9 +298,14 @@ def parse_stream(stream, **options)
         callback = -> (string, size, _) {
           raise "Expected size to be >= 0, got: #{size}" if size <= 0
 
-          if !(line = stream.gets(size - 1)).nil?
+          line = stream.gets(size)
+          if line.nil?
+            0
+          else
             source << line
-            string.write_string("#{line}\x00", line.bytesize + 1)
+            size = line.bytesize
+            string.write_string(line, size)
+            size
           end
         }
 

src/source.c

@@ -365,28 +365,42 @@ bool
 pm_source_stream_read(pm_source_t *source) {
     pm_buffer_t *buffer = source->stream.buffer;
 
+    /*
+     * We should use rb_enc_mbmaxlen() of stream.external_encoding for
+     * the max character length but we don't want to do it here for
+     * performance. We can detect the max character length in all
+     * available encodings by the following command line in the Ruby
+     * source directory:
+     *
+     *   grep -Erh "max (enc|byte) length" enc |
+     *     grep -Eo '[0-9]+' |
+     *     sort |
+     *     tail -n1 # => 6
+     */
+#define MAX_ENC_LEN 6
+
 #define LINE_SIZE 4096
-    char line[LINE_SIZE];
 
-    while (memset(line, '\n', LINE_SIZE), source->stream.fgets(line, LINE_SIZE, source->stream.stream) != NULL) {
-        size_t length = LINE_SIZE;
-        while (length > 0 && line[length - 1] == '\n') length--;
+    /*
+     * io.gets(size) may read size or more bytes when the last
+     * character is a multi-byte character. "+ MAX_ENC_LEN" is for the
+     * case.
+     */
+    char line[LINE_SIZE + MAX_ENC_LEN];
+    int length;
+
+    while ((length = source->stream.fgets(line, LINE_SIZE, source->stream.stream)) > 0) {
+        /* Append the line to the buffer. */
+        pm_buffer_append_string(buffer, line, (size_t)length);
 
-        if (length == LINE_SIZE) {
+        if (line[length - 1] != '\n') {
             /*
-             * 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.
+             * If it doesn't end with a newline, then we'll continue
+             * reading.
              */
-            length--;
-            pm_buffer_append_string(buffer, line, length);
             continue;
         }
 
-        /* Append the line to the buffer. */
-        length--;
-        pm_buffer_append_string(buffer, line, length);
-
         /*
          * Check if the line matches the __END__ marker. If it does, then stop
          * reading and return false. In most circumstances, this means we should
@@ -428,6 +442,8 @@ pm_source_stream_read(pm_source_t *source) {
 
 #undef LINE_SIZE
 
+#undef MAX_ENC_LEN
+
     source->stream.eof = true;
     source->source = (const uint8_t *) pm_buffer_value(buffer);
     source->length = pm_buffer_length(buffer);

test/prism/api/parse_stream_test.rb

@@ -114,5 +114,19 @@ def test_nul_bytes
       assert result.success?
       assert_equal 3, result.value.statements.body.length
     end
+
+    def test_long_line
+      # The important point of this test is that the character at 4095
+      # is a multi-byte character (U+3042 HIRAGANA LETTER A) in this
+      # case. gets(4095) may return 4095 or more bytes if the last
+      # character is a multi-byte character. For example,
+      # StringIO.new("\u3042").gets(1).bytesize is 3 not 1.
+      long_string = "\"" + ("a" * 4093) + "\u3042" + "\""
+      io = StringIO.new(long_string)
+      result = Prism.parse_stream(io)
+
+      assert result.success?
+      assert_equal long_string[1..-2], result.value.statements.body[0].content
+    end
   end
 end