chore(internal): protect SSE parsing pipeline from broken UTF-8 characters

Author: stainless-app[bot]

lib/openai/internal/transport/pooled_net_requester.rb

@@ -149,7 +149,7 @@ def execute(request)
                 break if finished
 
                 rsp.read_body do |bytes|
-                  y << bytes
+                  y << bytes.force_encoding(Encoding::BINARY)
                   break if finished
 
                   self.class.calibrate_socket_timeout(conn, deadline)

lib/openai/internal/util.rb

@@ -448,7 +448,7 @@ def initialize(src, &blk)
             else
               src
             end
-          @buf = String.new.b
+          @buf = String.new
           @blk = blk
         end
       end
@@ -460,7 +460,7 @@ class << self
         # @return [Enumerable<String>]
         def writable_enum(&blk)
           Enumerator.new do |y|
-            buf = String.new.b
+            buf = String.new
             y.define_singleton_method(:write) do
               self << buf.replace(_1)
               buf.bytesize
@@ -582,14 +582,35 @@ def encode_content(headers, body)
 
         # @api private
         #
+        # https://www.iana.org/assignments/character-sets/character-sets.xhtml
+        #
+        # @param content_type [String]
+        # @param text [String]
+        def force_charset!(content_type, text:)
+          charset = /charset=([^;\s]+)/.match(content_type)&.captures&.first
+
+          return unless charset
+
+          begin
+            encoding = Encoding.find(charset)
+            text.force_encoding(encoding)
+          rescue ArgumentError
+            nil
+          end
+        end
+
+        # @api private
+        #
+        # Assumes each chunk in stream has `Encoding::BINARY`.
+        #
         # @param headers [Hash{String=>String}, Net::HTTPHeader]
         # @param stream [Enumerable<String>]
         # @param suppress_error [Boolean]
         #
         # @raise [JSON::ParserError]
         # @return [Object]
         def decode_content(headers, stream:, suppress_error: false)
-          case headers["content-type"]
+          case (content_type = headers["content-type"])
           in %r{^application/(?:vnd\.api\+)?json}
             json = stream.to_a.join
             begin
@@ -606,11 +627,10 @@ def decode_content(headers, stream:, suppress_error: false)
           in %r{^text/event-stream}
             lines = decode_lines(stream)
             decode_sse(lines)
-          in %r{^text/}
-            stream.to_a.join
           else
-            # TODO: parsing other response types
-            StringIO.new(stream.to_a.join)
+            text = stream.to_a.join
+            force_charset!(content_type, text: text)
+            StringIO.new(text)
           end
         end
       end
@@ -675,12 +695,17 @@ def chain_fused(enum, &blk)
       class << self
         # @api private
         #
+        # Assumes Strings have been forced into having `Encoding::BINARY`.
+        #
+        # This decoder is responsible for reassembling lines split across multiple
+        # fragments.
+        #
         # @param enum [Enumerable<String>]
         #
         # @return [Enumerable<String>]
         def decode_lines(enum)
           re = /(\r\n|\r|\n)/
-          buffer = String.new.b
+          buffer = String.new
           cr_seen = nil
 
           chain_fused(enum) do |y|
@@ -711,6 +736,8 @@ def decode_lines(enum)
         #
         # https://html.spec.whatwg.org/multipage/server-sent-events.html#parsing-an-event-stream
         #
+        # Assumes that `lines` has been decoded with `#decode_lines`.
+        #
         # @param lines [Enumerable<String>]
         #
         # @return [Enumerable<Hash{Symbol=>Object}>]
@@ -734,7 +761,7 @@ def decode_sse(lines)
                 in "event"
                   current.merge!(event: value)
                 in "data"
-                  (current[:data] ||= String.new.b) << (value << "\n")
+                  (current[:data] ||= String.new) << (value << "\n")
                 in "id" unless value.include?("\0")
                   current.merge!(id: value)
                 in "retry" if /^\d+$/ =~ value

rbi/lib/openai/internal/util.rbi

@@ -215,6 +215,14 @@ module OpenAI
         def encode_content(headers, body); end
 
         # @api private
+        #
+        # https://www.iana.org/assignments/character-sets/character-sets.xhtml
+        sig { params(content_type: String, text: String).void }
+        def force_charset!(content_type, text:); end
+
+        # @api private
+        #
+        # Assumes each chunk in stream has `Encoding::BINARY`.
         sig do
           params(
             headers: T.any(T::Hash[String, String], Net::HTTPHeader),
@@ -263,12 +271,19 @@ module OpenAI
 
       class << self
         # @api private
+        #
+        # Assumes Strings have been forced into having `Encoding::BINARY`.
+        #
+        # This decoder is responsible for reassembling lines split across multiple
+        # fragments.
         sig { params(enum: T::Enumerable[String]).returns(T::Enumerable[String]) }
         def decode_lines(enum); end
 
         # @api private
         #
         # https://html.spec.whatwg.org/multipage/server-sent-events.html#parsing-an-event-stream
+        #
+        # Assumes that `lines` has been decoded with `#decode_lines`.
         sig do
           params(lines: T::Enumerable[String]).returns(T::Enumerable[OpenAI::Internal::Util::ServerSentEvent])
         end

sig/openai/internal/util.rbs

@@ -120,6 +120,8 @@ module OpenAI
         top body
       ) -> top
 
+      def self?.force_charset!: (String content_type, text: String) -> void
+
       def self?.decode_content: (
         ::Hash[String, String] headers,
         stream: Enumerable[String],

test/openai/internal/util_test.rb

@@ -368,6 +368,24 @@ def test_close_fused_sse_chain
   end
 end
 
+class OpenAI::Test::UtilContentDecodingTest < Minitest::Test
+  def test_charset
+    cases = {
+      "application/json" => Encoding::BINARY,
+      "application/json; charset=utf-8" => Encoding::UTF_8,
+      "charset=uTf-8 application/json; " => Encoding::UTF_8,
+      "charset=UTF-8; application/json; " => Encoding::UTF_8,
+      "charset=ISO-8859-1 ;application/json; " => Encoding::ISO_8859_1,
+      "charset=EUC-KR ;application/json; " => Encoding::EUC_KR
+    }
+    text = String.new.force_encoding(Encoding::BINARY)
+    cases.each do |content_type, encoding|
+      OpenAI::Internal::Util.force_charset!(content_type, text: text)
+      assert_equal(encoding, text.encoding)
+    end
+  end
+end
+
 class OpenAI::Test::UtilSseTest < Minitest::Test
   def test_decode_lines
     cases = {
@@ -381,7 +399,9 @@ def test_decode_lines
       %W[\na b\n\n] => %W[\n ab\n \n],
       %W[\na b] => %W[\n ab],
       %W[\u1F62E\u200D\u1F4A8] => %W[\u1F62E\u200D\u1F4A8],
-      %W[\u1F62E \u200D \u1F4A8] => %W[\u1F62E\u200D\u1F4A8]
+      %W[\u1F62E \u200D \u1F4A8] => %W[\u1F62E\u200D\u1F4A8],
+      ["\xf0\x9f".b, "\xa5\xba".b] => ["\xf0\x9f\xa5\xba".b],
+      ["\xf0".b, "\x9f".b, "\xa5".b, "\xba".b] => ["\xf0\x9f\xa5\xba".b]
     }
     eols = %W[\n \r \r\n]
     cases.each do |enum, expected|