From 8760b5e8c4e2c29a1f8714dcf491cef69c606ecf Mon Sep 17 00:00:00 2001
From: Florian Weingarten <flo@shopify.com>
Date: Thu, 30 May 2013 12:01:15 -0400
Subject: [PATCH 1/4] 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          | 10 +++++++++-
 lib/liquid/context.rb        | 21 ++++++++++++++-------
 lib/liquid/errors.rb         |  3 ++-
 lib/liquid/tags/assign.rb    |  6 ++++--
 lib/liquid/tags/capture.rb   |  1 +
 lib/liquid/template.rb       |  6 +++---
 test/liquid/template_test.rb | 23 +++++++++++++++++++++++
 7 files changed, 56 insertions(+), 14 deletions(-)

diff --git a/lib/liquid/block.rb b/lib/liquid/block.rb
index 8ea87b08..9f888811 100644
--- a/lib/liquid/block.rb
+++ b/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
diff --git a/lib/liquid/context.rb b/lib/liquid/context.rb
index 129b71a8..73a62663 100644
--- a/lib/liquid/context.rb
+++ b/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
diff --git a/lib/liquid/errors.rb b/lib/liquid/errors.rb
index b0add6f6..85cb3731 100644
--- a/lib/liquid/errors.rb
+++ b/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
diff --git a/lib/liquid/tags/assign.rb b/lib/liquid/tags/assign.rb
index 3540b76c..70b49cec 100644
--- a/lib/liquid/tags/assign.rb
+++ b/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  
diff --git a/lib/liquid/tags/capture.rb b/lib/liquid/tags/capture.rb
index 2f67a0b2..495a6f75 100644
--- a/lib/liquid/tags/capture.rb
+++ b/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
diff --git a/lib/liquid/template.rb b/lib/liquid/template.rb
index 1d019825..43c9318a 100644
--- a/lib/liquid/template.rb
+++ b/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
diff --git a/test/liquid/template_test.rb b/test/liquid/template_test.rb
index 92803ea5..4e912c98 100644
--- a/test/liquid/template_test.rb
+++ b/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

From 9075b428b143827c467e33fec2816423b1265b88 Mon Sep 17 00:00:00 2001
From: Florian Weingarten <flo@shopify.com>
Date: Fri, 31 May 2013 09:25:25 -0400
Subject: [PATCH 2/4] Resource limits: Don't raise Error but render error
 message (but abort after first error)

---
 lib/liquid/block.rb          |  4 ++--
 lib/liquid/template.rb       |  4 +++-
 test/liquid/template_test.rb | 21 +++++++++++++++++----
 3 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/lib/liquid/block.rb b/lib/liquid/block.rb
index 9f888811..efd149be 100644
--- a/lib/liquid/block.rb
+++ b/lib/liquid/block.rb
@@ -108,9 +108,9 @@ module Liquid
 
           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?
+          raise MemoryError.new("Memory limits exceeded") if context.resource_limits_reached?
           output << token_output
-        rescue Liquid::MemoryError => e
+        rescue MemoryError => e
           raise e
         rescue ::StandardError => e
           output << (context.handle_error(e))
diff --git a/lib/liquid/template.rb b/lib/liquid/template.rb
index 43c9318a..10a78049 100644
--- a/lib/liquid/template.rb
+++ b/lib/liquid/template.rb
@@ -120,9 +120,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
diff --git a/test/liquid/template_test.rb b/test/liquid/template_test.rb
index 4e912c98..bc1c9a32 100644
--- a/test/liquid/template_test.rb
+++ b/test/liquid/template_test.rb
@@ -72,26 +72,39 @@ class TemplateTest < Test::Unit::TestCase
     @global = nil
   end
 
-  def test_resource_limits
+  def test_resource_limits_render_length
     t = Template.parse("0123456789")
     t.resource_limits = { :render_length_limit => 5 }
-    assert_raises(Liquid::MemoryError) { t.render() }
+    assert_equal "Liquid error: Memory limits exceeded", t.render()
     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()
     t = Template.parse("{% for a in (1..100) %} foo {% endfor %}")
     t.resource_limits = { :render_score_limit => 50 }
-    assert_raises(Liquid::MemoryError) { t.render() }
+    assert_equal "Liquid error: Memory limits exceeded", t.render()
     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_raises(Liquid::MemoryError) { t.render() }
+    assert_equal "Liquid error: Memory limits exceeded", t.render()
     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()
+  end
 end # TemplateTest

From 2b17e24b16fdbeb887d37e4403993321965a7212 Mon Sep 17 00:00:00 2001
From: Florian Weingarten <flo@shopify.com>
Date: Fri, 31 May 2013 09:34:23 -0400
Subject: [PATCH 3/4] Mutate resource_limits hash to flag that the limit was
 reached (for outside observation)

---
 lib/liquid/block.rb          | 5 ++++-
 test/liquid/template_test.rb | 5 +++++
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/lib/liquid/block.rb b/lib/liquid/block.rb
index efd149be..a0a07e49 100644
--- a/lib/liquid/block.rb
+++ b/lib/liquid/block.rb
@@ -108,7 +108,10 @@ module Liquid
 
           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 MemoryError.new("Memory limits exceeded") if context.resource_limits_reached?
+          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
diff --git a/test/liquid/template_test.rb b/test/liquid/template_test.rb
index bc1c9a32..9a04e884 100644
--- a/test/liquid/template_test.rb
+++ b/test/liquid/template_test.rb
@@ -76,6 +76,7 @@ class TemplateTest < Test::Unit::TestCase
     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]
@@ -85,9 +86,11 @@ class TemplateTest < Test::Unit::TestCase
     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]
@@ -97,6 +100,7 @@ class TemplateTest < Test::Unit::TestCase
     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]
@@ -106,5 +110,6 @@ class TemplateTest < Test::Unit::TestCase
     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
 end # TemplateTest

From 1e8c081b428268ddb8bd23492236e65f2fee8a18 Mon Sep 17 00:00:00 2001
From: Florian Weingarten <flo@shopify.com>
Date: Fri, 31 May 2013 09:41:59 -0400
Subject: [PATCH 4/4] Create new resource_limits hash on Template
 initialization

---
 lib/liquid/template.rb       | 3 ++-
 test/liquid/template_test.rb | 8 ++++++++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/lib/liquid/template.rb b/lib/liquid/template.rb
index 10a78049..a4895d82 100644
--- a/lib/liquid/template.rb
+++ b/lib/liquid/template.rb
@@ -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,7 +89,7 @@ module Liquid
     #
     def render(*args)
       return '' if @root.nil?
-      
+
       context = case args.first
       when Liquid::Context
         args.shift
diff --git a/test/liquid/template_test.rb b/test/liquid/template_test.rb
index 9a04e884..6fb68e8c 100644
--- a/test/liquid/template_test.rb
+++ b/test/liquid/template_test.rb
@@ -112,4 +112,12 @@ class TemplateTest < Test::Unit::TestCase
     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