Add feature to optionally set headers after servlet execution (setHeadersAfterServlet) and to optionally append values (appendValues)

README.md

@@ -10,6 +10,24 @@ Add custom headers easily to your Tomcat response using param-name and param-val
     <filter>
         <filter-name>ResponseHeaderFilter</filter-name>
         <filter-class>pl.hordyjewiczmichal.ResponseHeaderFilter</filter-class>
+        <init-param>
+             <!-- optional parameter -->
+            <param-name>appendValues</param-name>
+            <!--
+              false: (default) filter sets a new header for every header value
+              true: filter appends values to the header in a comma-delimited list
+            -->
+            <param-value>false</param-value>
+        </init-param>
+        <init-param>
+             <!-- optional parameter -->
+            <param-name>setHeadersAfterServlet</param-name>
+            <!--
+              false: (default) filter sets headers *before* servlet execution
+              true: filter sets headers *after* servlet execution
+            -->
+            <param-value>false</param-value>
+        </init-param>
         <init-param>
             <param-name>Your-Header-1</param-name> <!-- put your any header name -->
             <param-value>
@@ -30,3 +48,7 @@ Add custom headers easily to your Tomcat response using param-name and param-val
     ```
 In `<init-param>` use `<param-name>` tag to add your custom header and below use `<param-value>` to add value to your header.
 You can use multiple values for your header in `<param-value>` - just separate them by a new line.
+
+The `appendValues` parameter controls whether the filter appends values to a header in a comma-delimited list or sets a new header for every header value (default).
+
+The `setHeadersAfterServlet` parameter controls whether headers are set before or after the Java application processes the request.

pom.xml

@@ -40,6 +40,14 @@
             <version>2.23.4</version>
             <scope>test</scope>
         </dependency>
+
+        <!-- https://mvnrepository.com/artifact/org.springframework/spring-test -->
+        <dependency>
+            <groupId>org.springframework</groupId>
+            <artifactId>spring-test</artifactId>
+            <version>5.3.29</version>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
 
@@ -52,6 +60,15 @@
                 <artifactId>maven-jar-plugin</artifactId>
                 <version>3.1.1</version>
             </plugin>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-compiler-plugin</artifactId>
+                <version>3.8.1</version>
+                <configuration>
+                    <source>9</source>
+                    <target>9</target>
+                </configuration>
+            </plugin>
         </plugins>
     </build>
 </project>

src/main/java/pl/hordyjewiczmichal/BufferedHttpServletResponseWrapper.java

@@ -0,0 +1,84 @@
+package pl.hordyjewiczmichal;
+
+import javax.servlet.ServletOutputStream;
+import javax.servlet.WriteListener;
+import javax.servlet.http.HttpServletResponse;
+import javax.servlet.http.HttpServletResponseWrapper;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.PrintWriter;
+
+/**
+ * A custom HttpServletResponseWrapper that buffers the response content.
+ * This allows for inspection and modification of the response headers and content
+ * before the response is committed to the client.
+ */
+public class BufferedHttpServletResponseWrapper extends HttpServletResponseWrapper {
+    private final ByteArrayOutputStream buffer = new ByteArrayOutputStream();
+    private final PrintWriter writer = new PrintWriter(buffer);
+
+    /**
+     * Constructs a response adaptor wrapping the given response.
+     *
+     * @param response the original HttpServletResponse to wrap
+     */
+    public BufferedHttpServletResponseWrapper(HttpServletResponse response) {
+        super(response);
+    }
+
+    /**
+     * Returns a ServletOutputStream suitable for writing binary data in the response.
+     *
+     * @return a ServletOutputStream for writing binary data
+     * @throws IOException if an input or output exception occurred
+     */
+    @Override
+    public ServletOutputStream getOutputStream() {
+        return new ServletOutputStream() {
+            @Override
+            public void write(int b) throws IOException {
+                buffer.write(b);
+            }
+
+            @Override
+            public boolean isReady() {
+                return true;
+            }
+
+            @Override
+            public void setWriteListener(WriteListener writeListener) {
+                // No-op implementation
+            }
+        };
+    }
+
+    /**
+     * Returns a PrintWriter object that can send character text to the client.
+     *
+     * @return a PrintWriter object that can return character data to the client
+     * @throws IOException if an input or output exception occurred
+     */
+    @Override
+    public PrintWriter getWriter() {
+        return writer;
+    }
+
+    /**
+     * Returns the buffered response content as a byte array.
+     *
+     * @return the buffered response content
+     */
+    public byte[] getBuffer() {
+        return buffer.toByteArray();
+    }
+
+    /**
+     * Flushes the buffer content to the wrapped response's output stream.
+     *
+     * @throws IOException if an input or output exception occurred
+     */
+    public void flushBuffer() throws IOException {
+        writer.flush();
+        buffer.writeTo(getResponse().getOutputStream());
+    }
+}

src/main/java/pl/hordyjewiczmichal/ResponseHeaderFilter.java

@@ -5,70 +5,83 @@
 import java.io.IOException;
 import java.util.*;
 
-@SuppressWarnings("WeakerAccess")
-public class ResponseHeaderFilter implements Filter
-{
-    private Map<String, List<String>> headers;
-
+public class ResponseHeaderFilter implements Filter {
+    private Map<String, List<String>> headersToSet = new HashMap<>();
+    private boolean setHeadersAfterServlet = false;
+    private boolean appendValues = false;
 
     @Override
-    public void init(FilterConfig config)
-    {
-        this.headers = new HashMap<>();
-
+    public void init(FilterConfig config) {
         Enumeration<String> paramNames = config.getInitParameterNames();
-
         if (paramNames == null) return;
 
-        while (paramNames.hasMoreElements())
-        {
-            String headerName = paramNames.nextElement();
-            String headerValues = config.getInitParameter(headerName); // can contain multi-values (separated by new line)
-
-            List<String> values = new ArrayList<>();
-
-            try
-            {
-                if ("".equals(headerValues)) throw new IllegalArgumentException();
+        this.setHeadersAfterServlet = Boolean.parseBoolean(config.getInitParameter("setHeadersAfterServlet"));
+        this.appendValues = Boolean.parseBoolean(config.getInitParameter("appendValues"));
 
-                StringTokenizer st = new StringTokenizer(headerValues, "\r\n");
-                while (st.hasMoreTokens())
-                {
-                    String nextValue = st.nextToken().trim();
-                    if("".equals(nextValue)) continue;
-                    values.add(nextValue);
-                }
+        while (paramNames.hasMoreElements()) {
+            String paramName = paramNames.nextElement();
+            if ("setHeadersAfterServlet".equals(paramName) || "appendValues".equals(paramName)) {
+                continue; // skip these metadata parameters
             }
-            catch (NullPointerException | IllegalArgumentException e)
-            {
-                continue; // skip if no value specified
+            String paramValue = config.getInitParameter(paramName);
+            if (paramValue != null && !paramValue.trim().isEmpty()) {
+                headersToSet.put(paramName, Arrays.asList(paramValue.split("\n")));
             }
-
-            this.headers.put(headerName, values);
         }
     }
 
     @Override
-    public void doFilter(ServletRequest req, ServletResponse resp, FilterChain chain) throws ServletException, IOException
-    {
-        HttpServletResponse response = (HttpServletResponse) resp;
-        Set<String> headerNames = this.getHeaders().keySet();
+    public void doFilter(ServletRequest req, ServletResponse resp, FilterChain chain) throws ServletException, IOException {
+        // use a wrapper to buffer the response content; otherwise the response will be committed when chain.doFilter
+        //  is called
+        BufferedHttpServletResponseWrapper response = new BufferedHttpServletResponseWrapper((HttpServletResponse) resp);
+
+        if (!setHeadersAfterServlet) {
+            setHeadersToSet(response);
+        }
+
+        chain.doFilter(req, response);
 
-        headerNames.forEach(headerName ->
-                this.getHeaders().get(headerName).forEach(val ->
-                        response.addHeader(headerName, val)));
+        if (setHeadersAfterServlet) {
+            setHeadersToSet(response);
+        }
 
-        chain.doFilter(req, resp);
+        // commit the response
+        response.flushBuffer();
+    }
+
+    private void setHeader(HttpServletResponse response, String name, String value, boolean append) {
+        if (append && response.containsHeader(name)) {
+            String existingValue = response.getHeader(name);
+            String newValue = String.join(", ", existingValue, value);
+            response.setHeader(name, newValue);
+        } else {
+            response.addHeader(name, value);
+        }
+    }
+
+    private void setHeadersToSet(HttpServletResponse response) {
+        headersToSet.forEach((headerName, values) -> {
+            values.forEach(value -> setHeader(response, headerName, value, appendValues));
+        });
     }
 
     @Override
-    public void destroy()
-    {
+    public void destroy() {
         // nothing to clean up.
     }
 
-    public Map<String, List<String>> getHeaders()
-    {
+    public Map<String, List<String>> getHeadersToSet() {
+        return headersToSet;
+    }
+
+    // helper for step debugging
+    public Map<String, String> getResponseHeaders(HttpServletResponse response) {
+        Map<String, String> headers = new HashMap<>();
+        Collection<String> headerNames = response.getHeaderNames();
+        for (String headerName : headerNames) {
+            headers.put(headerName, response.getHeader(headerName));
+        }
         return headers;
     }
-}
+}