Modifying response headers with custom modules

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35

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;
     }
   }
 }
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

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;
   }
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

@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
    }
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51

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], "");
     }
   }
 }
}