WICKET-7107 improved BaseWicketTester#WicketTesterServletWebResponse

wicket-core-tests/src/test/java/org/apache/wicket/csp/CSPRequestCycleListenerTest.java

@@ -0,0 +1,136 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.wicket.csp;
+
+
+import org.apache.wicket.Page;
+import org.apache.wicket.RestartResponseException;
+import org.apache.wicket.core.request.handler.PageProvider;
+import org.apache.wicket.core.request.handler.RenderPageRequestHandler.RedirectPolicy;
+import org.apache.wicket.markup.html.WebPage;
+import org.apache.wicket.protocol.http.WebApplication;
+import org.apache.wicket.util.tester.DummyHomePage;
+import org.apache.wicket.util.tester.WicketTestCase;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.MethodSource;
+
+import java.io.Serial;
+import java.util.List;
+
+class CSPRequestCycleListenerTest extends WicketTestCase
+{
+	public static List<Class<? extends Page>> pageClasses()
+	{
+		return List.of( DummyHomePage.class, RedirectPage.class, RestartPage.class, StatelessAutoRedirectPage.class, StatefulAutoRedirectPage.class );
+	}
+
+	protected WebApplication newApplication()
+	{
+		return new WebApplication()
+		{
+			@Override
+			public Class<? extends Page> getHomePage()
+			{
+				return DummyHomePage.class;
+			}
+
+			@Override
+			protected void init()
+			{
+				super.init();
+
+				getCspSettings().blocking().strict()
+						.add(CSPDirective.STYLE_SRC, CSPDirectiveSrcValue.SELF)
+						.add(CSPDirective.STYLE_SRC, "https://fonts.foo.bar/css")
+						.add(CSPDirective.FONT_SRC, "https://fonts.foo.bar");
+
+				mountPage("redirect", RedirectPage.class);
+			}
+		};
+	}
+
+	@ParameterizedTest(name="pageClass={arguments}")
+	@MethodSource("pageClasses")
+	<C extends Page> void strict(final Class<C> pageClass)
+	{
+		tester.startPage(pageClass);
+		tester.assertRenderedPage(DummyHomePage.class);
+
+		final String cspHeaderValue = tester.getLastResponse().getHeader("Content-Security-Policy");
+		Assertions.assertEquals("default-src 'none'; script-src 'strict-dynamic' 'NONCE'; style-src 'NONCE' 'self' https://fonts.foo.bar/css; img-src 'self'; connect-src 'self'; font-src 'self' https://fonts.foo.bar; manifest-src 'self'; child-src 'self'; base-uri 'self'; frame-src 'self'", cspHeaderValue.replaceAll("'nonce-[^']+'", "'NONCE'"));
+	}
+
+	public static class RestartPage extends WebPage
+	{
+		@Serial
+		private static final long serialVersionUID = 1L;
+
+		public RestartPage()
+		{
+			throw new RestartResponseException(
+					new PageProvider(DummyHomePage.class),
+					RedirectPolicy.NEVER_REDIRECT
+			);
+		}
+	}
+
+	public static class RedirectPage extends WebPage
+	{
+		@Serial
+		private static final long serialVersionUID = 1L;
+
+		public RedirectPage()
+		{
+			throw new RestartResponseException(
+					new PageProvider(DummyHomePage.class),
+					RedirectPolicy.ALWAYS_REDIRECT
+			);
+		}
+	}
+
+	public static class StatelessAutoRedirectPage extends WebPage
+	{
+		@Serial
+		private static final long serialVersionUID = 1L;
+
+		public StatelessAutoRedirectPage()
+		{
+			final DummyHomePage page = new DummyHomePage();
+			throw new RestartResponseException(
+					new PageProvider(page.setStatelessHint(true)),
+					RedirectPolicy.AUTO_REDIRECT
+			);
+		}
+	}
+
+	public static class StatefulAutoRedirectPage extends WebPage
+	{
+		@Serial
+		private static final long serialVersionUID = 1L;
+
+		public StatefulAutoRedirectPage()
+		{
+			final DummyHomePage page = new DummyHomePage();
+			throw new RestartResponseException(
+					new PageProvider(page.setStatelessHint(false)),
+					RedirectPolicy.AUTO_REDIRECT
+			);
+		}
+	}
+}

