Allow the agent to edit application response headers (#41)

Author: sada-sigsci

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Golang Module Release Notes
 
+## 1.14.0 2024-11-18
+
+* Allow the agent to edit application response headers
+
 ## 1.13.0 2023-07-06
 
 * Added new module configuration option for more granular inspection

VERSION

@@ -1 +1 @@
-1.13.0
+1.14.0

go.mod

@@ -1,10 +1,10 @@
 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
+require github.com/philhofer/fwd v1.1.3-0.20240916144458-20a13a1f6b7c // indirect

go.sum

@@ -1,6 +1,6 @@
-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/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=

module.go

@@ -12,9 +12,14 @@ import (
 	"sync"
 	"time"
 
+	"github.com/signalsciences/sigsci-module-golang/schema"
 	"github.com/signalsciences/tlstext"
 )
 
+type RPCMsgIn = schema.RPCMsgIn
+type RPCMsgIn2 = schema.RPCMsgIn2
+type RPCMsgOut = schema.RPCMsgOut
+
 // Module is an http.Handler that wraps an existing handler with
 // data collection and sends it to the Signal Sciences Agent for
 // inspection.
@@ -115,7 +120,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

@@ -6,6 +6,8 @@ import (
 	"io"
 	"net"
 	"net/http"
+
+	"github.com/signalsciences/sigsci-module-golang/schema"
 )
 
 // ResponseWriter is a http.ResponseWriter allowing extraction of data needed for inspection
@@ -24,12 +26,17 @@ type ResponseWriterFlusher interface {
 
 // NewResponseWriter returns a ResponseWriter or ResponseWriterFlusher depending on the base http.ResponseWriter.
 func NewResponseWriter(base http.ResponseWriter) ResponseWriter {
+	return newResponseWriter(base, nil)
+}
+
+func newResponseWriter(base http.ResponseWriter, actions []schema.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 +46,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 []schema.Action
 }
 
 // BaseResponseWriter returns the base http.ResponseWriter allowing access if needed
@@ -66,12 +74,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 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
+}
+
 // 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,7 +6,10 @@ import (
 	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
+	"reflect"
 	"testing"
+
+	"github.com/signalsciences/sigsci-module-golang/schema"
 )
 
 // testResponseRecorder is a httptest.ResponseRecorder without the Flusher interface
@@ -123,3 +126,31 @@ func TestResponseWriter(t *testing.T) {
 func TestResponseWriterFlusher(t *testing.T) {
 	testResponseWriter(t, NewResponseWriter(&testResponseRecorderFlusher{httptest.NewRecorder()}), 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 := []schema.Action{
+		{schema.AddHdr, []string{"csp", "src=abc"}},
+		{schema.SetHdr, []string{"content-type", "text/json"}},
+		{schema.DelHdr, []string{"x-powered-by"}},
+		{schema.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 renamed to schema/rpc.go

@@ -1,6 +1,6 @@
-package sigsci
+package schema
 
-//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