From 9d120ad66b120b7d458b14893574e6e403672957 Mon Sep 17 00:00:00 2001
From: Samuel Georges <sam@daftspunk.com>
Date: Sat, 2 Nov 2019 18:57:32 +1100
Subject: [PATCH] Add header and cookie support to ResponseMaker

---
 modules/backend/controllers/Auth.php    | 14 +++---
 modules/system/traits/ResponseMaker.php | 60 +++++++++++++++++++++++--
 2 files changed, 61 insertions(+), 13 deletions(-)

diff --git a/modules/backend/controllers/Auth.php b/modules/backend/controllers/Auth.php
index 1c8c4403f..7bd8d59ac 100644
--- a/modules/backend/controllers/Auth.php
+++ b/modules/backend/controllers/Auth.php
@@ -35,11 +35,6 @@ class Auth extends Controller
     {
         parent::__construct();
 
-        // $this->middleware(function ($request, $response) {
-        //     // Clear Cache and any previous data to fix Invalid security token issue, see github: #3707
-        //     $response->headers->set('Cache-Control', 'no-cache, no-store, must-revalidate');
-        // })->only('signin');
-
         $this->layout = 'auth';
     }
 
@@ -58,6 +53,9 @@ class Auth extends Controller
     {
         $this->bodyClass = 'signin';
 
+        // Clear Cache and any previous data to fix invalid security token issue
+        $this->setResponseHeader('Cache-Control', 'no-cache, no-store, must-revalidate');
+
         try {
             if (post('postback')) {
                 return $this->signin_onSubmit();
@@ -122,14 +120,12 @@ class Auth extends Controller
             BackendAuth::logout();
         }
 
-        $redirect = Backend::redirect('backend');
-
         // Add HTTP Header 'Clear Site Data' to purge all sensitive data upon signout
         if (Request::secure()) {
-            $redirect->header('Clear-Site-Data', 'cache, cookies, storage, executionContexts');
+            $this->setResponseHeader('Clear-Site-Data', 'cache, cookies, storage, executionContexts');
         }
 
-        return $redirect;
+        return Backend::redirect('backend');
     }
 
     /**
diff --git a/modules/system/traits/ResponseMaker.php b/modules/system/traits/ResponseMaker.php
index 9aa4bed4d..a4c9ff1de 100644
--- a/modules/system/traits/ResponseMaker.php
+++ b/modules/system/traits/ResponseMaker.php
@@ -1,6 +1,8 @@
 <?php namespace System\Traits;
 
 use Response;
+use Symfony\Component\HttpFoundation\HeaderBag;
+use Symfony\Component\HttpFoundation\Response as BaseResponse;
 
 /**
  * Response Maker Trait
@@ -21,10 +23,15 @@ trait ResponseMaker
      */
     protected $responseOverride = null;
 
+    /**
+     * @var Symfony\Component\HttpFoundation\HeaderBag
+     */
+    protected $responseHeaderBag = null;
+
     /**
      * Sets the status code for the current web response.
      * @param int $code Status code
-     * @return self
+     * @return $this
      */
     public function setStatusCode($code)
     {
@@ -45,6 +52,7 @@ trait ResponseMaker
      * Sets the response for the current page request cycle, this value takes priority
      * over the standard response prepared by the controller.
      * @param mixed $response Response object or string
+     * @return $this
      */
     public function setResponse($response)
     {
@@ -52,6 +60,46 @@ trait ResponseMaker
         return $this;
     }
 
+    /**
+     * Set a header on the Response.
+     *
+     * @param  string  $key
+     * @param  array|string  $values
+     * @param  bool    $replace
+     * @return $this
+     */
+    public function setResponseHeader($key, $values, $replace = true)
+    {
+        if ($this->responseHeaderBag === null) {
+            $this->responseHeaderBag = new HeaderBag;
+        }
+
+        $this->responseHeaderBag->set($key, $values, $replace);
+
+        return $this;
+    }
+
+    /**
+     * Add a cookie to the response.
+     *
+     * @param  \Symfony\Component\HttpFoundation\Cookie|mixed  $cookie
+     * @return $this
+     */
+    public function setResponseCookie($cookie)
+    {
+        if ($this->responseHeaderBag === null) {
+            $this->responseHeaderBag = new HeaderBag;
+        }
+
+        if (is_string($cookie) && function_exists('cookie')) {
+            $cookie = call_user_func_array('cookie', func_get_args());
+        }
+
+        $this->responseHeaderBag->setCookie($cookie);
+
+        return $this;
+    }
+
     /**
      * Prepares a response that considers overrides and custom responses.
      * @param mixed $contents
@@ -63,10 +111,14 @@ trait ResponseMaker
             $contents = $this->responseOverride;
         }
 
-        if (!is_string($contents)) {
-            return $contents;
+        if (is_string($contents)) {
+            $contents = Response::make($contents, $this->statusCode);
         }
 
-        return Response::make($contents, $this->statusCode);
+        if ($contents instanceof BaseResponse && $this->responseHeaderBag !== null) {
+            $contents = $contents->withHeaders($this->responseHeaderBag);
+        }
+
+        return $contents;
     }
 }