diff --git a/lib/liquid/block_body.rb b/lib/liquid/block_body.rb index 2456034..18b3062 100644 --- a/lib/liquid/block_body.rb +++ b/lib/liquid/block_body.rb @@ -170,7 +170,7 @@ module Liquid end idx += 1 - context.resource_limits.check_render_length(output.bytesize) + context.resource_limits.increment_write_score(output) end output diff --git a/lib/liquid/resource_limits.rb b/lib/liquid/resource_limits.rb index 4c50bff..70fac24 100644 --- a/lib/liquid/resource_limits.rb +++ b/lib/liquid/resource_limits.rb @@ -22,8 +22,16 @@ module Liquid raise_limits_reached if @assign_score_limit && @assign_score > @assign_score_limit end - def check_render_length(output_byte_size) - raise_limits_reached if @render_length_limit && output_byte_size > @render_length_limit + # 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 @@ -37,7 +45,18 @@ module Liquid def reset @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/capture.rb b/lib/liquid/tags/capture.rb index 5a2189d..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.increment_assign_score(capture_output.bytesize) + context.resource_limits.with_capture do + capture_output = render(context) + context.scopes.last[@to] = capture_output + end output end