Merge pull request #198 from Shopify/limit_resource_usage

Resource usage limits

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,15 @@ 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)
+          if context.resource_limits_reached?
+            context.resource_limits[:reached] = true
+            raise MemoryError.new("Memory limits exceeded")
+          end
+          output << token_output
+        rescue 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)
+    def initialize(environments = {}, outer_scope = {}, registers = {}, rethrow_errors = false, resource_limits = {})
-      @environments   = [environments].flatten
+      @environments    = [environments].flatten
-      @scopes         = [(outer_scope || {})]
+      @scopes          = [(outer_scope || {})]
-      @registers      = registers
+      @registers       = registers
-      @errors         = []
+      @errors          = []
-      @rethrow_errors = rethrow_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
@@ -50,6 +50,7 @@ module Liquid
 
     # creates a new <tt>Template</tt> from an array of tokens. Use <tt>Template.parse</tt> instead
     def initialize
+      @resource_limits = {}
     end
 
     # Parse source code.
@@ -88,14 +89,14 @@ module Liquid
     #
     def render(*args)
       return '' if @root.nil?
-      
+
       context = case args.first
       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
@@ -120,9 +121,11 @@ module Liquid
 
       begin
         # render the nodelist.
-        # for performance reasons we get a array back here. join will make a string out of it
+        # for performance reasons we get an array back here. join will make a string out of it.
         result = @root.render(context)
         result.respond_to?(:join) ? result.join : result
+      rescue Liquid::MemoryError => e
+        context.handle_error(e)
       ensure
         @errors = context.errors
       end

test/liquid/template_test.rb

@@ -71,4 +71,53 @@ class TemplateTest < Test::Unit::TestCase
     assert_equal '1', t.render(assigns)
     @global = nil
   end
+
+  def test_resource_limits_render_length
+    t = Template.parse("0123456789")
+    t.resource_limits = { :render_length_limit => 5 }
+    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()
+    assert_not_nil t.resource_limits[:render_length_current]
+  end
+
+  def test_resource_limits_render_score
+    t = Template.parse("{% for a in (1..10) %} {% for a in (1..10) %} foo {% endfor %} {% endfor %}")
+    t.resource_limits = { :render_score_limit => 50 }
+    assert_equal "Liquid error: Memory limits exceeded", t.render()
+    assert t.resource_limits[:reached]
+    t = Template.parse("{% for a in (1..100) %} foo {% endfor %}")
+    t.resource_limits = { :render_score_limit => 50 }
+    assert_equal "Liquid error: Memory limits exceeded", t.render()
+    assert t.resource_limits[:reached]
+    t.resource_limits = { :render_score_limit => 200 }
+    assert_equal (" foo " * 100), t.render()
+    assert_not_nil t.resource_limits[:render_score_current]
+  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()
+    assert_not_nil t.resource_limits[:assign_score_current]
+  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 }
+    assert_equal "Liquid error: Memory limits exceeded", t.render()
+    assert t.resource_limits[:reached]
+  end
+
+  def test_resource_limits_hash_in_template_gets_updated_even_if_no_limits_are_set
+    t = Template.parse("{% for a in (1..100) %} {% assign foo = 1 %} {% endfor %}")
+    t.render()
+    assert t.resource_limits[:assign_score_current] > 0
+    assert t.resource_limits[:render_score_current] > 0
+    assert t.resource_limits[:render_length_current] > 0
+  end
 end # TemplateTest