immunio 0.15.2

data/lib/immunio/cli.rb

@@ -0,0 +1,39 @@
+require 'thor'
+require 'fileutils'
+
+module Immunio
+  class CLI < Thor
+    desc 'init', 'Initializes an Immunio configuration in the current app'
+    method_option :key,
+                  type: :string,
+                  desc: 'The key generated for your app in Immunio'
+    method_option :secret,
+                  type: :string,
+                  desc: 'The secret generated for your app in Immunio'
+    def init
+      if File.exist?(config_file) && File.read(config_file) =~ /key:\s+\w+|secret:\s+\w+/
+        say 'Immunio already initialized.', :green
+        Kernel.exit 0
+      end
+
+      key    = options[:key]    || ask('Enter the key generated for your app in Immunio:')
+      secret = options[:secret] || ask('Enter the secret generated for your app in Immunio:')
+
+      FileUtils.mkdir_p(File.dirname(config_file))
+
+      File.open(config_file, 'a') do |f|
+        f.puts "key: #{key}"
+        f.puts "secret: #{secret}"
+      end
+
+      say "Credentials written to #{config_file}", :green
+    end
+
+    private
+
+    def config_file
+      root = defined?(Rails) ? Rails.root : Dir.pwd
+      File.join(root, 'config', 'immunio.yml')
+    end
+  end
+end

data/lib/immunio/context.rb

@@ -0,0 +1,114 @@
+module Immunio
+  module Context
+    RAILS_TEMPLATE_FILTER = Regexp.new("(.*(_erb|_haml))__+\\d+_\\d+(.*)")
+    # Cache for contexts (named in tribute to our buddy Adam Back who invented proof of work)
+    @@hash_cache = {}
+
+    # Calculate context hashes and a stack trace. Additional data, in the form
+    # of a String, may be provided to mix into the strict context hash.
+    def self.context(additional_data=nil)
+      # We can filter out at least the top two frames
+      cache_key = Digest::SHA1.hexdigest(caller(2).join())
+      if @@hash_cache.has_key?(cache_key) then
+        loose_context = @@hash_cache[cache_key]["loose_context"]
+        strict_context = @@hash_cache[cache_key]["strict_context"]
+        stack = @@hash_cache[cache_key]["stack"]
+        loose_stack = @@hash_cache[cache_key]["stack"]
+
+        if Immunio.agent.config.log_context_data
+          Immunio.logger.info {"Stack contexts from cache"}
+        end
+      else
+        # Use ropes as they're faster than string concatenation
+        loose_stack_rope = []
+        loose_context_rope = []
+        stack_rope = []
+        strict_context_rope = []
+
+        # drop the top frame as it's us, but retain the rest. Immunio frames
+        # are filtered by the Gem regex.
+        locations = caller(1).map do |frame|
+          frame = frame.split(":", 3)
+          {path: frame[0], line: frame[1], label: frame[2]}
+        end
+
+        locations.each do |frame|
+
+          # Filter frame names from template rendering to remove generated random bits
+          matchdata = RAILS_TEMPLATE_FILTER.match(frame[:label])
+          if matchdata != nil then
+            frame[:label] = matchdata[1] + matchdata[3]
+          end
+
+          # Reduce paths to be relative to root if possible, to allow
+          # relocation. If there's no rails root, or the path doesn't start with
+          # the rails root, just use the filename part.
+          if defined?(Rails) && defined?(Rails.root) &&
+             Rails.root && frame[:path].start_with?(Rails.root.to_s)
+            strict_path = frame[:path].sub(Rails.root.to_s, '')
+          else
+            strict_path = File.basename(frame[:path])
+          end
+
+          stack_rope << "\n" unless stack_rope.empty?
+          stack_rope << frame[:path]
+          stack_rope << ":"
+          stack_rope << frame[:line]
+          stack_rope << ":"
+          stack_rope << frame[:label]
+
+          strict_context_rope << "\n" unless strict_context_rope.empty?
+          strict_context_rope << strict_path
+          strict_context_rope << ":"
+          strict_context_rope << frame[:line]
+          strict_context_rope << ":"
+          strict_context_rope << frame[:label]
+
+          # Remove pathname from the loose context. The goal here is to prevent
+          # upgrading gem versions from changing the loose context key, so for instance
+          # users don't have to rebuild their whitelists every time they update a gem
+          loose_context_rope << "\n" unless loose_context_rope.empty?
+          loose_context_rope << File.basename(frame[:path])
+          loose_context_rope << ":"
+          loose_context_rope << frame[:label]
+
+          # build a second seperate rope for the stack that determines ou loose context key
+          # This includes filenames for usability -- just method names not being very good
+          # for display purposes...
+          loose_stack_rope << "\n" unless loose_stack_rope.empty?
+          loose_stack_rope << frame[:path]
+          loose_stack_rope << ":"
+          loose_stack_rope << frame[:label]
+
+        end
+        stack = stack_rope.join()
+        strict_stack = strict_context_rope.join()
+        loose_stack = loose_stack_rope.join()
+
+        if Immunio.agent.config.log_context_data
+          Immunio.logger.info {"Strict context stack:\n#{strict_stack}"}
+          Immunio.logger.info {"Loose context stack:\n#{loose_stack}"}
+        end
+
+        strict_context = Digest::SHA1.hexdigest(strict_stack)
+        loose_context = Digest::SHA1.hexdigest(loose_context_rope.join())
+        @@hash_cache[cache_key] = {
+            "strict_context" => strict_context,
+            "loose_context" => loose_context,
+            "stack" => stack,
+            "loose_stack" => loose_stack
+          }
+      end
+
+      # Mix in additional context data
+      unless additional_data.nil?
+        if Immunio.agent.config.log_context_data
+          Immunio.logger.info {"Additional context data:\n#{additional_data}"}
+        end
+        strict_context = Digest::SHA1.hexdigest(strict_context + additional_data)
+      end
+
+      return strict_context, loose_context, stack, loose_stack
+    end
+  end
+end

