how to add custom response HTTP header in Servlet Filter that depends on status code returned from app

Question

for logging (MDC) I need to set a custom header when the response is failed as the following:

public class MDCFilter implements Filter {
    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException {
        MDC.put("trackingId", UUID.randomUUID().toString());

        try {
            chain.doFilter(request, response);
        } catch (Throwable throwable) {
            //suppress
        }

        if (((HttpServletResponse) response).getStatus() >= 300) {
            ((HttpServletResponse) response).setHeader("x-tracking-id", MDC.get("trackingId"));
//            ((HttpServletResponse) response).getHeader("x-tracking-id"); //this returns null
        }
        MDC.remove("trackingId");
    }
}

but this does not work, no header is set. if I set the header before chain.doFilter it works, but I need to set this filter as late as possible, I need to know if the response status is OK or failed.

Answer 1

You need to set the header as soon as the status is set, not after when the response is already sent to the client. That's a point of no return. You can use a HttpServletResponseWrapper to decorate the behavior of setStatus() method.

For example,

@Override
public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException {
    try {
        MDC.put("trackingId", UUID.randomUUID().toString());
        chain.doFilter(request, new HttpServletResponseWrapper((HttpServletResponse) response) {
            @Override
            public void setStatus(int status) {
                super.setStatus(status);

                if (status >= 300) {
                    setHeader("x-tracking-id", MDC.get("trackingId"));
                }
            }
        });
    }
    finally {
        MDC.remove("trackingId");
    }
}

Note that I also fixed the sloppy exception handling.