diff --git a/lib/liquid/tags/assign.rb b/lib/liquid/tags/assign.rb index ff4ab40e..6d4f7d8d 100644 --- a/lib/liquid/tags/assign.rb +++ b/lib/liquid/tags/assign.rb @@ -45,11 +45,18 @@ module Liquid def assign_score_of(val) if val.instance_of?(String) val.bytesize - elsif val.instance_of?(Array) || val.instance_of?(Hash) + elsif val.instance_of?(Array) sum = 1 # Uses #each to avoid extra allocations. val.each { |child| sum += assign_score_of(child) } sum + elsif val.instance_of?(Hash) + sum = 1 + val.each do |key, entry_value| + sum += assign_score_of(key) + sum += assign_score_of(entry_value) + end + sum else 1 end diff --git a/test/integration/assign_test.rb b/test/integration/assign_test.rb index 99a32dcc..b956fd1e 100644 --- a/test/integration/assign_test.rb +++ b/test/integration/assign_test.rb @@ -47,4 +47,66 @@ class AssignTest < Minitest::Test assert Template.parse("{% assign foo = ('X' | downcase) %}") end end -end # AssignTest + + def test_assign_score_exceeding_resource_limit + t = Template.parse("{% assign foo = 42 %}{% assign bar = 23 %}") + t.resource_limits.assign_score_limit = 1 + assert_equal("Liquid error: Memory limits exceeded", t.render) + assert(t.resource_limits.reached?) + + t.resource_limits.assign_score_limit = 2 + assert_equal("", t.render!) + refute_nil(t.resource_limits.assign_score) + end + + def test_assign_score_exceeding_limit_from_composite_object + 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_assign_score_of_int + assert_equal(1, assign_score_of(123)) + end + + def test_assign_score_of_string_counts_bytes + assert_equal(3, assign_score_of('123')) + assert_equal(5, assign_score_of('12345')) + assert_equal(9, assign_score_of('すごい')) + end + + def test_assign_score_of_array + assert_equal(1, assign_score_of([])) + assert_equal(2, assign_score_of([123])) + assert_equal(6, assign_score_of([123, 'abcd'])) + end + + def test_assign_score_of_hash + assert_equal(1, assign_score_of({})) + assert_equal(5, assign_score_of('int' => 123)) + assert_equal(12, assign_score_of('int' => 123, 'str' => 'abcd')) + end + + private + + class ObjectWrapperDrop < Liquid::Drop + def initialize(obj) + @obj = obj + end + + def value + @obj + end + end + + def assign_score_of(obj) + context = Liquid::Context.new('drop' => ObjectWrapperDrop.new(obj)) + Liquid::Template.parse('{% assign obj = drop.value %}').render!(context) + context.resource_limits.assign_score + end +end diff --git a/test/integration/capture_test.rb b/test/integration/capture_test.rb index 39098dde..c2dc5269 100644 --- a/test/integration/capture_test.rb +++ b/test/integration/capture_test.rb @@ -49,4 +49,10 @@ class CaptureTest < Minitest::Test rendered = template.render! assert_equal("3-3", rendered.gsub(/\s/, '')) end -end # CaptureTest + + def test_increment_assign_score_by_bytes_not_characters + t = Template.parse("{% capture foo %}すごい{% endcapture %}") + t.render! + assert_equal(9, t.resource_limits.assign_score) + end +end diff --git a/test/integration/template_test.rb b/test/integration/template_test.rb index 246c6459..8cc3d790 100644 --- a/test/integration/template_test.rb +++ b/test/integration/template_test.rb @@ -135,38 +135,6 @@ class TemplateTest < Minitest::Test refute_nil(t.resource_limits.render_score) end - def test_resource_limits_assign_score - t = Template.parse("{% assign foo = 42 %}{% assign bar = 23 %}") - t.resource_limits.assign_score_limit = 1 - assert_equal("Liquid error: Memory limits exceeded", t.render) - assert(t.resource_limits.reached?) - - t.resource_limits.assign_score_limit = 2 - assert_equal("", t.render!) - refute_nil(t.resource_limits.assign_score) - end - - def test_resource_limits_assign_score_counts_bytes_not_characters - t = Template.parse("{% assign foo = 'すごい' %}") - t.render - assert_equal(9, t.resource_limits.assign_score) - - t = Template.parse("{% capture foo %}すごい{% endcapture %}") - t.render - assert_equal(9, 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