data/lib/immunio/errors.rb

@@ -0,0 +1,43 @@
+require 'singleton'
+
+module Immunio
+  # General agent error
+  class Error < RuntimeError; end
+
+  # Error to block a request in progress
+  # Ruby's `SecurityError` is outside the `StandardError` hierarchy, so it will
+  # only be caught if somebody decides to catch it explicitly.
+  #
+  # `StandardError` and `RuntimeError` are both caught by a default rescue
+  # block. This makes it too easy for a developer to catch and ignore one of
+  # our blocking errors unintentionally.
+  class BlockError < SecurityError; end
+
+  # Request was not allowed by hook
+  class RequestBlocked < BlockError
+    class Result
+      include Singleton
+
+      # PG activerecord exception results must have an error_field method before Rails 4.x
+      def error_field(_field)
+        ""
+      end
+    end
+
+    # PG activerecord exceptions must have a result method before Rails 4.x
+    def result
+      Result.instance
+    end
+  end
+
+  # Response overridden by hook
+  class OverrideResponse < BlockError
+    attr_reader :status, :headers, :body
+
+    def initialize(status, headers, body)
+      @status = status
+      @headers = headers
+      @body = body
+    end
+  end
+end

data/lib/immunio/immunio_ca.crt

@@ -0,0 +1,45 @@
+-----BEGIN CERTIFICATE-----
+MIIDrzCCApegAwIBAgIQCDvgVpBCRrGhdWrJWZHHSjANBgkqhkiG9w0BAQUFADBh
+MQswCQYDVQQGEwJVUzEVMBMGA1UEChMMRGlnaUNlcnQgSW5jMRkwFwYDVQQLExB3
+d3cuZGlnaWNlcnQuY29tMSAwHgYDVQQDExdEaWdpQ2VydCBHbG9iYWwgUm9vdCBD
+QTAeFw0wNjExMTAwMDAwMDBaFw0zMTExMTAwMDAwMDBaMGExCzAJBgNVBAYTAlVT
+MRUwEwYDVQQKEwxEaWdpQ2VydCBJbmMxGTAXBgNVBAsTEHd3dy5kaWdpY2VydC5j
+b20xIDAeBgNVBAMTF0RpZ2lDZXJ0IEdsb2JhbCBSb290IENBMIIBIjANBgkqhkiG
+9w0BAQEFAAOCAQ8AMIIBCgKCAQEA4jvhEXLeqKTTo1eqUKKPC3eQyaKl7hLOllsB
+CSDMAZOnTjC3U/dDxGkAV53ijSLdhwZAAIEJzs4bg7/fzTtxRuLWZscFs3YnFo97
+nh6Vfe63SKMI2tavegw5BmV/Sl0fvBf4q77uKNd0f3p4mVmFaG5cIzJLv07A6Fpt
+43C/dxC//AH2hdmoRBBYMql1GNXRor5H4idq9Joz+EkIYIvUX7Q6hL+hqkpMfT7P
+T19sdl6gSzeRntwi5m3OFBqOasv+zbMUZBfHWymeMr/y7vrTC0LUq7dBMtoM1O/4
+gdW7jVg/tRvoSSiicNoxBN33shbyTApOB6jtSj1etX+jkMOvJwIDAQABo2MwYTAO
+BgNVHQ8BAf8EBAMCAYYwDwYDVR0TAQH/BAUwAwEB/zAdBgNVHQ4EFgQUA95QNVbR
+TLtm8KPiGxvDl7I90VUwHwYDVR0jBBgwFoAUA95QNVbRTLtm8KPiGxvDl7I90VUw
+DQYJKoZIhvcNAQEFBQADggEBAMucN6pIExIK+t1EnE9SsPTfrgT1eXkIoyQY/Esr
+hMAtudXH/vTBH1jLuG2cenTnmCmrEbXjcKChzUyImZOMkXDiqw8cvpOp/2PV5Adg
+06O/nVsJ8dWO41P0jmP6P6fbtGbfYmbW0W5BjfIttep3Sp+dWOIrWcBAI+0tKIJF
+PnlUkiaY4IBIqDfv8NZ5YBberOgOzW6sRBc4L0na4UU+Krk2U886UAb3LujEV0ls
+YSEY1QSteDwsOoBrp+uvFRTp2InBuThs4pFsiv9kuXclVzDAGySj4dzp30d8tbQk
+CAUw7C29C79Fv1C5qfPrmAESrciIxpg0X40KPMbp1ZWVbd4=
+-----END CERTIFICATE-----
+
+-----BEGIN CERTIFICATE-----
+MIIDljCCAn6gAwIBAgIBATANBgkqhkiG9w0BAQUFADBcMQswCQYDVQQGEwJDQTEW
+MBQGA1UEChMNSW1tdW4uaW8gSW5jLjEVMBMGA1UECxMMd3d3LmltbXVuLmlvMR4w
+HAYDVQQDExVJbW11bi5pbyBJbmMuIFJvb3QgQ0EwHhcNMTQwMzI1MTQ0NTI3WhcN
+MjQwMzIyMTQ0NTI3WjBcMQswCQYDVQQGEwJDQTEWMBQGA1UEChMNSW1tdW4uaW8g
+SW5jLjEVMBMGA1UECxMMd3d3LmltbXVuLmlvMR4wHAYDVQQDExVJbW11bi5pbyBJ
+bmMuIFJvb3QgQ0EwggEiMA0GCSqGSIb3DQEBAQUAA4IBDwAwggEKAoIBAQDomhh6
+lGHL6IsOlK7TmFikinZ0ShPQQ8WHFbwoLiELGJSNGFKdiSnQICOkjTI6kdXCiOgk
+TYARs8Ty7e1rbCTiBZUV2d2Q1qktS+wHv6GYEukjkX+/yLLf65XsNQlbPheFuBCG
+Tvy8qU8PbdXQu35zLuCwpq8DAanCEXWANOAIYXOaIa0DMqRrVG5QeSfi5mxeXL5q
+Xb4FhqQbMwX7xkQiIB3NQCmVtplDSdMPfCvg/T97rC14XR8jKteRe2OkaLFeGHZp
+mROT2k4dTY4r8h2dV3EW8N4pQizDVpUfzNrgwoaKGccg6JiLLVeMQT7BJxjEi3Tt
+tSt7gf0cWGwMWHbbAgMBAAGjYzBhMA8GA1UdEwEB/wQFMAMBAf8wDgYDVR0PAQH/
+BAQDAgEGMB0GA1UdDgQWBBRUsJtzljk4lIWaEpQXfTR48b70SDAfBgNVHSMEGDAW
+gBRUsJtzljk4lIWaEpQXfTR48b70SDANBgkqhkiG9w0BAQUFAAOCAQEAS0PyWyXj
+1oF/pWKhY6g71UlcnnISXoutIEQxK0cxqzwjCd/62W5SlgxpSFHJmeuhMHKgT1XF
+PgretLnGSXft1ZmpAtUmHXJYxSFbHz8kv/S5rYW7GeJRTUEqJNGly9nPiAkbY2Uk
+nZSF2uY62CztoDTtESunAcSPN3BxoEEQ3GLg+PFJVFVApyXfjdK5Jc9lfMGH6vIX
+hnW1BIYyyHqaerL2VSFJBaxqDk9NeUmt3w77EdJ0CpRo03Gbw8g19aIrXr9u25lW
+Zm2+58fpRuZ2pdclnTWfjyKb00WCkqfInFLeMvXIEhA60TirhlPBRZ5+8D42I5Sd
+AcWwGci7HgH1+A==
+-----END CERTIFICATE-----

