Merge pull request rails#45614 from Shopify/faster-output-buffer

Speedup ActionView::OutputBuffer

Author: byroot

actionview/lib/action_view/buffers.rb

@@ -18,24 +18,64 @@ module ActionView
   #   sbuf << 5
   #   puts sbuf # => "hello\u0005"
   #
-  class OutputBuffer < ActiveSupport::SafeBuffer # :nodoc:
-    def initialize(*)
-      super
-      encode!
+  class OutputBuffer # :nodoc:
+    def initialize(buffer = "")
+      @buffer = String.new(buffer)
+      @buffer.encode!
+    end
+
+    delegate :length, :blank?, :encoding, :encode!, :force_encoding, to: :@buffer
+
+    def to_s
+      @buffer.html_safe
+    end
+    alias_method :html_safe, :to_s
+
+    def to_str
+      @buffer.dup
+    end
+
+    def html_safe?
+      true
     end
 
     def <<(value)
-      return self if value.nil?
-      super(value.to_s)
+      unless value.nil?
+        value = value.to_s
+        @buffer << if value.html_safe?
+          value
+        else
+          CGI.escapeHTML(value)
+        end
+      end
+      self
     end
     alias :append= :<<
 
+    def safe_concat(value)
+      @buffer << value
+      self
+    end
+    alias :safe_append= :safe_concat
+
     def safe_expr_append=(val)
       return self if val.nil?
-      safe_concat val.to_s
+      @buffer << val.to_s
+      self
     end
 
-    alias :safe_append= :safe_concat
+    def initialize_copy(other)
+      @buffer = other.to_str
+    end
+
+    # Don't use this
+    def slice!(range)
+      @buffer.slice!(range)
+    end
+
+    def ==(other)
+      other.class == self.class && @buffer == other.to_str
+    end
   end
 
   class StreamingBuffer # :nodoc:

actionview/lib/action_view/helpers/cache_helper.rb

@@ -285,6 +285,9 @@ def write_fragment_for(name, options)
         pos = output_buffer.length
         yield
         output_safe = output_buffer.html_safe?
+        # We need to modify the buffer in place to be deal with the view handlers
+        # like `builder` that don't access the buffer through `@output_buffer` but
+        # keep the initial reference.
         fragment = output_buffer.slice!(pos..-1)
         if output_safe
           self.output_buffer = output_buffer.class.new(output_buffer)

actionview/lib/action_view/helpers/capture_helper.rb

@@ -43,8 +43,14 @@ module CaptureHelper
       def capture(*args)
         value = nil
         buffer = with_output_buffer { value = yield(*args) }
-        if (string = buffer.presence || value) && string.is_a?(String)
-          ERB::Util.html_escape string
+
+        case string = buffer.presence || value
+        when OutputBuffer
+          string.to_s
+        when ActiveSupport::SafeBuffer
+          string
+        when String
+          ERB::Util.html_escape(string)
         end
       end
 

actionview/lib/action_view/renderer/abstract_renderer.rb

@@ -142,7 +142,7 @@ class RenderedTemplate # :nodoc:
       attr_reader :body, :template
 
       def initialize(body, template)
-        @body = body
+        @body = body.to_s
         @template = template
       end
 

actionview/test/output_buffer_test.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+require "abstract_unit"
+
+class TestOutputBuffer < ActiveSupport::TestCase
+  setup do
+    @buffer = ActionView::OutputBuffer.new
+  end
+
+  test "#<< maintains HTML safety" do
+    @buffer << "<script>alert('pwned!')</script>"
+    assert_predicate @buffer, :html_safe?
+    assert_predicate @buffer.to_s, :html_safe?
+    assert_equal "&lt;script&gt;alert(&#39;pwned!&#39;)&lt;/script&gt;", @buffer.to_s
+  end
+
+  test "#safe_append= bypasses HTML safety" do
+    @buffer.safe_append = "<p>This is fine</p>"
+    assert_predicate @buffer, :html_safe?
+    assert_predicate @buffer.to_s, :html_safe?
+    assert_equal "<p>This is fine</p>", @buffer.to_s
+  end
+
+  test "can be duped" do
+    @buffer << "Hello"
+    copy = @buffer.dup
+    copy << " World!"
+    assert_equal "Hello World!", copy.to_s
+    assert_equal "Hello", @buffer.to_s
+  end
+end

actionview/test/template/capture_helper_test.rb

@@ -182,7 +182,7 @@ def test_with_output_buffer_swaps_the_output_buffer_given_no_argument
     buffer = @av.with_output_buffer do
       @av.output_buffer << "."
     end
-    assert_equal ".", buffer
+    assert_equal ".", buffer.to_s
     assert_nil @av.output_buffer
   end
 
@@ -192,7 +192,7 @@ def test_with_output_buffer_swaps_the_output_buffer_with_an_argument
     @av.with_output_buffer(buffer) do
       @av.output_buffer << "."
     end
-    assert_equal "..", buffer
+    assert_equal "..", buffer.to_s
     assert_nil @av.output_buffer
   end
 
@@ -219,7 +219,7 @@ def test_with_output_buffer_sets_proper_encoding
 
   def test_with_output_buffer_does_not_assume_there_is_an_output_buffer
     assert_nil @av.output_buffer
-    assert_equal "", @av.with_output_buffer { }
+    assert_equal "", @av.with_output_buffer { }.to_s
   end
 
   def alt_encoding(output_buffer)