Added CompositeHeaderWriter

web/src/main/java/org/springframework/security/web/header/HeaderWriterFilter.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2018 the original author or authors.
+ * Copyright 2002-2019 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.
@@ -27,6 +27,7 @@
 import javax.servlet.http.HttpServletRequestWrapper;
 import javax.servlet.http.HttpServletResponse;
 
+import org.springframework.security.web.header.writers.CompositeHeaderWriter;
 import org.springframework.security.web.util.OnCommittedResponseWrapper;
 import org.springframework.util.Assert;
 import org.springframework.web.filter.OncePerRequestFilter;
@@ -38,55 +39,55 @@
  *
  * @author Marten Deinum
  * @author Josh Cummings
+ * @author Ankur Pathak
  * @since 3.2
- *
  */
 public class HeaderWriterFilter extends OncePerRequestFilter {
 
+
 	/**
-	 * Collection of {@link HeaderWriter} instances to write out the headers to the
+	 * Composite {@link CompositeHeaderWriter} containing Collection of {@link HeaderWriter} instances to write out the headers to the
 	 * response.
 	 */
-	private final List<HeaderWriter> headerWriters;
+	private final HeaderWriter headerWriter;
 
 	/**
 	 * Creates a new instance.
 	 *
 	 * @param headerWriters the {@link HeaderWriter} instances to write out headers to the
-	 * {@link HttpServletResponse}.
+	 *                      {@link HttpServletResponse}.
 	 */
 	public HeaderWriterFilter(List<HeaderWriter> headerWriters) {
 		Assert.notEmpty(headerWriters, "headerWriters cannot be null or empty");
-		this.headerWriters = headerWriters;
+		this.headerWriter = new CompositeHeaderWriter(headerWriters);
 	}
 
 	@Override
 	protected void doFilterInternal(HttpServletRequest request,
-			HttpServletResponse response, FilterChain filterChain)
-					throws ServletException, IOException {
+									HttpServletResponse response, FilterChain filterChain)
+			throws ServletException, IOException {
 
 		HeaderWriterResponse headerWriterResponse = new HeaderWriterResponse(request,
-				response, this.headerWriters);
+				response, this.headerWriter);
 		HeaderWriterRequest headerWriterRequest = new HeaderWriterRequest(request,
 				headerWriterResponse);
 
 		try {
 			filterChain.doFilter(headerWriterRequest, headerWriterResponse);
-		}
-		finally {
+		} finally {
 			headerWriterResponse.writeHeaders();
 		}
 	}
 
 	static class HeaderWriterResponse extends OnCommittedResponseWrapper {
 		private final HttpServletRequest request;
-		private final List<HeaderWriter> headerWriters;
+		private final HeaderWriter headerWriter;
 
 		HeaderWriterResponse(HttpServletRequest request, HttpServletResponse response,
-				List<HeaderWriter> headerWriters) {
+				HeaderWriter headerWriter) {
 			super(response);
 			this.request = request;
-			this.headerWriters = headerWriters;
+			this.headerWriter = headerWriter;
 		}
 
 		/*
@@ -105,9 +106,7 @@ protected void writeHeaders() {
 			if (isDisableOnResponseCommitted()) {
 				return;
 			}
-			for (HeaderWriter headerWriter : this.headerWriters) {
-				headerWriter.writeHeaders(this.request, getHttpResponse());
-			}
+			this.headerWriter.writeHeaders(this.request, getHttpResponse());
 		}
 
 		private HttpServletResponse getHttpResponse() {

web/src/main/java/org/springframework/security/web/header/writers/CompositeHeaderWriter.java

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2002-2019 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.
+ * 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.springframework.security.web.header.writers;
+
+import org.springframework.security.web.header.HeaderWriter;
+import org.springframework.util.Assert;
+
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+import java.util.List;
+
+/**
+ * A {@link HeaderWriter} that delegates to several other {@link HeaderWriter}s.
+ *
+ * @author Ankur Pathak
+ * @since 5.2
+ */
+public class CompositeHeaderWriter implements HeaderWriter {
+
+	private final List<HeaderWriter> headerWriters;
+
+	/**
+	 * Creates a new instance.
+	 *
+	 * @param headerWriters the {@link HeaderWriter} instances to write out headers to the
+	 *                      {@link HttpServletResponse}.
+	 */
+	public CompositeHeaderWriter(List<HeaderWriter> headerWriters) {
+		Assert.notEmpty(headerWriters, "headerWriters cannot be empty");
+		this.headerWriters = headerWriters;
+	}
+
+	@Override
+	public void writeHeaders(HttpServletRequest request, HttpServletResponse response) {
+		this.headerWriters.forEach(headerWriter -> headerWriter.writeHeaders(request, response));
+	}
+}

web/src/test/java/org/springframework/security/web/header/writers/CompositeHeaderWriterTests.java

@@ -0,0 +1,69 @@
+/*
+ * Copyright 2002-2019 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.
+ * 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.springframework.security.web.header.writers;
+
+import java.util.Arrays;
+
+import org.junit.Before;
+import org.junit.Test;
+
+import org.springframework.mock.web.MockHttpServletRequest;
+import org.springframework.mock.web.MockHttpServletResponse;
+import org.springframework.security.web.header.HeaderWriter;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Tests for class {@link CompositeHeaderWriter}/.
+ *
+ * @author Ankur Pathak
+ * @since 5.2
+ */
+public class CompositeHeaderWriterTests {
+	private MockHttpServletRequest request;
+
+	private MockHttpServletResponse response;
+
+	private CompositeHeaderWriter writer;
+
+	@Before
+	public void setup() {
+		this.request = new MockHttpServletRequest();
+		this.response = new MockHttpServletResponse();
+		HeaderWriter writerA = (request, response) -> {
+			if (!response.containsHeader("A")) {
+				response.setHeader("A", "a");
+			}
+		};
+
+		HeaderWriter writerB = (request, response) -> {
+			if (!response.containsHeader("B")) {
+				response.setHeader("B", "b");
+			}
+		};
+		this.writer = new CompositeHeaderWriter(Arrays.asList(writerA, writerB));
+	}
+
+
+	@Test
+	public void doCompositeHeaderWrite(){
+		this.writer.writeHeaders(request, response);
+		assertThat(this.response.containsHeader("A")).isTrue();
+		assertThat(this.response.containsHeader("B")).isTrue();
+		assertThat(this.response.getHeader("A")).isEqualTo("a");
+		assertThat(this.response.getHeader("B")).isEqualTo("b");
+	}
+}