Add optional resource usage limitations to number of rendering calls, length of rendering output and/or number of variable/capture assignments

lib/liquid/block.rb

@@ -90,6 +90,9 @@ module Liquid
 
     def render_all(list, context)
       output = []
+      context.resource_limits[:render_length_current] = 0
+      context.resource_limits[:render_score_current] += list.length
+
       list.each do |token|
         # Break out if we have any unhanded interrupts.
         break if context.has_interrupt?
@@ -103,7 +106,12 @@ module Liquid
             break
           end
 
-          output << (token.respond_to?(:render) ? token.render(context) : token)
+          token_output = (token.respond_to?(:render) ? token.render(context) : token)
+          context.resource_limits[:render_length_current] += (token_output.respond_to?(:length) ? token_output.length : 1)
+          raise Liquid::MemoryError, context.resource_limits if context.resource_limits_reached?
+          output << token_output
+        rescue Liquid::MemoryError => e
+          raise e
         rescue ::StandardError => e
           output << (context.handle_error(e))
         end

lib/liquid/context.rb

@@ -13,19 +13,26 @@ module Liquid
   #
   #   context['bob']  #=> nil  class Context
   class Context
-    attr_reader :scopes, :errors, :registers, :environments
+    attr_reader :scopes, :errors, :registers, :environments, :resource_limits
 
-    def initialize(environments = {}, outer_scope = {}, registers = {}, rethrow_errors = false)
-      @environments   = [environments].flatten
-      @scopes         = [(outer_scope || {})]
-      @registers      = registers
-      @errors         = []
-      @rethrow_errors = rethrow_errors
+    def initialize(environments = {}, outer_scope = {}, registers = {}, rethrow_errors = false, resource_limits = {})
+      @environments    = [environments].flatten
+      @scopes          = [(outer_scope || {})]
+      @registers       = registers
+      @errors          = []
+      @rethrow_errors  = rethrow_errors
+      @resource_limits = (resource_limits || {}).merge!({ :render_score_current => 0, :assign_score_current => 0 })
       squash_instance_assigns_with_environments
 
       @interrupts = []
     end
 
+    def resource_limits_reached?
+      (@resource_limits[:render_length_limit] && @resource_limits[:render_length_current] > @resource_limits[:render_length_limit]) ||
+      (@resource_limits[:render_score_limit]  && @resource_limits[:render_score_current]  > @resource_limits[:render_score_limit] ) ||
+      (@resource_limits[:assign_score_limit]  && @resource_limits[:assign_score_current]  > @resource_limits[:assign_score_limit] )
+    end
+
     def strainer
       @strainer ||= Strainer.create(self)
     end

lib/liquid/errors.rb

@@ -1,6 +1,6 @@
 module Liquid
   class Error < ::StandardError; end
-  
+
   class ArgumentError < Error; end
   class ContextError < Error; end
   class FilterNotFound < Error; end
@@ -8,4 +8,5 @@ module Liquid
   class StandardError < Error; end
   class SyntaxError < Error; end
   class StackLevelError < Error; end
+  class MemoryError < Error; end
 end

lib/liquid/tags/assign.rb

@@ -23,8 +23,10 @@ module Liquid
     end
   
     def render(context)
-       context.scopes.last[@to] = @from.render(context)
-       ''
+      val = @from.render(context)
+      context.scopes.last[@to] = val
+      context.resource_limits[:assign_score_current] += (val.respond_to?(:length) ? val.length : 1)
+      ''
     end 
   
   end  

lib/liquid/tags/capture.rb

@@ -27,6 +27,7 @@ module Liquid
     def render(context)
       output = super
       context.scopes.last[@to] = output
+      context.resource_limits[:assign_score_current] += (output.respond_to?(:length) ? output.length : 1)
       ''
     end
   end

lib/liquid/template.rb

@@ -14,7 +14,7 @@ module Liquid
   #   template.render('user_name' => 'bob')
   #
   class Template
-    attr_accessor :root
+    attr_accessor :root, :resource_limits
     @@file_system = BlankFileSystem.new
 
     class << self
@@ -93,9 +93,9 @@ module Liquid
       when Liquid::Context
         args.shift
       when Hash
-        Context.new([args.shift, assigns], instance_assigns, registers, @rethrow_errors)
+        Context.new([args.shift, assigns], instance_assigns, registers, @rethrow_errors, @resource_limits)
       when nil
-        Context.new(assigns, instance_assigns, registers, @rethrow_errors)
+        Context.new(assigns, instance_assigns, registers, @rethrow_errors, @resource_limits)
       else
         raise ArgumentError, "Expect Hash or Liquid::Context as parameter"
       end

test/liquid/template_test.rb

@@ -71,4 +71,27 @@ class TemplateTest < Test::Unit::TestCase
     assert_equal '1', t.render(assigns)
     @global = nil
   end
+
+  def test_resource_limits
+    t = Template.parse("0123456789")
+    t.resource_limits = { :render_length_limit => 5 }
+    assert_raises(Liquid::MemoryError) { t.render() }
+    t.resource_limits = { :render_length_limit => 10 }
+    assert_equal "0123456789", t.render()
+    assert_not_nil t.resource_limits[:render_length_current]
+
+    t = Template.parse("{% for a in (1..100) %} foo {% endfor %}")
+    t.resource_limits = { :render_score_limit => 50 }
+    assert_raises(Liquid::MemoryError) { t.render() }
+    t.resource_limits = { :render_score_limit => 200 }
+    assert_equal (" foo " * 100), t.render()
+    assert_not_nil t.resource_limits[:render_score_current]
+
+    t = Template.parse("{% assign foo = 42 %}{% assign bar = 23 %}")
+    t.resource_limits = { :assign_score_limit => 1 }
+    assert_raises(Liquid::MemoryError) { t.render() }
+    t.resource_limits = { :assign_score_limit => 2 }
+    assert_equal "", t.render()
+    assert_not_nil t.resource_limits[:assign_score_current]
+  end
 end # TemplateTest