Bound resource use when iterating ranges

Gemfile

@@ -32,7 +32,7 @@ group :test do
 end
 
 group :spec do
-  # Includes the merged specs from https://github.com/Shopify/liquid-spec/pull/144.
-  gem 'liquid-spec', github: 'Shopify/liquid-spec', ref: '8a308cb199a7d4e635ae399797560875732c8839'
+  # Includes the range resource-limit specs from https://github.com/Shopify/liquid-spec/pull/165.
+  gem 'liquid-spec', github: 'Shopify/liquid-spec', ref: '84bf25e0edbca5f7e530574788b0f875ab331f50'
   gem 'activesupport', require: false
 end

History.md

@@ -1,5 +1,9 @@
 # Liquid Change Log
 
+## Unreleased
+
+* Avoid materializing integer ranges in `for` and `tablerow` loops, and account for each visited range item.
+
 ## 5.13.0
 
 * Add TruffleRuby in CI [Benoit Daloze]

README.md

@@ -149,6 +149,13 @@ template.render!({ 'x' => 1}, { strict_variables: true })
 #=> Liquid::UndefinedVariable: Liquid error: undefined variable y
 ```
 
+### Resource limits
+
+`render_score_limit` and `cumulative_render_score_limit` account for each item visited by
+integer-range `for` and `tablerow` loops, including loops with empty bodies. This bounds range
+iteration work when a score limit is configured; `render_length_limit` only bounds generated
+output and does not by itself limit CPU work for output-free loops.
+
 ### Usage tracking
 
 To help track usages of a feature or code path in production, we have released opt-in usage tracking. To enable this, we provide an empty `Liquid:: Usage.increment` method which you can customize to your needs. The feature is well suited to https://github.com/Shopify/statsd-instrument. However, the choice of implementation is up to you.

lib/liquid.rb

@@ -83,6 +83,7 @@ module Liquid
 require 'liquid/expression'
 require 'liquid/template'
 require 'liquid/condition'
+require 'liquid/range_slice'
 require 'liquid/utils'
 require 'liquid/tokenizer'
 require 'liquid/parse_context'

lib/liquid/range_slice.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+module Liquid
+  class RangeSlice
+    attr_reader :length
+
+    def initialize(range, from, to, resource_limits)
+      range_length = range.end - range.begin
+      range_length += 1 unless range.exclude_end?
+      range_length = 0 if range_length.negative?
+
+      start = [from, 0].max
+      finish = [to || range_length, range_length].min
+
+      @first = range.begin + start
+      @length = [finish - start, 0].max
+      @direction = 1
+      @resource_limits = resource_limits
+    end
+
+    def empty?
+      @length.zero?
+    end
+
+    def each
+      return enum_for(:each) unless block_given?
+
+      value = @first
+      @length.times do
+        @resource_limits.increment_render_score(1)
+        yield value
+        value += @direction
+      end
+    end
+
+    def reverse!
+      unless empty?
+        @first += @direction * (@length - 1)
+        @direction = -@direction
+      end
+      self
+    end
+  end
+end

lib/liquid/tags/for.rb

@@ -130,7 +130,6 @@ def collection_segment(context)
       end
 
       collection = context.evaluate(@collection_name)
-      collection = collection.to_a if collection.is_a?(Range)
 
       limit_value = context.evaluate(@limit)
       to = if limit_value.nil?
@@ -139,7 +138,9 @@ def collection_segment(context)
         Utils.to_integer(limit_value) + from
       end
 
-      segment = Utils.slice_collection(collection, from, to)
+      segment = Utils.slice_collection_for_iteration(
+        collection, from, to, context.resource_limits, use_range_to_a: true
+      )
       segment.reverse! if @reversed
 
       offsets[@name] = from + segment.length

lib/liquid/tags/table_row.rb

@@ -85,12 +85,13 @@ def render_to_output_buffer(context, output)
       from = @attributes.key?('offset') ? to_integer(context.evaluate(@attributes['offset'])) : 0
       to = @attributes.key?('limit') ? from + to_integer(context.evaluate(@attributes['limit'])) : nil
 
-      collection = Utils.slice_collection(collection, from, to)
+      collection = Utils.slice_collection_for_iteration(collection, from, to, context.resource_limits, allow_endless: true)
       length     = collection.length
 
       cols = @attributes.key?('cols') ? to_integer(context.evaluate(@attributes['cols'])) : length
 
       output << "<tr class=\"row1\">\n"
+      context.resource_limits.increment_write_score(output)
       context.stack do
         tablerowloop = Liquid::TablerowloopDrop.new(length, cols)
         context['tablerowloop'] = tablerowloop
@@ -101,6 +102,7 @@ def render_to_output_buffer(context, output)
           output << "<td class=\"col#{tablerowloop.col}\">"
           super
           output << '</td>'
+          context.resource_limits.increment_write_score(output)
 
           # Handle any interrupts if they exist.
           if context.interrupt?
@@ -110,13 +112,15 @@ def render_to_output_buffer(context, output)
 
           if tablerowloop.col_last && !tablerowloop.last
             output << "</tr>\n<tr class=\"row#{tablerowloop.row + 1}\">"
+            context.resource_limits.increment_write_score(output)
           end
 
           tablerowloop.send(:increment!)
         end
       end
 
       output << "</tr>\n"
+      context.resource_limits.increment_write_score(output)
       output
     end
 

lib/liquid/utils.rb

@@ -13,6 +13,26 @@ def self.slice_collection(collection, from, to)
       end
     end
 
+    # This is intentionally separate from slice_collection, whose Array-returning
+    # behavior is used outside of the iteration tags.
+    def self.slice_collection_for_iteration(
+      collection, from, to, resource_limits, allow_endless: false, use_range_to_a: false
+    )
+      if integer_range?(collection)
+        RangeSlice.new(collection, from, to, resource_limits)
+      elsif collection.is_a?(Range)
+        if use_range_to_a && range_method_overridden?(collection, :to_a)
+          # For historically honored custom Range#to_a. Charge the resulting
+          # selection before buffering it, just as for a custom #each.
+          slice_collection_for_iteration_using_each(collection.to_a, from, to, resource_limits)
+        else
+          slice_range_using_each(collection, from, to, resource_limits, allow_endless: allow_endless)
+        end
+      else
+        slice_collection(collection, from, to)
+      end
+    end
+
     def self.slice_collection_using_each(collection, from, to)
       segments = []
       index    = 0
@@ -38,6 +58,55 @@ def self.slice_collection_using_each(collection, from, to)
       segments
     end
 
+    # Arithmetic slicing must not bypass a Range subclass's custom #each.
+    def self.integer_range?(collection)
+      collection.instance_of?(Range) && collection.begin.is_a?(Integer) && collection.end.is_a?(Integer)
+    end
+    private_class_method :integer_range?
+
+    # Preserve support for Ruby-supplied string ranges and custom Range#each.
+    # Their selected length cannot be inferred from integer bounds, but the tags
+    # need it before rendering for loop metadata, continuation offsets, and columns.
+    # Buffer the selection so we do not have to replay a potentially custom iterator.
+    def self.slice_range_using_each(collection, from, to, resource_limits, allow_endless:)
+      # TableRow historically accepted an endless subclass when its custom #each
+      # was finite, while For historically raised through Range#to_a.
+      if collection.end.nil? && !(allow_endless && (!to.nil? || range_method_overridden?(collection, :each)))
+        raise RangeError, "cannot convert endless range to an array"
+      end
+      if collection.begin.nil? && !range_method_overridden?(collection, :each)
+        raise TypeError, "can't iterate from NilClass"
+      end
+
+      slice_collection_for_iteration_using_each(collection, from, to, resource_limits)
+    end
+    private_class_method :slice_range_using_each
+
+    # Custom Range#each can make a nominally beginless range finite; standard
+    # beginless ranges were rejected before reaching this budgeted traversal.
+    def self.slice_collection_for_iteration_using_each(collection, from, to, resource_limits)
+      return [] if to && to <= from
+
+      segments = []
+      index = 0
+      collection.each do |item|
+        break if to && to <= index
+
+        # Charge preparation, including skipped offsets, before buffering; checking
+        # only while rendering the buffered values would leave this work unbudgeted.
+        resource_limits.increment_render_score(1)
+        segments << item if from <= index
+        index += 1
+      end
+      segments
+    end
+    private_class_method :slice_collection_for_iteration_using_each
+
+    def self.range_method_overridden?(collection, method_name)
+      collection.method(method_name).owner != Range.instance_method(method_name).owner
+    end
+    private_class_method :range_method_overridden?
+
     def self.to_integer(num)
       return num if num.is_a?(Integer)
       num = num.to_s

test/integration/tags/for_tag_test.rb

@@ -465,4 +465,85 @@ def test_for_cleans_up_registers
 
     assert(context.registers[:for_stack].empty?)
   end
+
+  def test_integer_range_is_not_materialized_and_charges_empty_iterations
+    range = bounded_integer_range_with_tripwires
+    template = Template.parse('{% for i in numbers %}{% endfor %}')
+    template.resource_limits.render_score_limit = 3
+
+    assert_raises(Liquid::MemoryError) { template.render!('numbers' => range) }
+
+    assert(template.resource_limits.reached?)
+    assert_equal(4, template.resource_limits.render_score)
+  end
+
+  def test_integer_range_uses_arithmetic_offsets_reversal_and_metadata
+    range = bounded_integer_range_with_tripwires
+    template = Template.parse(
+      '{% for i in numbers reversed offset:997 limit:2 %}{{ forloop.length }}:{{ i }}{% endfor %}',
+    )
+    template.resource_limits.render_score_limit = 10
+
+    assert_equal('2:9992:998', template.render!('numbers' => range))
+  end
+
+  def test_range_scores_cannot_be_bypassed_by_repeated_renders
+    template = Template.parse('{% for i in (1..2) %}{% endfor %}')
+    template.resource_limits.cumulative_render_score_limit = 3
+    assert_equal('', template.render!)
+    assert_raises(Liquid::MemoryError) { template.render! }
+  end
+
+  def test_range_break_charges_only_visited_items_and_preserves_full_metadata
+    range = bounded_integer_range_with_tripwires
+    template = Template.parse(
+      '{% for i in numbers reversed %}{{ forloop.length }}:{{ i }}{% break %}{% endfor %}',
+    )
+    template.resource_limits.render_score_limit = 10
+
+    assert_equal('1000:1000', template.render!('numbers' => range))
+    assert_operator(template.resource_limits.render_score, :<=, 10)
+  end
+
+  def test_range_subclass_uses_its_custom_each
+    range = Class.new(Range) do
+      def each
+        yield 10
+        yield 20
+      end
+    end.new(nil, 3)
+
+    assert_template_result('1020', '{% for i in numbers %}{{ i }}{% endfor %}', { 'numbers' => range })
+  end
+
+  def test_range_subclass_custom_to_a_is_honored_for_finite_and_open_bounds
+    range_class = Class.new(Range) do
+      def to_a
+        [42]
+      end
+    end
+
+    [range_class.new(1, 3), range_class.new(nil, 3), range_class.new(1, nil)].each do |range|
+      assert_template_result('42', '{% for i in numbers %}{{ i }}{% endfor %}', { 'numbers' => range })
+    end
+  end
+
+  def test_endless_range_remains_unsupported_with_a_limit
+    template = Template.parse('{% for i in numbers limit:2 %}{{ i }}{% endfor %}')
+
+    assert_raises(RangeError) { template.render!('numbers' => (1..)) }
+  end
+
+  def test_endless_range_subclass_with_custom_each_remains_unsupported
+    range = Class.new(Range) do
+      def each
+        yield 10
+        yield 20
+      end
+    end.new(1, nil)
+
+    assert_raises(RangeError) do
+      Template.parse('{% for i in numbers %}{{ i }}{% endfor %}').render!('numbers' => range)
+    end
+  end
 end