Add OutputBuffer#raw and #capture to reduce the need to swap the buffer

Right now many helpers have to deal with two modes of operation to
capture view output.

The main one is to swap the `@output_buffer` variable with a new buffer.
But since some view implementations such as `builder` keep a reference
on the buffer they were initialized with, this doesn't always work.

So additionally, the various capturing helpers also record the buffer
length prior to executing the block, and then `slice!` the buffer back
to its original size.

This is wasteful and make the code rather unclear.

Now that `OutputBuffer` is a delegator, I'd like to refactor all this
so that:

  - @output_buffer is no longer re-assigned
  - A single OutputBuffer instance is used for the entire response rendering
  - Instead capturing is done through `OutputBuffer#capture`

Once the above is achieved, it should allow us to enabled Erubi's
`:chain_appends` option and get some reduced template size and some
performance.

Not re-assigning `@output_buffer` will also allow template to access
the local variable instead of an instance variable, which is cheaper.

But more importantly, that should make the code easier to understand
and easier to be compatible with `StreamingBuffer`.

actionmailer/test/caching_test.rb

@@ -221,31 +221,8 @@ def self.output_buffer=; end
 
     cache_helper.stub :controller, controller do
       cache_helper.stub :output_buffer, output_buffer do
-        assert_called_with cache_helper, :output_buffer=, [output_buffer.class.new(output_buffer)] do
-          assert_nothing_raised do
-            cache_helper.send :fragment_for, "Test fragment name", "Test fragment", &Proc.new { nil }
-          end
-        end
-      end
-    end
-  end
-
-  def test_safe_buffer
-    output_buffer = ActiveSupport::SafeBuffer.new
-    controller = MockController.new
-    cache_helper = Class.new do
-      def self.controller; end
-      def self.output_buffer; end
-      def self.output_buffer=; end
-    end
-    cache_helper.extend(ActionView::Helpers::CacheHelper)
-
-    cache_helper.stub :controller, controller do
-      cache_helper.stub :output_buffer, output_buffer do
-        assert_called_with cache_helper, :output_buffer=, [output_buffer.class.new(output_buffer)] do
-          assert_nothing_raised do
-            cache_helper.send :fragment_for, "Test fragment name", "Test fragment", &Proc.new { nil }
-          end
+        assert_nothing_raised do
+          cache_helper.send :fragment_for, "Test fragment name", "Test fragment", &Proc.new { nil }
         end
       end
     end

actionpack/test/controller/caching_test.rb

@@ -344,31 +344,8 @@ def self.output_buffer=; end
 
     cache_helper.stub :controller, controller do
       cache_helper.stub :output_buffer, output_buffer do
-        assert_called_with cache_helper, :output_buffer=, [output_buffer.class.new(output_buffer)] do
-          assert_nothing_raised do
-            cache_helper.send :fragment_for, "Test fragment name", "Test fragment", &Proc.new { nil }
-          end
-        end
-      end
-    end
-  end
-
-  def test_safe_buffer
-    output_buffer = ActiveSupport::SafeBuffer.new
-    controller = MockController.new
-    cache_helper = Class.new do
-      def self.controller; end
-      def self.output_buffer; end
-      def self.output_buffer=; end
-    end
-    cache_helper.extend(ActionView::Helpers::CacheHelper)
-
-    cache_helper.stub :controller, controller do
-      cache_helper.stub :output_buffer, output_buffer do
-        assert_called_with cache_helper, :output_buffer=, [output_buffer.class.new(output_buffer)] do
-          assert_nothing_raised do
-            cache_helper.send :fragment_for, "Test fragment name", "Test fragment", &Proc.new { nil }
-          end
+        assert_nothing_raised do
+          cache_helper.send :fragment_for, "Test fragment name", "Test fragment", &Proc.new { nil }
         end
       end
     end

actionview/lib/action_view/buffers.rb

@@ -20,19 +20,19 @@ module ActionView
   #
   class OutputBuffer # :nodoc:
     def initialize(buffer = "")
-      @buffer = String.new(buffer)
-      @buffer.encode!
+      @raw_buffer = String.new(buffer)
+      @raw_buffer.encode!
     end
 
-    delegate :length, :blank?, :encoding, :encode!, :force_encoding, to: :@buffer
+    delegate :length, :blank?, :encoding, :encode!, :force_encoding, to: :@raw_buffer
 
     def to_s
-      @buffer.html_safe
+      @raw_buffer.html_safe
     end
     alias_method :html_safe, :to_s
 
     def to_str
-      @buffer.dup
+      @raw_buffer.dup
     end
 
     def html_safe?
@@ -42,7 +42,7 @@ def html_safe?
     def <<(value)
       unless value.nil?
         value = value.to_s
-        @buffer << if value.html_safe?
+        @raw_buffer << if value.html_safe?
           value
         else
           CGI.escapeHTML(value)
@@ -53,28 +53,54 @@ def <<(value)
     alias :append= :<<
 
     def safe_concat(value)
-      @buffer << value
+      @raw_buffer << value
       self
     end
     alias :safe_append= :safe_concat
 
     def safe_expr_append=(val)
       return self if val.nil?
-      @buffer << val.to_s
+      @raw_buffer << val.to_s
       self
     end
 
     def initialize_copy(other)
