Allow the agent to edit application response headers

go.mod

@@ -1,10 +1,8 @@
 module github.com/signalsciences/sigsci-module-golang
 
-go 1.13
+go 1.21
 
 require (
 	github.com/signalsciences/tlstext v1.2.0
-	github.com/tinylib/msgp v1.1.0
+	github.com/tinylib/msgp v1.2.4
 )
-
-require github.com/philhofer/fwd v1.0.0 // indirect

go.sum

@@ -1,6 +1,73 @@
-github.com/philhofer/fwd v1.0.0 h1:UbZqGr5Y38ApvM/V/jEljVxwocdweyH+vmYvRPBnbqQ=
-github.com/philhofer/fwd v1.0.0/go.mod h1:gk3iGcWd9+svBvR0sR+KPcfE+RNWozjowpeBVG3ZVNU=
+github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
+github.com/philhofer/fwd v1.1.3-0.20240916144458-20a13a1f6b7c h1:dAMKvw0MlJT1GshSTtih8C2gDs04w8dReiOGXrGLNoY=
+github.com/philhofer/fwd v1.1.3-0.20240916144458-20a13a1f6b7c/go.mod h1:RqIHx9QI14HlwKwm98g9Re5prTQ6LdeRQn+gXJFxsJM=
 github.com/signalsciences/tlstext v1.2.0 h1:ps1ZCoDz93oMK0ySe7G/2J0dpTT32cN20U+/xy0S7uk=
 github.com/signalsciences/tlstext v1.2.0/go.mod h1:DKD8bjL8ZiedHAgWtcpgZm6TBAnmAlImuyJX2whrm3k=
-github.com/tinylib/msgp v1.1.0 h1:9fQd+ICuRIu/ue4vxJZu6/LzxN0HwMds2nq/0cFvxHU=
-github.com/tinylib/msgp v1.1.0/go.mod h1:+d+yLhGm8mzTaHzB+wgMYrodPfmZrzkirds8fDWklFE=
+github.com/tinylib/msgp v1.2.4 h1:yLFeUGostXXSGW5vxfT5dXG/qzkn4schv2I7at5+hVU=
+github.com/tinylib/msgp v1.2.4/go.mod h1:ykjzy2wzgrlvpDCRc4LA8UXy6D8bzMSuAF3WD57Gok0=
+github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
+golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
+golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
+golang.org/x/crypto v0.13.0/go.mod h1:y6Z2r+Rw4iayiXXAIxJIDAJ1zMW4yaTpebo8fPOliYc=
+golang.org/x/crypto v0.19.0/go.mod h1:Iy9bg/ha4yyC70EfRS8jz+B6ybOBKMaSxLj6P6oBDfU=
+golang.org/x/crypto v0.23.0/go.mod h1:CKFgDieR+mRhux2Lsu27y0fO304Db0wZe70UKqHu0v8=
+golang.org/x/crypto v0.24.0/go.mod h1:Z1PMYSOR5nyMcyAVAIQSKCDwalqy85Aqn1x3Ws4L5DM=
+golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
+golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
+golang.org/x/mod v0.12.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
+golang.org/x/mod v0.15.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
+golang.org/x/mod v0.17.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
+golang.org/x/mod v0.18.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
+golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
+golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
+golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=
+golang.org/x/net v0.6.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs=
+golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg=
+golang.org/x/net v0.15.0/go.mod h1:idbUs1IY1+zTqbi8yxTbhexhEEk5ur9LInksu6HrEpk=
+golang.org/x/net v0.21.0/go.mod h1:bIjVDfnllIU7BJ2DNgfnXvpSvtn8VRwhlsaeUTyUS44=
+golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM=
+golang.org/x/net v0.26.0/go.mod h1:5YKkiSynbBIh3p6iOc/vibscux0x38BZDkn8sCUPxHE=
+golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.3.0/go.mod h1:FU7BRWz2tNW+3quACPkgCx/L+uEAv1htQ0V83Z9Rj+Y=
+golang.org/x/sync v0.6.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
+golang.org/x/sync v0.7.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
+golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
+golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/sys v0.20.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/sys v0.21.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/telemetry v0.0.0-20240228155512-f48c80bd79b2/go.mod h1:TeRTkGYfJXctD9OcfyVLyj2J3IxLnKwHJR8f4D8a3YE=
+golang.org/x/telemetry v0.0.0-20240521205824-bda55230c457/go.mod h1:pRgIJT+bRLFKnoM1ldnzKoxTIn14Yxz928LQRYYgIN0=
+golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
+golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
+golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k=
+golang.org/x/term v0.8.0/go.mod h1:xPskH00ivmX89bAKVGSKKtLOWNx2+17Eiy94tnKShWo=
+golang.org/x/term v0.12.0/go.mod h1:owVbMEjm3cBLCHdkQu9b1opXd4ETQWc3BhuQGKgXgvU=
+golang.org/x/term v0.17.0/go.mod h1:lLRBjIVuehSbZlaOtGMbcMncT+aqLLLmKrsjNrUguwk=
+golang.org/x/term v0.20.0/go.mod h1:8UkIAJTvZgivsXaD6/pH6U9ecQzZ45awqEOzuCvwpFY=
+golang.org/x/term v0.21.0/go.mod h1:ooXLefLobQVslOqselCNF4SxFAaoS6KujMbsGzSDmX0=
+golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
+golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
+golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
+golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
+golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
+golang.org/x/text v0.13.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE=
+golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
+golang.org/x/text v0.15.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
+golang.org/x/text v0.16.0/go.mod h1:GhwF1Be+LQoKShO3cGOHzqOgRrGaYc9AvblQOmPVHnI=
+golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
+golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
+golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
+golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU=
+golang.org/x/tools v0.13.0/go.mod h1:HvlwmtVNQAhOuCjW7xxvovg8wbNq7LwfXh/k7wXUl58=
+golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d/go.mod h1:aiJjzUbINMkxbQROHiO6hDPo2LHcIPhhQsa9DLh0yGk=
+golang.org/x/tools v0.22.0/go.mod h1:aCwcsjqvq7Yqt6TNyX7QMU2enbQ/Gt0bo6krSeEri+c=
+golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=

