Modifying response headers with custom modules

You can modify your custom module to support adding and editing HTTP response headers. Support for modifying response headers is a requirement for Fastly Client-Side Protection. You can also use this functionality to normalize response headers across multiple web applications.

To add support for response header modification, update your custom module to:

save the response actions that are returned as part of the PreRequest function.
apply the saved headers to responses once the Next-Gen WAF agent finishes request processing.

Limitations and considerations

The framework that your custom module integrates with may have limitations that impact response header modification. For example, some web serves and frameworks do not support deleting a header key in the response. In those cases, we have set the value to an empty string.

Request flow

When you use the Next-Gen WAF to modify response headers, the request flow is as follows:

The web server or framework passes the request to your custom Next-Gen WAF module.
The module calls the PreRequest function in the Next-Gen WAF agent. The agent returns the RPCMsgOut structure, which may contain one or more actions to perform on the response. The actions tell the module which response headers to add or modify.
The module stores these actions while the agent processes the request.
The module passes the request to the agent for evaluation.
If the request is allowed, the agent forwards the request to the origin. The origin sends a response to the agent and the agent forwards that response to the module.
Based on the response actions the module saved during step two, the module adds and edits the appropriate response headers.
The module sends the response to the web client that initiated the request.

Saving response actions

With a custom module, the PreRequest function returns response actions in the RPCMsgOut structure as the RespActions array. To save these actions while the agent processes a request, complete the following:

Java module
Go module

For a custom Java module, you need to unpack the MessagePack response. The following example shows how to unpack the actions from the response:

public class RPCMsgOut {
 public Action[] responseActions;

 void unpack(MessageUnpacker u) throws Exception {
   int mlen = u.unpackMapHeader();
   for (int j = 0; j < mlen; j++) {
     String value = u.unpackString();
     switch (value) {
       case "RespActions":
         if (u.tryUnpackNil())
           break;
         int actionsLen = u.unpackArrayHeader();
         if (actionsLen == 0)
           break;
         responseActions = new Action[actionsLen];
         for (int k = 0; k < actionsLen; k++) {
           if (u.unpackArrayHeader() != 2) {
             throw new IllegalArgumentException("Invalid tuple length for RespActions");
           }
           int code = u.unpackInt();
           int argsLen = u.unpackArrayHeader();
           String[] args = new String[argsLen];
           for (int i = 0; i < argsLen; i++) {
             args[i] = u.unpackString();
           }
           responseActions[k] = new Action(code, args);
         }
         break;
       default:
         u.skipValue();
         break;
     }
   }
 }
}

The Action class is as follows:

public class Action {

   public static final int AddHdr = 1;
   public static final int SetHdr = 2;
   public static final int SetNEHdr = 3;
   public static final int DelHdr = 4;

   private int code;
   public String[] args;

   public Action() {
       // Default constructor
   }

   public Action(int code, String[] args) {
       this.code = code;
       this.args = args;
   }

   public String[] getArgs() {
       return args;
   }

   public int getCode() {
       return code;
   }
}

For a custom Go module, you need to add the RespActions structure. The following example shows what the structure looks like:

type RPCMsgOut struct {
    ...
    RespActions []Action // Slice of response actions
}

const (
    AddHdr int8 = iota + 1 // value 1, add the header
    SetHdr               // value 2, replace any existing header
    SetNEHdr             // value 3, add the header if not exist
    DelHdr               // value 4, delete the header
)

type Action struct {
    Code uint16
    Args []string
}

Setting response headers

To update your custom module to act on the saved response actions, complete the following:

Java module
Go module

For a custom Java module, you need to implement a Jakarta handler to modify the response headers as demonstrated in the following code snippet.

IMPORTANT: If you're using a framework other than Java, you'll need to modify this code.

