diff --git a/lib/liquid/block_body.rb b/lib/liquid/block_body.rb index c9dbfbd..18b3062 100644 --- a/lib/liquid/block_body.rb +++ b/lib/liquid/block_body.rb @@ -155,12 +155,10 @@ module Liquid end def render_to_output_buffer(context, output) - context.resource_limits.render_score += @nodelist.length + context.resource_limits.increment_render_score(@nodelist.length) idx = 0 while (node = @nodelist[idx]) - previous_output_size = output.bytesize - if node.instance_of?(String) output << node else @@ -172,7 +170,7 @@ module Liquid end idx += 1 - raise_if_resource_limits_reached(context, output.bytesize - previous_output_size) + context.resource_limits.increment_write_score(output) end output @@ -184,17 +182,13 @@ module Liquid node.render_to_output_buffer(context, output) rescue UndefinedVariable, UndefinedDropMethod, UndefinedFilter => e context.handle_error(e, node.line_number) + rescue MemoryError + raise rescue ::StandardError => e line_number = node.is_a?(String) ? nil : node.line_number output << context.handle_error(e, line_number) end - def raise_if_resource_limits_reached(context, length) - context.resource_limits.render_length += length - return unless context.resource_limits.reached? - raise MemoryError, "Memory limits exceeded" - end - def create_variable(token, parse_context) token.scan(ContentOfVariable) do |content| markup = content.first diff --git a/lib/liquid/resource_limits.rb b/lib/liquid/resource_limits.rb index 9c898f0..70fac24 100644 --- a/lib/liquid/resource_limits.rb +++ b/lib/liquid/resource_limits.rb @@ -2,8 +2,8 @@ module Liquid class ResourceLimits - attr_accessor :render_length, :render_score, :assign_score, - :render_length_limit, :render_score_limit, :assign_score_limit + attr_accessor :render_length_limit, :render_score_limit, :assign_score_limit + attr_reader :render_score, :assign_score def initialize(limits) @render_length_limit = limits[:render_length_limit] @@ -12,14 +12,51 @@ module Liquid reset end + def increment_render_score(amount) + @render_score += amount + raise_limits_reached if @render_score_limit && @render_score > @render_score_limit + end + + def increment_assign_score(amount) + @assign_score += amount + raise_limits_reached if @assign_score_limit && @assign_score > @assign_score_limit + end + + # update either render_length or assign_score based on whether or not the writes are captured + def increment_write_score(output) + if (last_captured = @last_capture_length) + captured = output.bytesize + increment = captured - last_captured + @last_capture_length = captured + increment_assign_score(increment) + elsif @render_length_limit && output.bytesize > @render_length_limit + raise_limits_reached + end + end + + def raise_limits_reached + @reached_limit = true + raise MemoryError, "Memory limits exceeded" + end + def reached? - (@render_length_limit && @render_length > @render_length_limit) || - (@render_score_limit && @render_score > @render_score_limit) || - (@assign_score_limit && @assign_score > @assign_score_limit) + @reached_limit end def reset - @render_length = @render_score = @assign_score = 0 + @reached_limit = false + @last_capture_length = nil + @render_score = @assign_score = 0 + end + + def with_capture + old_capture_length = @last_capture_length + begin + @last_capture_length = 0 + yield + ensure + @last_capture_length = old_capture_length + end end end end diff --git a/lib/liquid/tags/assign.rb b/lib/liquid/tags/assign.rb index 3a3f683..404b005 100644 --- a/lib/liquid/tags/assign.rb +++ b/lib/liquid/tags/assign.rb @@ -27,7 +27,7 @@ module Liquid def render_to_output_buffer(context, output) val = @from.render(context) context.scopes.last[@to] = val - context.resource_limits.assign_score += assign_score_of(val) + context.resource_limits.increment_assign_score(assign_score_of(val)) output end diff --git a/lib/liquid/tags/capture.rb b/lib/liquid/tags/capture.rb index 425032d..3eb63bb 100644 --- a/lib/liquid/tags/capture.rb +++ b/lib/liquid/tags/capture.rb @@ -25,9 +25,10 @@ module Liquid end def render_to_output_buffer(context, output) - capture_output = render(context) - context.scopes.last[@to] = capture_output - context.resource_limits.assign_score += capture_output.bytesize + context.resource_limits.with_capture do + capture_output = render(context) + context.scopes.last[@to] = capture_output + end output end diff --git a/test/integration/template_test.rb b/test/integration/template_test.rb index 1c12ff6..246c645 100644 --- a/test/integration/template_test.rb +++ b/test/integration/template_test.rb @@ -111,13 +111,12 @@ class TemplateTest < Minitest::Test def test_resource_limits_render_length t = Template.parse("0123456789") - t.resource_limits.render_length_limit = 5 + t.resource_limits.render_length_limit = 9 assert_equal("Liquid error: Memory limits exceeded", t.render) assert(t.resource_limits.reached?) t.resource_limits.render_length_limit = 10 assert_equal("0123456789", t.render!) - refute_nil(t.resource_limits.render_length) end def test_resource_limits_render_score @@ -180,36 +179,33 @@ class TemplateTest < Minitest::Test t.render! assert(t.resource_limits.assign_score > 0) assert(t.resource_limits.render_score > 0) - assert(t.resource_limits.render_length > 0) end def test_render_length_persists_between_blocks t = Template.parse("{% if true %}aaaa{% endif %}") - t.resource_limits.render_length_limit = 7 + t.resource_limits.render_length_limit = 3 assert_equal("Liquid error: Memory limits exceeded", t.render) - t.resource_limits.render_length_limit = 8 + t.resource_limits.render_length_limit = 4 assert_equal("aaaa", t.render) t = Template.parse("{% if true %}aaaa{% endif %}{% if true %}bbb{% endif %}") - t.resource_limits.render_length_limit = 13 + t.resource_limits.render_length_limit = 6 assert_equal("Liquid error: Memory limits exceeded", t.render) - t.resource_limits.render_length_limit = 14 + t.resource_limits.render_length_limit = 7 assert_equal("aaaabbb", t.render) t = Template.parse("{% if true %}a{% endif %}{% if true %}b{% endif %}{% if true %}a{% endif %}{% if true %}b{% endif %}{% if true %}a{% endif %}{% if true %}b{% endif %}") t.resource_limits.render_length_limit = 5 assert_equal("Liquid error: Memory limits exceeded", t.render) - t.resource_limits.render_length_limit = 11 - assert_equal("Liquid error: Memory limits exceeded", t.render) - t.resource_limits.render_length_limit = 12 + t.resource_limits.render_length_limit = 6 assert_equal("ababab", t.render) end def test_render_length_uses_number_of_bytes_not_characters t = Template.parse("{% if true %}すごい{% endif %}") - t.resource_limits.render_length_limit = 10 + t.resource_limits.render_length_limit = 8 assert_equal("Liquid error: Memory limits exceeded", t.render) - t.resource_limits.render_length_limit = 18 + t.resource_limits.render_length_limit = 9 assert_equal("すごい", t.render) end @@ -219,7 +215,6 @@ class TemplateTest < Minitest::Test t.render!(context) assert(context.resource_limits.assign_score > 0) assert(context.resource_limits.render_score > 0) - assert(context.resource_limits.render_length > 0) end def test_can_use_drop_as_context