Refactor Jetty SameSiteSupplier cookie support to use a Handler

Update `JettyServletWebServerFactory` so that the `SimeSiteSupplier`
support is handled using a `Handler` rather than a `HttpStream.Wrapper`.

Closes gh-37809

spring-boot-project/spring-boot/src/main/java/org/springframework/boot/web/embedded/jetty/JettyServletWebServerFactory.java

@@ -33,7 +33,7 @@ import java.util.EventListener;
 import java.util.Iterator;
 import java.util.LinkedHashSet;
 import java.util.List;
-import java.util.ListIterator;
+import java.util.Objects;
 import java.util.Set;
 import java.util.Spliterator;
 import java.util.UUID;
@@ -52,8 +52,10 @@ import org.eclipse.jetty.ee10.webapp.AbstractConfiguration;
 import org.eclipse.jetty.ee10.webapp.Configuration;
 import org.eclipse.jetty.ee10.webapp.WebAppContext;
 import org.eclipse.jetty.ee10.webapp.WebInfConfiguration;
+import org.eclipse.jetty.http.CookieCompliance;
 import org.eclipse.jetty.http.HttpCookie;
 import org.eclipse.jetty.http.HttpField;
+import org.eclipse.jetty.http.HttpFields;
 import org.eclipse.jetty.http.HttpFields.Mutable;
 import org.eclipse.jetty.http.HttpHeader;
 import org.eclipse.jetty.http.MimeTypes;
@@ -68,7 +70,6 @@ import org.eclipse.jetty.server.Handler;
 import org.eclipse.jetty.server.HttpConfiguration;
 import org.eclipse.jetty.server.HttpConnectionFactory;
 import org.eclipse.jetty.server.HttpCookieUtils;
-import org.eclipse.jetty.server.HttpStream;
 import org.eclipse.jetty.server.Request;
 import org.eclipse.jetty.server.Response;
 import org.eclipse.jetty.server.Server;
@@ -775,6 +776,8 @@ public class JettyServletWebServerFactory extends AbstractServletWebServerFactor
 	 */
 	private static class SuppliedSameSiteCookieHandlerWrapper extends Handler.Wrapper {
 
+		private static final SetCookieParser setCookieParser = SetCookieParser.newInstance();
+
 		private final List<CookieSameSiteSupplier> suppliers;
 
 		SuppliedSameSiteCookieHandlerWrapper(List<CookieSameSiteSupplier> suppliers) {
@@ -783,62 +786,74 @@ public class JettyServletWebServerFactory extends AbstractServletWebServerFactor
 
 		@Override
 		public boolean handle(Request request, Response response, Callback callback) throws Exception {
-			request.addHttpStreamWrapper((stream) -> new SameSiteCookieHttpStreamWrapper(stream, request));
+			SuppliedSameSiteCookieResponse wrappedResponse = new SuppliedSameSiteCookieResponse(request, response);
-			return super.handle(request, response, callback);
+			return super.handle(request, wrappedResponse, callback);
 		}
 
-		private final class SameSiteCookieHttpStreamWrapper extends HttpStream.Wrapper {
+		private class SuppliedSameSiteCookieResponse extends Response.Wrapper {
 
-			private static final SetCookieParser setCookieParser = SetCookieParser.newInstance();
+			private HttpFields.Mutable wrappedHeaders;
 
-			private final Request request;
+			SuppliedSameSiteCookieResponse(Request request, Response wrapped) {
+				super(request, wrapped);
+				this.wrappedHeaders = new SuppliedSameSiteCookieHeaders(
+						request.getConnectionMetaData().getHttpConfiguration().getResponseCookieCompliance(),
+						wrapped.getHeaders());
+			}
 
-			private SameSiteCookieHttpStreamWrapper(HttpStream wrapped, Request request) {
+			@Override
-				super(wrapped);
+			public Mutable getHeaders() {
-				this.request = request;
+				return this.wrappedHeaders;
+			}
+
+		}
+
+		private class SuppliedSameSiteCookieHeaders extends HttpFields.Mutable.Wrapper {
+
+			private final CookieCompliance compliance;
+
+			SuppliedSameSiteCookieHeaders(CookieCompliance compliance, HttpFields.Mutable fields) {
+				super(fields);
+				this.compliance = compliance;
 			}
 
 			@Override
-			public void prepareResponse(Mutable headers) {
+			public HttpField onAddField(HttpField field) {
-				super.prepareResponse(headers);
+				return (field.getHeader() != HttpHeader.SET_COOKIE) ? field : onAddSetCookieField(field);
-				ListIterator<HttpField> headerFields = headers.listIterator();
-				while (headerFields.hasNext()) {
-					HttpField updatedField = applySameSiteIfNecessary(headerFields.next());
-					if (updatedField != null) {
-						headerFields.set(updatedField);
-					}
-				}
 			}
 
-			private HttpField applySameSiteIfNecessary(HttpField headerField) {
+			private HttpField onAddSetCookieField(HttpField field) {
-				if (headerField.getHeader() != HttpHeader.SET_COOKIE) {
+				HttpCookie cookie = setCookieParser.parse(field.getValue());
-					return null;
+				SameSite sameSite = (cookie != null) ? getSameSite(cookie) : null;
-				}
-				HttpCookie cookie = setCookieParser.parse(headerField.getValue());
-				if (cookie == null) {
-					return null;
-				}
-				SameSite sameSite = getSameSite(cookie);
 				if (sameSite == null) {
-					return null;
+					return field;
 				}
-				return new HttpCookieUtils.SetCookieHttpField(
+				HttpCookie updatedCookie = buildCookieWithUpdatedSameSite(cookie, sameSite);
-						HttpCookie.build(cookie)
+				return new HttpCookieUtils.SetCookieHttpField(updatedCookie, this.compliance);
-							.sameSite(org.eclipse.jetty.http.HttpCookie.SameSite.from(sameSite.name()))
+			}
-							.build(),
+
-						this.request.getConnectionMetaData().getHttpConfiguration().getResponseCookieCompliance());
+			private HttpCookie buildCookieWithUpdatedSameSite(HttpCookie cookie, SameSite sameSite) {
+				return HttpCookie.build(cookie)
+					.sameSite(org.eclipse.jetty.http.HttpCookie.SameSite.from(sameSite.name()))
+					.build();
 			}
 
 			private SameSite getSameSite(HttpCookie cookie) {
+				return getSameSite(asServletCookie(cookie));
+			}
+
+			private SameSite getSameSite(Cookie cookie) {
+				return SuppliedSameSiteCookieHandlerWrapper.this.suppliers.stream()
+					.map((supplier) -> supplier.getSameSite(cookie))
+					.filter(Objects::nonNull)
+					.findFirst()
+					.orElse(null);
+			}
+
+			private Cookie asServletCookie(HttpCookie cookie) {
 				Cookie servletCookie = new Cookie(cookie.getName(), cookie.getValue());
 				cookie.getAttributes().forEach(servletCookie::setAttribute);
-				for (CookieSameSiteSupplier supplier : SuppliedSameSiteCookieHandlerWrapper.this.suppliers) {
+				return servletCookie;
-					SameSite sameSite = supplier.getSameSite(servletCookie);
-					if (sameSite != null) {
-						return sameSite;
-					}
-				}
-				return null;
 			}
 
 		}