@Override
   public void handle(String target, Request request,
                      HttpServletRequest httpRequest, HttpServletResponse httpResponse)
           throws IOException, ServletException {
       if (isInvalidRequest(target, request, httpRequest, httpResponse)) {
           return;
       }
       RPCMsgIn msgIn = filterService.newRPCMsgIn(httpRequest);
       if (request.isSecure()) {
           SSLSession sslSession = (SSLSession) request.getAttribute("org.eclipse.jetty.servlet.request.ssl_session");
           msgIn.TLSCipher = sslSession.getCipherSuite();
           msgIn.TLSProtocol = sslSession.getProtocol();
       }

       httpRequest = filterService.readRequestBody(httpRequest, httpResponse, msgIn);
       RPCMsgOut rpcResponse = filterService.getPreRequest(httpRequest, msgIn);

       try {
           /* Fail safe on no response from agent */
           if (rpcResponse == null) {
               getHandler().handle(target, request, httpRequest, httpResponse);
               return;
           }

           /*
            * Add any request headers returned from agent to the http request
            *  Set the rpc request/response attributes on http request
            */

           if (rpcResponse.requestHeaders != null && rpcResponse.requestHeaders.length > 0) {
               httpRequest = new HeaderDecoratorRequestWrapper(httpRequest, rpcResponse.requestHeaders);
           }

           /*
           * Check if there are response actions to apply to the response and then apply them.
           */

           if (rpcResponse.responseActions != null) {
               // Use CustomRequest to handle response actions in onResponseCommit
               com.signalsciences.jetty.JettyRequestWrapper customRequest = new JettyRequestWrapper(
                       HttpConnection.getCurrentConnection().getHttpChannel().getRequest(),
                       rpcResponse.responseActions
               );
               httpResponse = new ResponseWrapper(httpResponse, rpcResponse.responseActions);
               request = customRequest;
           }
       } catch (Exception e) {
        // Handle exceptions
    }
}

The ResponseWrapper class is responsible for implementing the response header actions.

public class ResponseWrapper extends HttpServletResponseWrapper {
 private boolean headersModified = false;
 private Action[] actions;

 public ResponseWrapper(HttpServletResponse resp, Action[] actions) {
   super(resp);
   this.actions = actions;
 }

 @Override
 public void flushBuffer() throws IOException {
   ensureHeadersHandled();
   super.flushBuffer();
 }

 @Override
 public ServletOutputStream getOutputStream() throws IOException {
   ensureHeadersHandled();
   return super.getOutputStream();
 }

 @Override
 public PrintWriter getWriter() throws IOException {
   ensureHeadersHandled();
   return super.getWriter();
 }

 public void ensureHeadersHandled() {
   if (!headersModified) {
     headersModified = true;
     handleResponseActions();
   }
 }

 private void handleResponseActions() {
   for (Action action : actions) {
     String[] args = action.args;
     if (action.getCode() == Action.AddHdr) {
       super.addHeader(args[0], args[1]);
     } else if (action.getCode() == Action.SetHdr) {
       super.setHeader(args[0], args[1]);
     } else if (action.getCode() == Action.SetNEHdr) {
       if (!super.containsHeader(args[0])) {
         super.addHeader(args[0], args[1]);
       }
     } else if (action.getCode() == Action.DelHdr) {
       super.setHeader(args[0], "");
     }
   }
 }
}

For a custom Java module, you need to hook into the WriteHeader function and use the function to modify the headers according to the actions we received previously. The following code sample provides an example for how to do this.

// Check if we have response actions before calling through to
// WriteHeader in net/http
func (w *responseRecorder) WriteHeader(status int) {
   if w.actions != nil {
       w.mergeHeader()
   }
   w.code = status
   w.base.WriteHeader(status)
}

// Merge in our header modifications using the appropriate functions in
// the Header type in the net/http package
func (w *responseRecorder) mergeHeader() {
   hdr := w.base.Header()
   for _, a := range w.actions {
       switch a.Code {
       case schema.AddHdr:
           hdr.Add(a.Args[0], a.Args[1])
       case schema.SetHdr:
           hdr.Set(a.Args[0], a.Args[1])
       case schema.SetNEHdr:
           if len(hdr.Get(a.Args[0])) == 0 {
               hdr.Set(a.Args[0], a.Args[1])
           }
       case schema.DelHdr:
           hdr.Del(a.Args[0])
       }
   }
   w.actions = nil
}

Examples

For examples on how to implement support for receiving and acting on response actions, check out these working implementations:

Java module directory for v2.7.0 and above
Go module repository and relevant pull request

About the module

Network services

Security

Compute

Quick start

Building blocks

Integrations

Tutorials

Demos

Use Cases

Code Examples

Starter Kits

Modifying response headers with custom modules

Limitations and considerations

Request flow

Saving response actions

Setting response headers

Examples

Network services

Security

Compute

Quick start

Building blocks

Integrations

Tutorials

Demos

Use Cases

Code Examples

Starter Kits

Modifying response headers with custom modules

Limitations and considerations

Request flow

Saving response actions

Setting response headers

Examples

Related content