-      @buffer = other.to_str
+      @raw_buffer = other.to_str
     end
 
-    # Don't use this
-    def slice!(range)
-      @buffer.slice!(range)
+    def capture
+      new_buffer = +""
+      old_buffer, @raw_buffer = @raw_buffer, new_buffer
+      yield
+      new_buffer.html_safe
+    ensure
+      @raw_buffer = old_buffer
     end
 
     def ==(other)
-      other.class == self.class && @buffer == other.to_str
+      other.class == self.class && @raw_buffer == other.to_str
+    end
+
+    def raw
+      RawOutputBuffer.new(self)
+    end
+
+    attr_reader :raw_buffer
+  end
+
+  class RawOutputBuffer # :nodoc:
+    def initialize(buffer)
+      @buffer = buffer
+    end
+
+    def <<(value)
+      unless value.nil?
+        @buffer.raw_buffer << value.to_s
+      end
+    end
+
+    def raw
+      self
     end
   end
 
@@ -96,12 +122,43 @@ def safe_concat(value)
     end
     alias :safe_append= :safe_concat
 
+    def capture
+      buffer = +""
+      old_block, @block = @block, ->(value) { buffer << value }
+      yield
+      buffer.html_safe
+    ensure
+      @block = old_block
+    end
+
     def html_safe?
       true
     end
 
     def html_safe
       self
     end
+
+    def raw
+      RawStreamingBuffer.new(self)
+    end
+
+    attr_reader :block
+  end
+
+  class RawStreamingBuffer # :nodoc:
+    def initialize(buffer)
+      @buffer = buffer
+    end
+
+    def <<(value)
+      unless value.nil?
+        @buffer.block.call(value.to_s)
+      end
+    end
+
+    def raw
+      self
+    end
   end
 end

actionview/lib/action_view/helpers/cache_helper.rb

@@ -281,17 +281,8 @@ def read_fragment_for(name, options)
         controller.read_fragment(name, options)
       end
 
-      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)
-        end
+      def write_fragment_for(name, options, &block)
+        fragment = output_buffer.capture(&block)
         controller.write_fragment(name, fragment, options)
       end
 

actionview/lib/action_view/template/handlers/builder.rb

@@ -7,10 +7,9 @@ class Builder
 
       def call(template, source)
         require_engine
-        "xml = ::Builder::XmlMarkup.new(:indent => 2);" \
-          "self.output_buffer = xml.target!;" +
-          source +
-          ";xml.target!;"
+        "xml = ::Builder::XmlMarkup.new(indent: 2, target: output_buffer.raw);" \
+          "#{source};" \
+          "output_buffer.to_s"
       end
 
       private

actionview/test/buffers_test.rb

@@ -0,0 +1,89 @@
+# frozen_string_literal: true
+
+require "abstract_unit"
+
+module SharedBufferTests
+  def self.included(test_case)
+    test_case.test "#<< maintains HTML safety" do
+      @buffer << "<script>alert('pwned!')</script>"
+      assert_predicate @buffer, :html_safe?
+      assert_predicate output, :html_safe?
+      assert_equal "&lt;script&gt;alert(&#39;pwned!&#39;)&lt;/script&gt;", output
+    end
+
+    test_case.test "#safe_append= bypasses HTML safety" do
+      @buffer.safe_append = "<p>This is fine</p>"
+      assert_predicate @buffer, :html_safe?
+      assert_predicate output, :html_safe?
+      assert_equal "<p>This is fine</p>", output
+    end
+
+    test_case.test "#raw allow to bypass HTML escaping" do
+      raw_buffer = @buffer.raw
+      raw_buffer << "<script>alert('pwned!')</script>"
+      assert_predicate @buffer, :html_safe?
+      assert_predicate output, :html_safe?
+      assert_equal "<script>alert('pwned!')</script>", output
+    end
+
+    test_case.test "#capture allow to intercept writes" do
+      @buffer << "Hello"
+      result = @buffer.capture do
+        @buffer << "George!"
+      end
+      assert_equal "George!", result
+      assert_predicate result, :html_safe?
+
+      @buffer << " World!"
+      assert_equal "Hello World!", output
+    end
+
+    test_case.test "#raw respects #capture" do
+      @buffer << "Hello"
+      raw_buffer = @buffer.raw
+      result = @buffer.capture do
+        raw_buffer << "George!"
+      end
+      assert_equal "George!", result
+      assert_predicate result, :html_safe?
+
+      @buffer << " World!"
+      assert_equal "Hello World!", output
+    end
+  end
+end
+
+class TestOutputBuffer < ActiveSupport::TestCase
+  include SharedBufferTests
+
+  setup do
+    @buffer = ActionView::OutputBuffer.new
+  end
+
+  test "can be duped" do
+    @buffer << "Hello"
+    copy = @buffer.dup
+    copy << " World!"
+    assert_equal "Hello World!", copy.to_s
+    assert_equal "Hello", output
+  end
+
+  private
+    def output
+      @buffer.to_s
+    end
+end
+
+class TestStreamingBuffer < ActiveSupport::TestCase
+  include SharedBufferTests
+
+  setup do
+    @raw_buffer = +""
+    @buffer = ActionView::StreamingBuffer.new(@raw_buffer.method(:<<))
+  end
+
+  private
+    def output
+      @raw_buffer.html_safe
+    end
+end