data/lib/immunio/logger.rb

@@ -0,0 +1,87 @@
+require 'logger'
+require 'stringio'
+
+module Immunio
+  # Subclass global Logger class to add TRACE level
+  class Logger < ::Logger
+    module Severity
+      TRACE = -1
+    end
+    include Severity
+
+    def trace?; @level <= TRACE; end
+
+    def trace(progname = nil, &block)
+      add(TRACE, nil, progname, &block)
+    end
+
+    def format_severity(severity)
+      SEV_LABEL[severity] || 'ANY'
+    end
+
+  private
+    SEV_LABEL = Array.new(::Logger::SEV_LABEL)
+    SEV_LABEL[-1] = 'TRACE'
+  end
+
+  attr_reader :logger
+
+  def self.create_startup_logger
+    @startup_messages = StringIO.new
+    @logger = Logger.new @startup_messages
+
+    setup_logger_formatter
+  end
+
+  def self.setup_logger_formatter
+    logger.formatter = proc do |severity, datetime, _progname, msg|
+      "[#{datetime}] #{severity}: #{msg}\n"
+    end
+  end
+
+  def self.switch_to_real_logger(log_file, log_level)
+    # Have we already switched to real logger?
+    return if !defined?(@startup_messages)
+
+    if log_file == "STDOUT"
+      @logger = Logger.new $stdout
+    elsif log_file == "STDERR"
+      @logger = Logger.new $stderr
+    else
+      path = Pathname.new(log_file)
+      begin
+        FileUtils.mkdir_p path.dirname unless File.exist? path.dirname
+
+        file = File.open path, 'a'
+        file.binmode
+        file.sync = true
+
+        @logger = Logger.new file
+        log_file = path.realpath
+      rescue StandardError => e
+        logger.warn "Failed to open #{log_file} (#{path.realdirpath}) for logging (#{e.message})"
+        @logger = Logger.new $stderr
+        log_file = "STDERR"
+      end
+    end
+
+    # Dump saved log messages during startup to real log
+    logger << @startup_messages.string
+    remove_instance_variable(:@startup_messages)
+
+    setup_logger_formatter
+
+    begin
+      logger.level = Logger.const_get(log_level.to_s.upcase)
+    rescue
+      logger.level = Logger::DEBUG
+      logger.debug "Failed to interpret log level #{log_level}, falling back to debug"
+    end
+
+    logger.debug "Logging to #{log_file}"
+  end
+
+  def self.logger
+    @logger
+  end
+end

