Middleware to highlight twice-escaped HTML

gistfile1.rb

class HighlightEscapedHTML
  def initialize(app)
    @app = app
  end
  
  def call(env)
    dup._call(env)
  end
  
  def _call(env)
    status, @headers, @response = @app.call(env)
    
    response = []
    @headers['Content-Length'] = each { |chunk| response << chunk }.to_s
    
    [status, @headers, response]
  end
  
  def each(&block)
    if @headers["Content-Type"].include? "text/html"
      @escaped_html_size = @escaped_html_count = 0
      
      @response.each do |chunk|
        block.call highlight_escaped_html(chunk)
      end
      
      length = @headers['Content-Length'].to_i
      
      if @escaped_html_count > 0
        length += @escaped_html_size
        
        summary = "<script>document.title += ' — #{@escaped_html_count} XSS bugs';</script>"
        block.call(summary)
        length += summary.size
      end
      
      length
    else
      @response.each(&block)
      @headers['Content-Length']
    end
  end
  
  private
  
  def highlight_escaped_html content
    content.gsub(/&lt;\/?\w.*?&gt;|&amp;\w+;/) do |escaped_html|
      hint = "<span style=\"color: red; background: yellow; text-decoration: blink;\">#{escaped_html}<script>console.log('#{escaped_html}');</script></span>"
      @escaped_html_count += 1
      @escaped_html_size += hint.size - escaped_html.size
      hint
    end
  end
end