From 96a9cf0d2dc95c76f71b5842f057971c7df359e3 Mon Sep 17 00:00:00 2001
From: Josh Cummings <3627351+jzheaux@users.noreply.github.com>
Date: Thu, 5 Dec 2024 09:52:06 -0700
Subject: [PATCH] Restore Previous Behavior for Servlet 5

Closes gh-16173
---
 .../web/csrf/CookieCsrfTokenRepository.java   | 17 +++++--
 .../csrf/CookieCsrfTokenRepositoryTests.java  | 44 ++++++++++++++++++-
 2 files changed, 57 insertions(+), 4 deletions(-)

diff --git a/web/src/main/java/org/springframework/security/web/csrf/CookieCsrfTokenRepository.java b/web/src/main/java/org/springframework/security/web/csrf/CookieCsrfTokenRepository.java
index 7463f6aa15..546ccd94d0 100644
--- a/web/src/main/java/org/springframework/security/web/csrf/CookieCsrfTokenRepository.java
+++ b/web/src/main/java/org/springframework/security/web/csrf/CookieCsrfTokenRepository.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2023 the original author or authors.
+ * Copyright 2012-2024 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -23,6 +23,7 @@ import jakarta.servlet.http.Cookie;
 import jakarta.servlet.http.HttpServletRequest;
 import jakarta.servlet.http.HttpServletResponse;
 
+import org.springframework.http.HttpHeaders;
 import org.springframework.http.ResponseCookie;
 import org.springframework.util.Assert;
 import org.springframework.util.StringUtils;
@@ -97,8 +98,18 @@ public final class CookieCsrfTokenRepository implements CsrfTokenRepository {
 
 		this.cookieCustomizer.accept(cookieBuilder);
 
-		Cookie cookie = mapToCookie(cookieBuilder.build());
-		response.addCookie(cookie);
+		ResponseCookie responseCookie = cookieBuilder.build();
+		if (!StringUtils.hasLength(responseCookie.getSameSite())) {
+			Cookie cookie = mapToCookie(responseCookie);
+			response.addCookie(cookie);
+		}
+		else if (request.getServletContext().getMajorVersion() > 5) {
+			Cookie cookie = mapToCookie(responseCookie);
+			response.addCookie(cookie);
+		}
+		else {
+			response.addHeader(HttpHeaders.SET_COOKIE, responseCookie.toString());
+		}
 
 		// Set request attribute to signal that response has blank cookie value,
 		// which allows loadToken to return null when token has been removed
diff --git a/web/src/test/java/org/springframework/security/web/csrf/CookieCsrfTokenRepositoryTests.java b/web/src/test/java/org/springframework/security/web/csrf/CookieCsrfTokenRepositoryTests.java
index 05f73a0d89..bc2088e53c 100644
--- a/web/src/test/java/org/springframework/security/web/csrf/CookieCsrfTokenRepositoryTests.java
+++ b/web/src/test/java/org/springframework/security/web/csrf/CookieCsrfTokenRepositoryTests.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2023 the original author or authors.
+ * Copyright 2002-2024 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,16 +17,20 @@
 package org.springframework.security.web.csrf;
 
 import jakarta.servlet.http.Cookie;
+import jakarta.servlet.http.HttpServletResponse;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.http.HttpHeaders;
 import org.springframework.mock.web.MockHttpServletRequest;
 import org.springframework.mock.web.MockHttpServletResponse;
+import org.springframework.mock.web.MockServletContext;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.springframework.security.web.csrf.CsrfTokenAssert.assertThatCsrfToken;
@@ -447,6 +451,44 @@ class CookieCsrfTokenRepositoryTests {
 		assertThat(tokenCookie.isHttpOnly()).isEqualTo(Boolean.FALSE);
 	}
 
+	// gh-16173
+	@Test
+	void saveTokenWhenSameSiteAndServletVersion5ThenUsesAddHeader() {
+		HttpServletResponse response = mock(HttpServletResponse.class);
+		((MockServletContext) this.request.getServletContext()).setMajorVersion(5);
+		this.repository.setCookieCustomizer((builder) -> builder.sameSite("Strict"));
+		CsrfToken token = this.repository.generateToken(this.request);
+		this.repository.saveToken(token, this.request, response);
+		verify(response, never()).addCookie(any(Cookie.class));
+		verify(response).addHeader(any(), any());
+	}
+
+	// gh-16173
+	@Test
+	void saveTokenWhenSameSiteAndServletVersion6OrHigherThenUsesAddCookie() {
+		HttpServletResponse response = mock(HttpServletResponse.class);
+		this.repository.setCookieCustomizer((builder) -> builder.sameSite("Strict"));
+		CsrfToken token = this.repository.generateToken(this.request);
+		this.repository.saveToken(token, this.request, response);
+		verify(response).addCookie(any(Cookie.class));
+		verify(response, never()).addHeader(any(), any());
+	}
+
+	// gh-16173
+	@Test
+	void saveTokenWhenNoSameSiteThenUsesAddCookie() {
+		HttpServletResponse response = mock(HttpServletResponse.class);
+		CsrfToken token = this.repository.generateToken(this.request);
+		this.repository.saveToken(token, this.request, response);
+		verify(response).addCookie(any(Cookie.class));
+		verify(response, never()).addHeader(any(), any());
+		((MockServletContext) this.request.getServletContext()).setMajorVersion(5);
+		response = mock(HttpServletResponse.class);
+		this.repository.saveToken(token, this.request, response);
+		verify(response).addCookie(any(Cookie.class));
+		verify(response, never()).addHeader(any(), any());
+	}
+
 	@Test
 	void setCookieNameNullIllegalArgumentException() {
 		assertThatIllegalArgumentException().isThrownBy(() -> this.repository.setCookieName(null));