wicket-core/src/main/java/org/apache/wicket/csp/CSPRequestCycleListener.java

@@ -18,6 +18,7 @@
 
 import org.apache.wicket.request.IRequestHandler;
 import org.apache.wicket.request.IRequestHandlerDelegate;
+import org.apache.wicket.request.Url;
 import org.apache.wicket.request.cycle.IRequestCycleListener;
 import org.apache.wicket.request.cycle.RequestCycle;
 import org.apache.wicket.request.http.WebResponse;
@@ -39,14 +40,7 @@ public CSPRequestCycleListener(ContentSecurityPolicySettings settings)
 	}
 
 	@Override
-	public void onRequestHandlerResolved(RequestCycle cycle, IRequestHandler handler)
-	{
-		// WICKET-7028- this is needed for redirect to buffer use case.
-		protect(cycle, handler);
-	}
-
-	@Override
-	public void onRequestHandlerExecuted(RequestCycle cycle, IRequestHandler handler)
+	public void onUrlMapped(RequestCycle cycle, IRequestHandler handler, Url url)
 	{
 		protect(cycle, handler);
 	}

wicket-tester/src/main/java/org/apache/wicket/util/tester/BaseWicketTester.java

@@ -24,9 +24,13 @@
 import java.lang.reflect.Constructor;
 import java.lang.reflect.Method;
 import java.nio.charset.Charset;
+import java.text.DateFormat;
 import java.text.ParseException;
 import java.time.Duration;
+import java.time.Instant;
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Date;
 import java.util.Enumeration;
 import java.util.HashMap;
 import java.util.List;
@@ -36,6 +40,7 @@
 import java.util.Set;
 import java.util.UUID;
 import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 
 import org.apache.wicket.Application;
 import org.apache.wicket.Component;
@@ -125,6 +130,7 @@
 import org.apache.wicket.util.lang.Generics;
 import org.apache.wicket.util.resource.StringResourceStream;
 import org.apache.wicket.util.string.Strings;
+import org.apache.wicket.util.value.ValueMap;
 import org.apache.wicket.util.visit.IVisit;
 import org.apache.wicket.util.visit.IVisitor;
 import org.slf4j.Logger;
@@ -2860,6 +2866,7 @@ private static class WicketTesterServletWebResponse extends ServletWebResponse
 			IMetaDataBufferingWebResponse
 	{
 		private List<Cookie> cookies = new ArrayList<Cookie>();
+		private final ValueMap headers = new ValueMap();
 
 		private WicketTesterServletWebResponse(ServletWebRequest request,
 			MockHttpServletResponse response)
@@ -2881,6 +2888,10 @@ public void writeMetaData(WebResponse webResponse)
 			{
 				webResponse.addCookie(cookie);
 			}
+			for (String name : headers.keySet())
+			{
+				webResponse.setHeader(name, Arrays.stream(headers.getStringArray(name)).collect(Collectors.joining(";")));
+			}
 		}
 
 		@Override
@@ -2896,6 +2907,28 @@ public void sendRedirect(String url)
 				throw new RuntimeException(e);
 			}
 		}
+
+		@Override
+		public void setHeader(String name, String value)
+		{
+			super.setHeader(name, value);
+			headers.put(name, value);
+		}
+
+		@Override
+		public void addHeader(String name, String value)
+		{
+			super.addHeader(name, value);
+			headers.add(name, value);
+		}
+
+		@Override
+		public void reset()
+		{
+			super.reset();
+			cookies.clear();
+			headers.clear();
+		}
 	}
 
 	private class LastPageRecordingPageRendererProvider implements IPageRendererProvider