data/lib/immunio/plugins/action_dispatch.rb

@@ -0,0 +1,45 @@
+module Immunio
+  module CookieHooks
+    extend ActiveSupport::Concern
+
+    included do
+      # TODO: should anything be checked to make sure @parent_jar exists
+      if method_defined? :[]  # Not sure when this wouldn't exist.
+        # The following won't work because of the names:
+        # alias_method_chain :[], :immunio if method_defined? :[]
+        alias_method :lookup_without_immunio, :[]
+        alias_method :[], :lookup_with_immunio
+      end
+
+    end
+
+    def lookup_with_immunio(name)
+      Request.time "plugin", "#{Module.nesting[0]}::#{__method__}" do
+        raw_cookie_value = @parent_jar[name]
+        cookie_value = Request.pause "plugin", "#{Module.nesting[0]}::#{__method__}" do
+          lookup_without_immunio(name)
+        end
+        if !raw_cookie_value.nil? and cookie_value.nil?
+          Immunio.run_hook! "action_dispatch", "bad_cookie", key: name,
+                                                             value: raw_cookie_value
+        end
+        cookie_value
+      end
+    end
+  end
+end
+
+class ActionDispatch::Cookies
+  if defined? SignedCookieJar
+    SignedCookieJar.send :include, Immunio::CookieHooks
+  end
+  if defined? UpgradeLegacySignedCookieJar
+    UpgradeLegacySignedCookieJar.send :include, Immunio::CookieHooks
+  end
+  if defined? EncryptedCookieJar
+    EncryptedCookieJar.send :include, Immunio::CookieHooks
+  end
+  if defined? UpgradeLegacyEncryptedCookieJar
+    UpgradeLegacyEncryptedCookieJar.send :include, Immunio::CookieHooks
+  end
+end

