You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@nifi.apache.org by th...@apache.org on 2022/12/05 19:10:22 UTC

[nifi] branch main updated: NIFI-10899 Added SameSite Policy to Application Cookies

This is an automated email from the ASF dual-hosted git repository.

thenatog pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/nifi.git


The following commit(s) were added to refs/heads/main by this push:
     new 45a31c7286 NIFI-10899 Added SameSite Policy to Application Cookies
45a31c7286 is described below

commit 45a31c7286b89a12487054078c9f1adea18b0fcb
Author: exceptionfactory <ex...@apache.org>
AuthorDate: Tue Nov 29 14:04:10 2022 -0600

    NIFI-10899 Added SameSite Policy to Application Cookies
    
    - Added __Secure prefix to Application Cookie Names
    
    Signed-off-by: Nathan Gough <th...@gmail.com>
    
    This closes #6735.
---
 .../web/security/cookie/ApplicationCookieName.java | 24 ++++--
 ...licationCookieName.java => SameSitePolicy.java} | 26 +++---
 .../cookie/StandardApplicationCookieService.java   | 17 +++-
 .../csrf/StandardCookieCsrfTokenRepository.java    | 58 +++++++------
 .../StandardApplicationCookieServiceTest.java      | 33 ++++----
 .../StandardCookieCsrfTokenRepositoryTest.java     | 94 +++++++++-------------
 6 files changed, 133 insertions(+), 119 deletions(-)

