Merge pull request #695 from Shopify/assign-score

Take nested values into account for assign score

lib/liquid/tags/assign.rb

@@ -23,16 +23,28 @@ module Liquid
     def render(context)
       val = @from.render(context)
       context.scopes.last[@to] = val
-
+      context.resource_limits.assign_score += assign_score_of(val)
-      inc = val.instance_of?(String) || val.instance_of?(Array) || val.instance_of?(Hash) ? val.length : 1
-      context.resource_limits.assign_score += inc
-
       ''.freeze
     end
 
     def blank?
       true
     end
+
+    private
+
+    def assign_score_of(val)
+      if val.instance_of?(String)
+        val.length
+      elsif val.instance_of?(Array) || val.instance_of?(Hash)
+        sum = 1
+        # Uses #each to avoid extra allocations.
+        val.each { |child| sum += assign_score_of(child) }
+        sum
+      else
+        1
+      end
+    end
   end
 
   Template.register_tag('assign'.freeze, Assign)

test/integration/template_test.rb

@@ -133,6 +133,17 @@ class TemplateTest < Minitest::Test
     refute_nil t.resource_limits.assign_score
   end
 
+  def test_resource_limits_assign_score_nested
+    t = Template.parse("{% assign foo = 'aaaa' | reverse %}")
+
+    t.resource_limits.assign_score_limit = 3
+    assert_equal "Liquid error: Memory limits exceeded", t.render
+    assert t.resource_limits.reached?
+
+    t.resource_limits.assign_score_limit = 5
+    assert_equal "", t.render!
+  end
+
   def test_resource_limits_aborts_rendering_after_first_error
     t = Template.parse("{% for a in (1..100) %} foo1 {% endfor %} bar {% for a in (1..100) %} foo2 {% endfor %}")
     t.resource_limits.render_score_limit = 50