data/lib/immunio/plugins/action_view.rb

@@ -0,0 +1,431 @@
+# Hook into ActionView rendering to inject Immunio's hooks.
+require 'securerandom'
+
+module Immunio
+  # Renders templates by filtering them through Immunio's hook handlers.
+  class Template
+    attr_accessor :vars
+
+    def initialize(template)
+      @template = template
+      @next_var_id = 0
+      @next_template_id = 0
+      @vars = {}
+      @scheduled_fragments_writes = []
+    end
+
+    def id
+      (@template.respond_to?(:virtual_path) && @template.virtual_path) || (@template.respond_to?(:source) && @template.source)
+    end
+
+    def ==(other)
+      self.class === other && id == other.id
+    end
+
+    def has_source?
+      @template.respond_to?(:source) && !@template.source.nil?
+    end
+
+    def is_text?
+      @template.formats.first == :text
+    end
+
+    def load_source(context)
+      return if !@template.respond_to?(:source) || !@template.source.nil?
+
+      # @template is a virtual template that doesn't contain the source. We need
+      # to try to load the source. But, the virtual template doesn't know the
+      # original format of the source template file. Grab the original format
+      # from the view context and override the default of just :html when
+      # when looking up the template.
+      old_formats = context.lookup_context.formats
+      begin
+        context.lookup_context.formats = @template.formats
+        refreshed = @template.refresh(context)
+      ensure
+        context.lookup_context.formats = old_formats
+      end
+
+      return if refreshed.nil?
+
+      @template.instance_variable_set :@source, refreshed.source
+    end
+
+    def template_sha
+      # A template might have a source but it might be nil.
+      @template_sha ||= begin
+        Digest::SHA1.hexdigest(@template.source) if has_source?
+      end
+    end
+
+    def compiled?
+      @template.instance_variable_get :@compiled
+    end
+
+    # Generate the next var unique ID to be used in a template.
+    def next_var_id
+      id = @next_var_id
+      @next_var_id += 1
+      id
+    end
+
+    def next_template_id
+      id = @next_template_id
+      @next_template_id += 1
+      id
+    end
+
+    def get_nonce
+      # Generate a two byte CSRNG nonce to make our substitutions unpreictable
+      # Why only 2 bytes? The nonce is per render, so the odds of guessing it are very low
+      # and entropy is finite so we don't want to drain the random pool unnecessarily
+      @nonce ||= SecureRandom.hex(2)
+    end
+
+    def mark_var(content, code, template_id, file, line, escape)
+      id = Template.next_var_id
+      nonce = Template.get_nonce
+      Template.vars[id.to_s] = {
+        template_sha: template_sha,
+        template_id: template_id.to_s,
+        nonce: nonce,
+        code: code,
+        file: file,
+        line: line
+      }
+
+      rval = ""
+      # NOTE: What happens here is pretty funky to preserve the html_safe SafeBuffer behaviour in ruby.
+      # If escaped is true we directly concatenate the content between two SafeBuffers. This will cause
+      # escaping if content is not itself a SafeBuffer.
+      # Otherwise we explicitly convert to a string, and convert that to a SafeBuffer to ensure that
+      # for instance no escaping is performed on the contents of a <%== %> Erubis interpolation.
+      if escape and not is_text? then
+        # explicitly convert (w/ escapes) and mark safe things that aren't String (SafeBuffer is_a String also)
+        # `to_s` is used to render any object passed to a template.
+        # It is called internally when appending to ActionView::OutputBuffer.
+        # We force rendering to get the actual string.
+        # This has no impact if `rendered` is already a string.
+        content = content.to_s.html_safe unless content.is_a? String
+        # As a failsafe, just return the content if it already contains our markers. This can occur when
+        # a helper calls render partial to generate a component of a page. Both render calls are root level
+        # templates from our perspective.
+        if content =~ /\{immunio-var:\d+:#{nonce}\}/ then
+          # don't add markers.
+          Immunio.logger.debug {"WARNING: ActionView not marking interpolation which already contains markers: \"#{content}\""}
+          rval = content
+        else
+          rval = "{immunio-var:#{id}:#{nonce}}".html_safe + content + "{/immunio-var:#{id}:#{nonce}}".html_safe
+        end
+      else
+        content = "" if content.nil?
+        # See comment above
+        if content =~ /\{immunio-var:\\d+:#{nonce}\}/ then
+          # don't add markers.
+          Immunio.logger.debug {"WARNING: ActionView not marking interpolation which already contains markers: \"#{content}\""}
+          rval = content.html_safe
+        else
+          rval = "{immunio-var:#{id}:#{nonce}}".html_safe + content.html_safe + "{/immunio-var:#{id}:#{nonce}}".html_safe
+        end
+      end
+      rval
+    end
+
+    def mark_and_defer_fragment_write(key, content, options)
+      id = @scheduled_fragments_writes.size
+      nonce = Template.get_nonce
+      @scheduled_fragments_writes << [key, content, options]
+      "{immunio-fragment:#{id}:#{nonce}}#{content}{/immunio-fragment:#{id}:#{nonce}}"
+    end
+
+    def render(context)
+      load_source context
+      # Don't handle templates with no source (inline text templates).
+      if not has_source? then
+        rendered = yield
+        rendered.instance_variable_set("@__immunio_processed", true)
+        return rendered
+      end
+
+      begin
+        root = true if rendering_stack.length == 0
+
+        rendering_stack.push self
+        # Calculate SHA1 of this template.
+        template_sha
+        Immunio.logger.debug {"ActionView rendering template with sha #{@template_sha}, root: #{root}"}
+        rendered = yield
+        rendered.instance_variable_set("@__immunio_processed", true)
+
+        if root
+          # This is the root template. Let ActionView render it, and then look
+          # for XSS.
+          rendered = rendered.to_str
+          # Rendering done!
+          result = run_hook! "template_render_done", rendered: rendered, vars: @vars
+
+          # We use the return value from the hook handler if present.
+          rendered = result.fetch("rendered") { rendered.dup }
+
+          remove_var_markers! rendered
+
+          # If some fragments were marked to be cached, commit their content to cache.
+          write_and_remove_fragments! context, rendered
+
+          rendered.html_safe
+        else
+          # This is a partial template. Just render it.
+          rendered
+        end
+      ensure
+        top_template = rendering_stack.pop
+        unless top_template == self
+          raise Error, "Unexpected Immunio::Template on rendering stack. Expected #{id}, got #{top_template.try :id}."
+        end
+      end
+    end
+
+    # Generate code injected in templates to wrap everything inside `<%= ... %>`.
+    def self.generate_render_var_code(code, escape)
+      template = Template.current
+      if template
+        template_id = template.next_template_id
+        "(__immunio_result = (#{code}); Immunio::Template.render_var(#{code.strip.inspect}, __immunio_result, #{template_id}, __FILE__, __LINE__, #{escape}))"
+      else
+        code
+      end
+    end
+
+    def self.render_var(code, rendered, template_id, file, line, escape)
+      if rendered.instance_variable_get("@__immunio_processed") then
+        # Ignore buffers marked as __immunio_processed in render as these are full templates or partials
+        return rendered
+      elsif code =~ /yield( .*)?/
+        # Ignore yielded blocks inside layouts
+        return rendered
+      end
+      template = Template.current
+      if template
+        rendered = template.mark_var rendered, code, template_id, file, line, escape
+      end
+      rendered.html_safe
+    end
+
+    def self.current
+      rendering_stack.last
+    end
+
+    def self.next_var_id
+      rendering_stack.first.next_var_id
+    end
+
+    def self.vars
+      rendering_stack.first.vars
+    end
+
+    def self.get_nonce
+      rendering_stack.first.get_nonce
+    end
+
+    # Save fragment info to the root template only
+    def self.mark_and_defer_fragment_write(*args)
+      rendering_stack.first.mark_and_defer_fragment_write(*args)
+    end
+
+    private
+      # Stack of the templates currently being rendered.
+      def self.rendering_stack
+        Thread.current["immunio.rendering_stack"] ||= []
+      end
+
+      def rendering_stack
+        self.class.rendering_stack
+      end
+
+      def run_hook!(name, meta={})
+        default_meta = {
+          template_sha: template_sha,
+          name: (@template.respond_to?(:virtual_path) && @template.virtual_path) || nil,
+          origin: @template.identifier,
+          nonce: Template.get_nonce
+        }
+        Immunio.run_hook! "action_view", name, default_meta.merge(meta)
+      end
+
+      def write_and_remove_fragments!(context, content)
+        # Rails tests do use the context as the view context sometimes.
+        if context.is_a? ActionController::Base
+          controller = context
+        elsif context.respond_to? :controller
+          controller = context.controller
+        else
+          # Some rails unit tests don't have a controller...
+          remove_all_markers! content
+          return
+        end
+
+        # Iterate to handle nested fragments. Child fragments have lower ids than their parents.
+        nonce = Template.get_nonce
+        @scheduled_fragments_writes.each_with_index do |(key, _, options), id|
+          # Remove the markers ...
+          content.sub!(/\{immunio-fragment:#{id}:#{nonce}\}(.*)\{\/immunio-fragment:#{id}:#{nonce}\}/m) do
+            # The escaped content inside the markers ($1), is written to cache.
+            output = $1
+            remove_all_markers! output
+            controller.write_fragment_without_immunio key, output, options
+            output
+          end
+        end
+        # To be extra safe strip all markers from content
+        remove_all_markers! content
+      end
+
+      def remove_var_markers!(input)
+        nonce = Template.get_nonce
+        # TODO is this the fastest way to remove the markers? Needs benchmarking ...
+        input.gsub!(/\{\/?immunio-var:\d+:#{nonce}\}/, "")
+      end
+
+      def remove_all_markers!(input)
+        input.gsub!(/\{\/?immunio-(fragment|var):\d+:[a-zA-Z0-9]+\}/, "")
+      end
+  end
+
+  # Regexp to test for blocks (... do) in the Ruby code of templates.
+  BLOCK_EXPR = ActionView::Template::Handlers::Erubis::BLOCK_EXPR
+
+  # Hooks for the ERB template engine.
+  # (Default one used in Rails).
+  module ErubisHooks
+    extend ActiveSupport::Concern
+
+    included do
+      alias_method_chain :add_expr, :immunio
+    end
+
+    def add_expr_with_immunio(src, code, indicator)
+      # Wrap expressions in the templates to track their rendered value.
+      # Do not wrap expressions with blocks, eg.: <%= form_tag do %>
+      # TODO should we support blocks?
+      Request.time "plugin", "#{Module.nesting[0]}::#{__method__}" do
+        unless code =~ BLOCK_EXPR
+          # escape unless we see the == indicator
+          escape = !(indicator == '==')
+          code = Immunio::Template.generate_render_var_code(code, escape)
+        end
+        Request.pause "plugin", "#{Module.nesting[0]}::#{__method__}" do
+          add_expr_without_immunio(src, code, indicator)
+        end
+      end
+    end
+  end
+
+  # Hooks for the HAML template engine.
+  module HamlHooks
+    extend ActiveSupport::Concern
+
+    included do
+      alias_method_chain :push_script, :immunio
+    end
+
+    def push_script_with_immunio(code, opts = {}, &block)
+      # Wrap expressions in the templates to track their rendered value.
+      Request.time "plugin", "#{Module.nesting[0]}::#{__method__}" do
+        if code !~ BLOCK_EXPR
+          # escape if we're told to by HAML
+          code = Immunio::Template.generate_render_var_code(code, opts[:escape_html])
+        end
+        Request.pause "plugin", "#{Module.nesting[0]}::#{__method__}" do
+          push_script_without_immunio(code, opts, &block)
+        end
+      end
+    end
+  end
+
+  # Hook for the `ActionView::TemplateRenderer`. These are called for root
+  # templates.
+  module TemplateRendererHooks
+    extend ActiveSupport::Concern
+
+    included do
+      alias_method_chain :render_template, :immunio
+    end
+
+    def render_template_with_immunio(template, *args)
+      Request.time "plugin", "#{Module.nesting[0]}::#{__method__}" do
+        renderer = Template.new(template)
+
+        renderer.render @view do
+          Request.pause "plugin", "#{Module.nesting[0]}::#{__method__}" do
+            render_template_without_immunio(template, *args)
+          end
+        end
+      end
+    end
+  end
+
+  # Hook for the `ActionView::Template`. These are called for non-root
+  # templates.
+  module TemplateHooks
+    extend ActiveSupport::Concern
+
+    included do
+      alias_method_chain :render, :immunio
+    end
+
+    def render_with_immunio(context, *args, &block)
+      Request.time "plugin", "#{Module.nesting[0]}::#{__method__}" do
+        renderer = Template.new(self)
+
+        renderer.render context do
+          Request.pause "plugin", "#{Module.nesting[0]}::#{__method__}" do
+            render_without_immunio(context, *args, &block)
+          end
+        end
+      end
+    end
+  end
+
+  # Hook for `ActionController::Caching::Fragments` responsible for handling the `<% cache do %>...` in templates.
+  module FragmentCachingHooks
+    extend ActiveSupport::Concern
+
+    included do
+      alias_method_chain :write_fragment, :immunio
+    end
+
+    def write_fragment_with_immunio(key, content, options = nil)
+      return content unless cache_configured?
+
+      template = Template.current
+      if template
+        # We're rendering a template. Defer caching 'till we get the escaped content from the hook handler.
+        content = Template.mark_and_defer_fragment_write(key, content, options)
+      else
+        # Not rendering a template. Ignore.
+        # Shouldn't happen. But, just to be safe in case fragment caching is used in the controller for something else.
+        content = write_fragment_without_immunio(key, content, options)
+      end
+
+      content
+    end
+  end
+end
+
+# Add XSS hooks if enabled
+if Immunio::agent.plugin_enabled?("xss") then
+  # Hook into template engines.
+  ActionView::Template::Handlers::Erubis.send :include, Immunio::ErubisHooks
+
+  ActiveSupport.on_load(:after_initialize) do
+    # Wait after Rails initialization to patch custom template engines.
+    if defined? Haml::Compiler
+      Haml::Compiler.send :include, Immunio::HamlHooks
+    end
+  end
+
+  # Hook into rendering process of Rails.
+  ActionView::TemplateRenderer.send :include, Immunio::TemplateRendererHooks
+  ActionView::Template.send :include, Immunio::TemplateHooks
+  ActionController::Caching::Fragments.send :include, Immunio::FragmentCachingHooks
+end