module.go

@@ -115,7 +115,7 @@ func (m *Module) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 		return
 	}
 
-	rw := NewResponseWriter(w)
+	rw := NewResponseWriter(w, out.RespActions)
 
 	wafresponse := out.WAFResponse
 	switch {

responsewriter.go

@@ -23,13 +23,14 @@ type ResponseWriterFlusher interface {
 }
 
 // NewResponseWriter returns a ResponseWriter or ResponseWriterFlusher depending on the base http.ResponseWriter.
-func NewResponseWriter(base http.ResponseWriter) ResponseWriter {
+func NewResponseWriter(base http.ResponseWriter, actions []Action) ResponseWriter {
 	// NOTE: according to net/http docs, if WriteHeader is not called explicitly,
 	// the first call to Write will trigger an implicit WriteHeader(http.StatusOK).
 	// this is why the default code is 200 and it only changes if WriteHeader is called.
 	w := &responseRecorder{
-		base: base,
-		code: 200,
+		base:    base,
+		code:    200,
+		actions: actions,
 	}
 	if _, ok := w.base.(http.Flusher); ok {
 		return &responseRecorderFlusher{w}
@@ -39,9 +40,10 @@ func NewResponseWriter(base http.ResponseWriter) ResponseWriter {
 
 // responseRecorder wraps a base http.ResponseWriter allowing extraction of additional inspection data
 type responseRecorder struct {
-	base http.ResponseWriter
-	code int
-	size int64
+	base    http.ResponseWriter
+	code    int
+	size    int64
+	actions []Action
 }
 
 // BaseResponseWriter returns the base http.ResponseWriter allowing access if needed
@@ -66,12 +68,37 @@ func (w *responseRecorder) Header() http.Header {
 
 // WriteHeader writes the header, recording the status code for inspection
 func (w *responseRecorder) WriteHeader(status int) {
+	if w.actions != nil {
+		w.mergeHeader()
+	}
 	w.code = status
 	w.base.WriteHeader(status)
 }
 
+func (w *responseRecorder) mergeHeader() {
+	hdr := w.base.Header()
+	for _, a := range w.actions {
+		switch a.Code {
+		case AddHdr:
+			hdr.Add(a.Args[0], a.Args[1])
+		case SetHdr:
+			hdr.Set(a.Args[0], a.Args[1])
+		case SetNEHdr:
+			if len(hdr.Get(a.Args[0])) == 0 {
+				hdr.Set(a.Args[0], a.Args[1])
+			}
+		case DelHdr:
+			hdr.Del(a.Args[0])
+		}
+	}
+	w.actions = nil
+}
+
 // Write writes data, tracking the length written for inspection
 func (w *responseRecorder) Write(b []byte) (int, error) {
+	if w.actions != nil {
+		w.mergeHeader()
+	}
 	w.size += int64(len(b))
 	return w.base.Write(b)
 }

responsewriter_test.go

@@ -6,6 +6,7 @@ import (
 	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
+	"reflect"
 	"testing"
 )
 
@@ -116,10 +117,38 @@ func testResponseWriter(t *testing.T, w ResponseWriter, flusher bool) {
 
 // TestResponseWriter tests a non-flusher ResponseWriter
 func TestResponseWriter(t *testing.T) {
-	testResponseWriter(t, NewResponseWriter(&testResponseRecorder{httptest.NewRecorder()}), false)
+	testResponseWriter(t, NewResponseWriter(&testResponseRecorder{httptest.NewRecorder()}, nil), false)
 }
 
 // TestResponseWriterFlusher tests a flusher ResponseWriter
 func TestResponseWriterFlusher(t *testing.T) {
-	testResponseWriter(t, NewResponseWriter(&testResponseRecorderFlusher{httptest.NewRecorder()}), true)
+	testResponseWriter(t, NewResponseWriter(&testResponseRecorderFlusher{httptest.NewRecorder()}, nil), true)
+}
+
+func TestResponseHeaders(t *testing.T) {
+
+	resp := &httptest.ResponseRecorder{
+		HeaderMap: http.Header{
+			"X-Powered-By": []string{"aa"},
+			"Content-Type": []string{"text/plain"},
+			"X-Report":     []string{"bb"},
+		},
+	}
+	actions := []Action{
+		{AddHdr, []string{"csp", "src=abc"}},
+		{SetHdr, []string{"content-type", "text/json"}},
+		{DelHdr, []string{"x-powered-by"}},
+		{SetNEHdr, []string{"x-report", "cc"}},
+	}
+	NewResponseWriter(resp, actions).Write([]byte("foo"))
+
+	got := resp.Header()
+	expected := http.Header{
+		"Csp":          []string{"src=abc"},
+		"Content-Type": []string{"text/json"},
+		"X-Report":     []string{"bb"},
+	}
+	if !reflect.DeepEqual(got, expected) {
+		t.Fatalf("expected %v, got %v", expected, got)
+	}
 }

rpc.go

@@ -1,6 +1,6 @@
 package sigsci
 
-//go:generate msgp -unexported -tests=false
+//go:generate go run github.com/tinylib/msgp@v1.2.4 -unexported -tests=false
 
 //
 // This is for messages to and from the agent
@@ -34,8 +34,22 @@ type RPCMsgIn struct {
 // RPCMsgOut is sent back to the webserver
 type RPCMsgOut struct {
 	WAFResponse    int32
-	RequestID      string      `json:",omitempty"` // Set if the server expects an UpdateRequest with this ID (UUID)
-	RequestHeaders [][2]string `json:",omitempty"` // Any additional information in the form of additional request headers
+	RequestID      string      `json:",omitempty"`                  // Set if the server expects an UpdateRequest with this ID (UUID)
+	RequestHeaders [][2]string `json:",omitempty"`                  // Any additional information in the form of additional request headers
+	RespActions    []Action    `json:",omitempty" msg:",omitempty"` // Add or Delete application response headers
+}
+
+const (
+	AddHdr int8 = iota + 1
+	SetHdr
+	SetNEHdr
+	DelHdr
+)
+
+//msgp:tuple Action
+type Action struct {
+	Code int8
+	Args []string
 }
 
 // RPCMsgIn2 is a follow-up message from the webserver to the Agent