diff --git a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/cookie/ApplicationCookieName.java b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/cookie/ApplicationCookieName.java
index dbbea5c9bb..13476d6d2c 100644
--- a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/cookie/ApplicationCookieName.java
+++ b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/cookie/ApplicationCookieName.java
@@ -22,21 +22,35 @@ import org.apache.nifi.web.security.http.SecurityCookieName;
  * Application Cookie Names
  */
 public enum ApplicationCookieName {
-    AUTHORIZATION_BEARER(SecurityCookieName.AUTHORIZATION_BEARER.getName()),
+    /** Authorization Bearer contains signed JSON Web Token and requires Strict Same Site handling */
+    AUTHORIZATION_BEARER(SecurityCookieName.AUTHORIZATION_BEARER.getName(), SameSitePolicy.STRICT),
 
-    LOGOUT_REQUEST_IDENTIFIER("nifi-logout-request-identifier"),
+    /** Cross-Site Request Forgery mitigation token requires Strict Same Site handling */
+    REQUEST_TOKEN(SecurityCookieName.REQUEST_TOKEN.getName(), SameSitePolicy.STRICT),
 
-    OIDC_REQUEST_IDENTIFIER("nifi-oidc-request-identifier"),
+    /** Logout Requests can interact with external identity providers requiring no Same Site restrictions */
+    LOGOUT_REQUEST_IDENTIFIER("__Secure-Logout-Request-Identifier", SameSitePolicy.NONE),
 
-    SAML_REQUEST_IDENTIFIER("nifi-saml-request-identifier");
+    /** OpenID Connect Requests use external identity providers requiring no Same Site restrictions */
+    OIDC_REQUEST_IDENTIFIER("__Secure-OIDC-Request-Identifier", SameSitePolicy.NONE),
+
+    /** SAML Requests use external identity providers requiring no Same Site restrictions */
+    SAML_REQUEST_IDENTIFIER("__Secure-SAML-Request-Identifier", SameSitePolicy.NONE);
 
     private final String cookieName;
 
-    ApplicationCookieName(final String cookieName) {
+    private final SameSitePolicy sameSitePolicy;
+
+    ApplicationCookieName(final String cookieName, final SameSitePolicy sameSitePolicy) {
         this.cookieName = cookieName;
+        this.sameSitePolicy = sameSitePolicy;
     }
 
     public String getCookieName() {
         return cookieName;
     }
+
+    public SameSitePolicy getSameSitePolicy() {
+        return sameSitePolicy;
+    }
 }
diff --git a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/cookie/ApplicationCookieName.java b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/cookie/SameSitePolicy.java
similarity index 59%
copy from nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/cookie/ApplicationCookieName.java
copy to nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/cookie/SameSitePolicy.java
index dbbea5c9bb..fc5a79f540 100644
--- a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/cookie/ApplicationCookieName.java
+++ b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/cookie/SameSitePolicy.java
@@ -16,27 +16,23 @@
  */
 package org.apache.nifi.web.security.cookie;
 
-import org.apache.nifi.web.security.http.SecurityCookieName;
-
 /**
- * Application Cookie Names
+ * Cookie SameSite attribute policy
  */
-public enum ApplicationCookieName {
-    AUTHORIZATION_BEARER(SecurityCookieName.AUTHORIZATION_BEARER.getName()),
-
-    LOGOUT_REQUEST_IDENTIFIER("nifi-logout-request-identifier"),
-
-    OIDC_REQUEST_IDENTIFIER("nifi-oidc-request-identifier"),
+public enum SameSitePolicy {
+    /** Instructs browsers to limit sending cookies to first-party-initiated requests */
+    STRICT("Strict"),
 
-    SAML_REQUEST_IDENTIFIER("nifi-saml-request-identifier");
+    /** Instructs browsers to send cookies for both first-party and cross-site requests */
+    NONE("None");
 
-    private final String cookieName;
+    private final String policy;
 
-    ApplicationCookieName(final String cookieName) {
-        this.cookieName = cookieName;
+    SameSitePolicy(final String policy) {
+        this.policy = policy;
     }
 
-    public String getCookieName() {
-        return cookieName;
+    public String getPolicy() {
+        return policy;
     }
 }
diff --git a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/cookie/StandardApplicationCookieService.java b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/cookie/StandardApplicationCookieService.java
index 01862fb403..a1b2cb20e1 100644
--- a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/cookie/StandardApplicationCookieService.java
+++ b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/cookie/StandardApplicationCookieService.java
@@ -43,8 +43,6 @@ public class StandardApplicationCookieService implements ApplicationCookieServic
 
     private static final String DEFAULT_PATH = "/";
 
-    private static final String SAME_SITE_STRICT = "Strict";
-
     private static final boolean SECURE_ENABLED = true;
 
     private static final boolean HTTP_ONLY_ENABLED = true;
@@ -77,7 +75,6 @@ public class StandardApplicationCookieService implements ApplicationCookieServic
     @Override
     public void addSessionCookie(final URI resourceUri, final HttpServletResponse response, final ApplicationCookieName applicationCookieName, final String value) {
         final ResponseCookie.ResponseCookieBuilder responseCookieBuilder = getCookieBuilder(resourceUri, applicationCookieName, value, MAX_AGE_SESSION);
-        responseCookieBuilder.sameSite(SAME_SITE_STRICT);
         setResponseCookie(response, responseCookieBuilder.build());
         logger.debug("Added Session Cookie [{}] URI [{}]", applicationCookieName.getCookieName(), resourceUri);
     }
@@ -110,15 +107,27 @@ public class StandardApplicationCookieService implements ApplicationCookieServic
         logger.debug("Removed Cookie [{}] URI [{}]", applicationCookieName.getCookieName(), resourceUri);
     }
 
-    private ResponseCookie.ResponseCookieBuilder getCookieBuilder(final URI resourceUri,
+    /**
+     * Get Response Cookie Builder with standard properties
+     *
+     * @param resourceUri Resource URI containing path and domain
+     * @param applicationCookieName Application Cookie Name to be used
+     * @param value Cookie value
+     * @param maxAge Max Age
+     * @return Response Cookie Builder
+     */
+    protected ResponseCookie.ResponseCookieBuilder getCookieBuilder(final URI resourceUri,
                                              final ApplicationCookieName applicationCookieName,
                                              final String value,
                                              final Duration maxAge) {
         Objects.requireNonNull(resourceUri, "Resource URI required");
         Objects.requireNonNull(applicationCookieName, "Response Cookie Name required");
+
+        final SameSitePolicy sameSitePolicy = applicationCookieName.getSameSitePolicy();
         return ResponseCookie.from(applicationCookieName.getCookieName(), value)
                 .path(getCookiePath(resourceUri))
                 .domain(resourceUri.getHost())
+                .sameSite(sameSitePolicy.getPolicy())
                 .secure(SECURE_ENABLED)
                 .httpOnly(HTTP_ONLY_ENABLED)
                 .maxAge(maxAge);
diff --git a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/csrf/StandardCookieCsrfTokenRepository.java b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/csrf/StandardCookieCsrfTokenRepository.java
index 9e5d64b9aa..5dbd3634e1 100644
--- a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/csrf/StandardCookieCsrfTokenRepository.java
+++ b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/csrf/StandardCookieCsrfTokenRepository.java
@@ -16,9 +16,13 @@
  */
 package org.apache.nifi.web.security.csrf;
 
+import org.apache.nifi.web.security.cookie.ApplicationCookieName;
+import org.apache.nifi.web.security.cookie.ApplicationCookieService;
+import org.apache.nifi.web.security.cookie.StandardApplicationCookieService;
 import org.apache.nifi.web.security.http.SecurityCookieName;
 import org.apache.nifi.web.security.http.SecurityHeader;
 import org.apache.nifi.web.util.RequestUriBuilder;
+import org.springframework.http.ResponseCookie;
 import org.springframework.security.web.csrf.CsrfToken;
 import org.springframework.security.web.csrf.CsrfTokenRepository;
 import org.springframework.security.web.csrf.DefaultCsrfToken;
@@ -29,6 +33,7 @@ import javax.servlet.http.Cookie;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 import java.net.URI;
+import java.time.Duration;
 import java.util.UUID;
 
 /**
@@ -37,15 +42,7 @@ import java.util.UUID;
 public class StandardCookieCsrfTokenRepository implements CsrfTokenRepository {
     private static final String REQUEST_PARAMETER = "requestToken";
 
-    private static final String ROOT_PATH = "/";
-
-    private static final String EMPTY = "";
-
-    private static final boolean SECURE_ENABLED = true;
-
-    private static final int MAX_AGE_EXPIRED = 0;
-
-    private static final int MAX_AGE_SESSION = -1;
+    private static final ApplicationCookieService applicationCookieService = new CsrfApplicationCookieService();
 
     /**
      * Generate CSRF Token or return current Token when present in HTTP Servlet Request Cookie header
@@ -71,16 +68,13 @@ public class StandardCookieCsrfTokenRepository implements CsrfTokenRepository {
      */
     @Override
     public void saveToken(final CsrfToken csrfToken, final HttpServletRequest httpServletRequest, final HttpServletResponse httpServletResponse) {
-        final String token = csrfToken == null ? EMPTY : csrfToken.getToken();
-        final int maxAge = csrfToken == null ? MAX_AGE_EXPIRED : MAX_AGE_SESSION;
-
-        final Cookie cookie = new Cookie(SecurityCookieName.REQUEST_TOKEN.getName(), token);
-        cookie.setSecure(SECURE_ENABLED);
-        cookie.setMaxAge(maxAge);
-
-        final String cookiePath = getCookiePath(httpServletRequest);
-        cookie.setPath(cookiePath);
-        httpServletResponse.addCookie(cookie);
+        final URI uri = RequestUriBuilder.fromHttpServletRequest(httpServletRequest).build();
+        if (csrfToken == null) {
+            applicationCookieService.removeCookie(uri, httpServletResponse, ApplicationCookieName.REQUEST_TOKEN);
+        } else {
+            final String token = csrfToken.getToken();
+            applicationCookieService.addSessionCookie(uri, httpServletResponse, ApplicationCookieName.REQUEST_TOKEN, token);
+        }
     }
 
     /**
@@ -104,10 +98,26 @@ public class StandardCookieCsrfTokenRepository implements CsrfTokenRepository {
         return UUID.randomUUID().toString();
     }
 
-    private String getCookiePath(final HttpServletRequest httpServletRequest) {
-        final RequestUriBuilder requestUriBuilder = RequestUriBuilder.fromHttpServletRequest(httpServletRequest);
-        requestUriBuilder.path(ROOT_PATH);
-        final URI uri = requestUriBuilder.build();
-        return uri.getPath();
+    private static class CsrfApplicationCookieService extends StandardApplicationCookieService {
+        private static final boolean HTTP_ONLY_DISABLED = false;
+
+        /**
+         * Get Response Cookie Builder with HttpOnly disabled allowing JavaScript to read value for subsequent requests
+         *
+         * @param resourceUri Resource URI containing path and domain
+         * @param applicationCookieName Application Cookie Name to be used
+         * @param value Cookie value
+         * @param maxAge Max Age
+         * @return Response Cookie Builder
+         */
+        @Override
+        protected ResponseCookie.ResponseCookieBuilder getCookieBuilder(final URI resourceUri,
+                                                                        final ApplicationCookieName applicationCookieName,
+                                                                        final String value,
+                                                                        final Duration maxAge) {
+            final ResponseCookie.ResponseCookieBuilder builder = super.getCookieBuilder(resourceUri, applicationCookieName, value, maxAge);
+            builder.httpOnly(HTTP_ONLY_DISABLED);
+            return builder;
+        }
     }
 }
diff --git a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/test/java/org/apache/nifi/web/security/cookie/StandardApplicationCookieServiceTest.java b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/test/java/org/apache/nifi/web/security/cookie/StandardApplicationCookieServiceTest.java
index c36af64f33..7b6d8b78db 100644
--- a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/test/java/org/apache/nifi/web/security/cookie/StandardApplicationCookieServiceTest.java
+++ b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/test/java/org/apache/nifi/web/security/cookie/StandardApplicationCookieServiceTest.java
@@ -17,13 +17,13 @@
 package org.apache.nifi.web.security.cookie;
 
 import org.apache.nifi.util.StringUtils;
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Captor;
 import org.mockito.Mock;
-import org.mockito.junit.MockitoJUnitRunner;
+import org.mockito.junit.jupiter.MockitoExtension;
 import org.springframework.mock.web.MockCookie;
 
 import javax.servlet.http.Cookie;
@@ -34,14 +34,14 @@ import java.net.URI;
 import java.util.Optional;
 import java.util.UUID;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
-@RunWith(MockitoJUnitRunner.class)
+@ExtendWith(MockitoExtension.class)
 public class StandardApplicationCookieServiceTest {
     private static final String DOMAIN = "localhost.localdomain";
 
@@ -59,7 +59,7 @@ public class StandardApplicationCookieServiceTest {
 
     private static final int REMOVE_MAX_AGE = 0;
 
-    private static final String SAME_SITE_STRICT = "SameSite=Strict";
+    private static final String SAME_SITE = "SameSite";
 
     private static final String COOKIE_VALUE = UUID.randomUUID().toString();
 
@@ -80,7 +80,7 @@ public class StandardApplicationCookieServiceTest {
     @Captor
     private ArgumentCaptor<String> cookieArgumentCaptor;
 
-    @Before
+    @BeforeEach
     public void setService() {
         service = new StandardApplicationCookieService();
         resourceUri = URI.create(RESOURCE_URI);
@@ -113,7 +113,6 @@ public class StandardApplicationCookieServiceTest {
 
         final String setCookieHeader = cookieArgumentCaptor.getValue();
         assertAddCookieMatches(setCookieHeader, ROOT_PATH, SESSION_MAX_AGE);
-        assertTrue("SameSite not found", setCookieHeader.endsWith(SAME_SITE_STRICT));
     }
 
     @Test
@@ -124,7 +123,6 @@ public class StandardApplicationCookieServiceTest {
 
         final String setCookieHeader = cookieArgumentCaptor.getValue();
         assertAddCookieMatches(setCookieHeader, CONTEXT_PATH, SESSION_MAX_AGE);
-        assertTrue("SameSite not found", setCookieHeader.endsWith(SAME_SITE_STRICT));
     }
 
     @Test
@@ -175,10 +173,11 @@ public class StandardApplicationCookieServiceTest {
     }
 
     private void assertCookieMatches(final String setCookieHeader, final Cookie cookie, final String path) {
-        assertEquals("Cookie Name not matched", COOKIE_NAME.getCookieName(), cookie.getName());
-        assertEquals("Path not matched", path, cookie.getPath());
-        assertEquals("Domain not matched", DOMAIN, cookie.getDomain());
-        assertTrue("HTTP Only not matched", cookie.isHttpOnly());
-        assertTrue("Secure not matched", cookie.getSecure());
+        assertEquals(COOKIE_NAME.getCookieName(), cookie.getName(), "Cookie Name not matched");
+        assertEquals(path, cookie.getPath(), "Path not matched");
+        assertEquals(DOMAIN, cookie.getDomain(), "Domain not matched");
+        assertTrue(cookie.isHttpOnly(), "HTTP Only not matched");
+        assertTrue(cookie.getSecure(), "Secure not matched");
+        assertTrue(setCookieHeader.contains(SAME_SITE), "SameSite not found");
     }
 }
diff --git a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/test/java/org/apache/nifi/web/security/csrf/StandardCookieCsrfTokenRepositoryTest.java b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/test/java/org/apache/nifi/web/security/csrf/StandardCookieCsrfTokenRepositoryTest.java
index 9616d95bcc..74a5f7e198 100644
--- a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/test/java/org/apache/nifi/web/security/csrf/StandardCookieCsrfTokenRepositoryTest.java
+++ b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/test/java/org/apache/nifi/web/security/csrf/StandardCookieCsrfTokenRepositoryTest.java
@@ -20,30 +20,20 @@ import org.apache.nifi.web.security.http.SecurityCookieName;
 import org.apache.nifi.web.util.WebUtils;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
-import org.junit.jupiter.api.extension.ExtendWith;
-import org.mockito.ArgumentCaptor;
-import org.mockito.Captor;
-import org.mockito.Mock;
-import org.mockito.junit.jupiter.MockitoExtension;
+import org.springframework.mock.web.MockHttpServletRequest;
+import org.springframework.mock.web.MockHttpServletResponse;
+import org.springframework.mock.web.MockServletContext;
 import org.springframework.security.web.csrf.CsrfToken;
 
-import javax.servlet.ServletContext;
 import javax.servlet.http.Cookie;
-import javax.servlet.http.HttpServletRequest;
-import javax.servlet.http.HttpServletResponse;
 
 import java.util.UUID;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
-import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
-import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
 
-@ExtendWith(MockitoExtension.class)
 public class StandardCookieCsrfTokenRepositoryTest {
     private static final String ALLOWED_CONTEXT_PATHS_PARAMETER = "allowedContextPaths";
 
@@ -55,8 +45,6 @@ public class StandardCookieCsrfTokenRepositoryTest {
 
     private static final String CONTEXT_PATH = "/context-path";
 
-    private static final String COOKIE_CONTEXT_PATH = CONTEXT_PATH + ROOT_PATH;
-
     private static final String HTTPS = "https";
 
     private static final String HOST = "localhost";
@@ -65,23 +53,21 @@ public class StandardCookieCsrfTokenRepositoryTest {
 
     private static final String EMPTY = "";
 
-    @Mock
-    private HttpServletRequest request;
+    private static final String SET_COOKIE_HEADER = "Set-Cookie";
 
-    @Mock
-    private HttpServletResponse response;
+    private static final String SAME_SITE = "SameSite";
 
-    @Mock
-    private ServletContext servletContext;
+    private MockHttpServletRequest request;
 
-    @Captor
-    private ArgumentCaptor<Cookie> cookieArgumentCaptor;
+    private MockHttpServletResponse response;
 
     private StandardCookieCsrfTokenRepository repository;
 
     @BeforeEach
     public void setRepository() {
         this.repository = new StandardCookieCsrfTokenRepository();
+        this.request = new MockHttpServletRequest();
+        this.response = new MockHttpServletResponse();
     }
 
     @Test
@@ -95,7 +81,7 @@ public class StandardCookieCsrfTokenRepositoryTest {
     public void testGenerateTokenCookieFound() {
         final String token = UUID.randomUUID().toString();
         final Cookie cookie = new Cookie(SecurityCookieName.REQUEST_TOKEN.getName(), token);
-        when(request.getCookies()).thenReturn(new Cookie[]{cookie});
+        request.setCookies(cookie);
 
         final CsrfToken csrfToken = repository.generateToken(request);
         assertNotNull(csrfToken);
@@ -106,7 +92,7 @@ public class StandardCookieCsrfTokenRepositoryTest {
     public void testLoadToken() {
         final String token = UUID.randomUUID().toString();
         final Cookie cookie = new Cookie(SecurityCookieName.REQUEST_TOKEN.getName(), token);
-        when(request.getCookies()).thenReturn(new Cookie[]{cookie});
+        request.setCookies(cookie);
 
         final CsrfToken csrfToken = repository.loadToken(request);
         assertNotNull(csrfToken);
@@ -115,31 +101,22 @@ public class StandardCookieCsrfTokenRepositoryTest {
 
     @Test
     public void testSaveToken() {
-        when(request.getServletContext()).thenReturn(servletContext);
-
         final CsrfToken csrfToken = repository.generateToken(request);
         repository.saveToken(csrfToken, request, response);
 
-        verify(response).addCookie(cookieArgumentCaptor.capture());
-        final Cookie cookie = cookieArgumentCaptor.getValue();
-        assertCookieEquals(csrfToken, cookie);
+        final Cookie cookie = assertCookieFound();
+        final String setCookieHeader = response.getHeader(SET_COOKIE_HEADER);
+        assertCookieEquals(csrfToken.getToken(), MAX_AGE_SESSION, cookie, setCookieHeader);
         assertEquals(ROOT_PATH, cookie.getPath());
     }
 
     @Test
     public void testSaveTokenNullCsrfToken() {
-        when(request.getServletContext()).thenReturn(servletContext);
-
         repository.saveToken(null, request, response);
 
-        verify(response).addCookie(cookieArgumentCaptor.capture());
-        final Cookie cookie = cookieArgumentCaptor.getValue();
-        assertEquals(ROOT_PATH, cookie.getPath());
-        assertEquals(EMPTY, cookie.getValue());
-        assertEquals(MAX_AGE_EXPIRED, cookie.getMaxAge());
-        assertTrue(cookie.getSecure());
-        assertFalse(cookie.isHttpOnly());
-        assertNull(cookie.getDomain());
+        final Cookie cookie = assertCookieFound();
+        final String setCookieHeader = response.getHeader(SET_COOKIE_HEADER);
+        assertCookieEquals(EMPTY, MAX_AGE_EXPIRED, cookie, setCookieHeader);
     }
 
     @Test
@@ -147,27 +124,36 @@ public class StandardCookieCsrfTokenRepositoryTest {
         this.repository = new StandardCookieCsrfTokenRepository();
 
         final CsrfToken csrfToken = repository.generateToken(request);
-        when(request.getHeader(eq(WebUtils.PROXY_SCHEME_HTTP_HEADER))).thenReturn(HTTPS);
-        when(request.getHeader(eq(WebUtils.PROXY_HOST_HTTP_HEADER))).thenReturn(HOST);
-        when(request.getHeader(eq(WebUtils.PROXY_PORT_HTTP_HEADER))).thenReturn(PORT);
-        when(request.getHeader(eq(WebUtils.PROXY_CONTEXT_PATH_HTTP_HEADER))).thenReturn(CONTEXT_PATH);
 
-        when(servletContext.getInitParameter(eq(ALLOWED_CONTEXT_PATHS_PARAMETER))).thenReturn(CONTEXT_PATH);
-        when(request.getServletContext()).thenReturn(servletContext);
+        request.addHeader(WebUtils.PROXY_SCHEME_HTTP_HEADER, HTTPS);
+        request.addHeader(WebUtils.PROXY_HOST_HTTP_HEADER, HOST);
+        request.addHeader(WebUtils.PROXY_PORT_HTTP_HEADER, PORT);
+        request.addHeader(WebUtils.PROXY_CONTEXT_PATH_HTTP_HEADER, CONTEXT_PATH);
+
+        final MockServletContext servletContext = (MockServletContext) request.getServletContext();
+        servletContext.setInitParameter(ALLOWED_CONTEXT_PATHS_PARAMETER, CONTEXT_PATH);
 
         repository.saveToken(csrfToken, request, response);
 
-        verify(response).addCookie(cookieArgumentCaptor.capture());
-        final Cookie cookie = cookieArgumentCaptor.getValue();
-        assertCookieEquals(csrfToken, cookie);
-        assertEquals(COOKIE_CONTEXT_PATH, cookie.getPath());
+        final Cookie cookie = assertCookieFound();
+        final String setCookieHeader = response.getHeader(SET_COOKIE_HEADER);
+        assertCookieEquals(csrfToken.getToken(), MAX_AGE_SESSION, cookie, setCookieHeader);
+        assertEquals(CONTEXT_PATH, cookie.getPath());
+    }
+
+    private Cookie assertCookieFound() {
+        final Cookie cookie = response.getCookie(SecurityCookieName.REQUEST_TOKEN.getName());
+        assertNotNull(cookie);
+        return cookie;
     }
 
-    private void assertCookieEquals(final CsrfToken csrfToken, final Cookie cookie) {
-        assertEquals(csrfToken.getToken(), cookie.getValue());
-        assertEquals(MAX_AGE_SESSION, cookie.getMaxAge());
+    private void assertCookieEquals(final String token, final int maxAge, final Cookie cookie, final String setCookieHeader) {
+        assertNotNull(setCookieHeader);
+        assertEquals(token, cookie.getValue());
+        assertEquals(maxAge, cookie.getMaxAge());
         assertTrue(cookie.getSecure());
         assertFalse(cookie.isHttpOnly());
-        assertNull(cookie.getDomain());
+        assertEquals(HOST, cookie.getDomain());
+        assertTrue(setCookieHeader.contains(SAME_SITE), "SameSite not found");
     }
 }