Fix CSRF tokens not being set on diactoros responses

Update typehints and implementation to be compatible with diactoros
response. Thankfully Cookie now has a `toHeaderValue()` method which
enables CSRF tokens to be set on diactoros response objects.

src/Http/Middleware/CsrfProtectionMiddleware.php

@@ -26,7 +26,6 @@ use Psr\Http\Message\ResponseInterface;
 use Psr\Http\Message\ServerRequestInterface;
 use Psr\Http\Server\MiddlewareInterface;
 use Psr\Http\Server\RequestHandlerInterface;
-use Zend\Diactoros\Response\RedirectResponse;
 
 /**
  * Provides CSRF protection & validation.
@@ -121,9 +120,6 @@ class CsrfProtectionMiddleware implements MiddlewareInterface
             $request = $request->withAttribute('csrfToken', $token);
             /** @var mixed $response */
             $response = $handler->handle($request);
-            if ($response instanceof RedirectResponse) {
-                return $response;
-            }
 
             return $this->_addTokenCookie($token, $request, $response);
         }
@@ -184,10 +180,10 @@ class CsrfProtectionMiddleware implements MiddlewareInterface
      *
      * @param string $token The token to add.
      * @param \Psr\Http\Message\ServerRequestInterface $request The request to validate against.
-     * @param \Cake\Http\Response $response The response.
+     * @param \Psr\Http\Message\ResponseInterface $response The response.
      * @return \Cake\Http\Response $response Modified response.
      */
-    protected function _addTokenCookie(string $token, ServerRequestInterface $request, Response $response): Response
+    protected function _addTokenCookie(string $token, ServerRequestInterface $request, ResponseInterface $response): ResponseInterface
     {
         $cookie = Cookie::create(
             $this->_config['cookieName'],
@@ -199,8 +195,11 @@ class CsrfProtectionMiddleware implements MiddlewareInterface
                 'httponly' => $this->_config['httpOnly'],
             ]
         );
+        if ($response instanceof Response) {
+            return $response->withCookie($cookie);
+        }
 
-        return $response->withCookie($cookie);
+        return $response->withAddedHeader('Set-Cookie', $cookie->toHeaderValue());
     }
 
     /**

tests/TestCase/Http/Middleware/CsrfProtectionMiddlewareTest.php

@@ -23,6 +23,7 @@ use Cake\TestSuite\TestCase;
 use Psr\Http\Message\ServerRequestInterface;
 use TestApp\Http\TestRequestHandler;
 use Zend\Diactoros\Response\RedirectResponse;
+use Zend\Diactoros\Response as DiactorosResponse;
 
 /**
  * Test for CsrfProtection
@@ -121,28 +122,41 @@ class CsrfProtectionMiddlewareTest extends TestCase
     }
 
     /**
-     * Test that the CSRF tokens are not set for redirect responses
+     * Test that the CSRF tokens are set for redirect responses
      *
      * @return void
      */
-    public function testRedirectResponseCookiesNotSet()
+    public function testRedirectResponseCookies()
     {
         $request = new ServerRequest([
             'environment' => ['REQUEST_METHOD' => 'GET'],
         ]);
-        $expectedResponse = new RedirectResponse('/');
-        $handler = new TestRequestHandler(function ($request) use ($expectedResponse) {
+        $handler = new TestRequestHandler(function () {
+            return new RedirectResponse('/');
+        });
 
-            return $expectedResponse;
+        $middleware = new CsrfProtectionMiddleware();
+        $response = $middleware->process($request, $handler);
+        $this->assertStringContainsString('csrfToken=', $response->getHeaderLine('Set-Cookie'));
+    }
+
+    /**
+     * Test that the CSRF tokens are set for diactoros responses
+     *
+     * @return void
+     */
+    public function testDiactorosResponseCookies()
+    {
+        $request = new ServerRequest([
+            'environment' => ['REQUEST_METHOD' => 'GET'],
+        ]);
+        $handler = new TestRequestHandler(function () {
+            return new DiactorosResponse();
         });
 
-        $middleware = $this->getMockBuilder(CsrfProtectionMiddleware::class)
-            ->onlyMethods(['_addTokenCookie'])
-            ->getMock();
-        $middleware->expects($this->never())
-            ->method('_addTokenCookie');
+        $middleware = new CsrfProtectionMiddleware();
         $response = $middleware->process($request, $handler);
-        $this->assertSame($expectedResponse, $response);
+        $this->assertStringContainsString('csrfToken=', $response->getHeaderLine('Set-Cookie'));
     }
 
     /**