From f8a4edc9cd3214b5d2647c09bfa5c9dadb287e0a Mon Sep 17 00:00:00 2001
From: Robin Daugherty <robin@robindaugherty.net>
Date: Sun, 13 Sep 2020 10:22:42 -0400
Subject: [PATCH] Add CSRF to error page and internal calls

---
 lib/better_errors/error_page.rb      |  2 +-
 lib/better_errors/middleware.rb      | 40 +++++++++++++++++++++++-----
 lib/better_errors/templates/main.erb |  2 ++
 3 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/lib/better_errors/error_page.rb b/lib/better_errors/error_page.rb
index 1f049f39..05a7e7bf 100644
--- a/lib/better_errors/error_page.rb
+++ b/lib/better_errors/error_page.rb
@@ -26,7 +26,7 @@ def id
       @id ||= SecureRandom.hex(8)
     end
 
-    def render(template_name = "main")
+    def render(template_name = "main", csrf_token = nil)
       binding.eval(self.class.template(template_name).src)
     rescue => e
       # Fix the backtrace, which doesn't identify the template that failed (within Better Errors).
diff --git a/lib/better_errors/middleware.rb b/lib/better_errors/middleware.rb
index 262cf3b2..b8cc1c91 100644
--- a/lib/better_errors/middleware.rb
+++ b/lib/better_errors/middleware.rb
@@ -39,6 +39,8 @@ def self.allow_ip!(addr)
     allow_ip! "127.0.0.0/8"
     allow_ip! "::1/128" rescue nil # windows ruby doesn't have ipv6 support
 
+    CSRF_TOKEN_COOKIE_NAME = 'BetterErrors-CSRF-Token'
+
     # A new instance of BetterErrors::Middleware
     #
     # @param app      The Rack app/middleware to wrap with Better Errors
@@ -89,11 +91,14 @@ def protected_app_call(env)
     end
 
     def show_error_page(env, exception=nil)
+      request = Rack::Request.new(env)
+      csrf_token = request.cookies[CSRF_TOKEN_COOKIE_NAME] || SecureRandom.uuid
+
       type, content = if @error_page
         if text?(env)
           [ 'plain', @error_page.render('text') ]
         else
-          [ 'html', @error_page.render ]
+          [ 'html', @error_page.render('main', csrf_token) ]
         end
       else
         [ 'html', no_errors_page ]
@@ -104,12 +109,21 @@ def show_error_page(env, exception=nil)
         status_code = ActionDispatch::ExceptionWrapper.new(env, exception).status_code
       end
 
-      [status_code, { "Content-Type" => "text/#{type}; charset=utf-8" }, [content]]
+      headers = {
+        "Content-Type" => "text/#{type}; charset=utf-8",
+      }
+      response = Rack::Response.new(content, status_code, headers)
+
+      unless request.cookies[CSRF_TOKEN_COOKIE_NAME]
+        response.set_cookie(CSRF_TOKEN_COOKIE_NAME, value: csrf_token, httponly: true, same_site: :strict)
+      end
+
+      response.finish
     end
 
     def text?(env)
       env["HTTP_X_REQUESTED_WITH"] == "XMLHttpRequest" ||
-      !env["HTTP_ACCEPT"].to_s.include?('html')
+        !env["HTTP_ACCEPT"].to_s.include?('html')
     end
 
     def log_exception
@@ -133,9 +147,15 @@ def internal_call(env, opts)
       return no_errors_json_response unless @error_page
       return invalid_error_json_response if opts[:id] != @error_page.id
 
-      env["rack.input"].rewind
-      response = @error_page.send("do_#{opts[:method]}", JSON.parse(env["rack.input"].read))
-      [200, { "Content-Type" => "text/plain; charset=utf-8" }, [JSON.dump(response)]]
+      request = Rack::Request.new(env)
+      return invalid_csrf_token_json_response unless request.cookies[CSRF_TOKEN_COOKIE_NAME]
+
+      request.body.rewind
+      body = JSON.parse(request.body.read)
+      return invalid_csrf_token_json_response unless request.cookies[CSRF_TOKEN_COOKIE_NAME] == body['csrfToken']
+
+      response = @error_page.send("do_#{opts[:method]}", body)
+      [200, { "Content-Type" => "application/json; charset=utf-8" }, [JSON.dump(response)]]
     end
 
     def no_errors_page
@@ -170,5 +190,13 @@ def invalid_error_json_response
           "and the exception is no longer available in memory.",
       )]]
     end
+
+    def invalid_csrf_token_json_response
+      [200, { "Content-Type" => "text/plain; charset=utf-8" }, [JSON.dump(
+        error: "Invalid CSRF Token",
+        explanation: "The browser session might have been cleared, " +
+          "or something went wrong.",
+      )]]
+    end
   end
 end
diff --git a/lib/better_errors/templates/main.erb b/lib/better_errors/templates/main.erb
index a945d99c..1ab68559 100644
--- a/lib/better_errors/templates/main.erb
+++ b/lib/better_errors/templates/main.erb
@@ -800,6 +800,7 @@
 (function() {
 
     var OID = "<%= id %>";
+    var csrfToken = "<%= csrf_token %>";
 
     var previousFrame = null;
     var previousFrameInfo = null;
@@ -810,6 +811,7 @@
         var req = new XMLHttpRequest();
         req.open("POST", "//" + window.location.host + <%== uri_prefix.gsub("<", "&lt;").inspect %> + "/__better_errors/" + OID + "/" + method, true);
         req.setRequestHeader("Content-Type", "application/json");
+        opts.csrfToken = csrfToken;
         req.send(JSON.stringify(opts));
         req.onreadystatechange = function() {
             if(req.readyState == 4) {