Use AcceptResponseHandler in goproxy https CONNECT hook (stripe#199)

* pipe AcceptResponseHandler into new goproxy hook

* update comment

* go mod vendor

* unit test

* use smokescreenctx in acceptresponsehandler

* fix unit tests

go.mod

@@ -12,7 +12,7 @@ require (
 	github.com/rs/xid v1.2.1
 	github.com/sirupsen/logrus v1.9.0
 	github.com/stretchr/testify v1.8.0
-	github.com/stripe/goproxy v0.0.0-20220308202309-3f1dfba6d1a4
+	github.com/stripe/goproxy v0.0.0-20230801191332-fabc3ecb7251
 	golang.org/x/net v0.7.0
 	gopkg.in/urfave/cli.v1 v1.20.0
 	gopkg.in/yaml.v2 v2.4.0

go.sum

@@ -218,6 +218,8 @@ github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PK
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
 github.com/stripe/goproxy v0.0.0-20220308202309-3f1dfba6d1a4 h1:6lqpRCXAhigpiMu6aZYTHryDmLrhIND0H4iUKEntN1M=
 github.com/stripe/goproxy v0.0.0-20220308202309-3f1dfba6d1a4/go.mod h1:hF2CVgH4++5ijZiy9grGVP8Fsi4u+SMOtbnIKYbMUjY=
+github.com/stripe/goproxy v0.0.0-20230801191332-fabc3ecb7251 h1:wR1exp7OglR0ctk8yWPVp1oTOuyaLUlJv3/Wlbvbw64=
+github.com/stripe/goproxy v0.0.0-20230801191332-fabc3ecb7251/go.mod h1:hF2CVgH4++5ijZiy9grGVP8Fsi4u+SMOtbnIKYbMUjY=
 github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.32/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=

pkg/smokescreen/config.go

@@ -82,8 +82,8 @@ type Config struct {
 	// Custom handler to allow clients to modify reject responses
 	RejectResponseHandler func(*http.Response)
 
-	// Custom handler to allow clients to modify accept responses
-	AcceptResponseHandler func(*http.Response)
+	// Custom handler to allow clients to modify successful CONNECT responses
+	AcceptResponseHandler func(*smokescreenContext, *http.Response) error
 
 	// UnsafeAllowPrivateRanges inverts the default behavior, telling smokescreen to allow private IP
 	// ranges by default (exempting loopback and unicast ranges)

pkg/smokescreen/smokescreen.go

@@ -546,7 +546,7 @@ func BuildProxy(config *Config) *goproxy.ProxyHttpServer {
 				resp.Header.Del(errorHeader)
 			}
 			if sctx.cfg.AcceptResponseHandler != nil {
-				sctx.cfg.AcceptResponseHandler(resp)
+				sctx.cfg.AcceptResponseHandler(sctx, resp)
 			}
 		}
 
@@ -559,6 +559,19 @@ func BuildProxy(config *Config) *goproxy.ProxyHttpServer {
 		logProxy(config, pctx)
 		return resp
 	})
+
+	// This function will be called on the response to a successful https CONNECT request.
+	// The goproxy OnResponse() function above is only called for non-https responses.
+	if config.AcceptResponseHandler != nil {
+		proxy.ConnectRespHandler = func(pctx *goproxy.ProxyCtx, resp *http.Response) error {
+			sctx, ok := pctx.UserData.(*smokescreenContext)
+			if !ok {
+				return fmt.Errorf("goproxy ProxyContext missing required UserData *smokescreenContext")
+			}
+			return config.AcceptResponseHandler(sctx, resp)
+		}
+	}
+
 	return proxy
 }
 

pkg/smokescreen/smokescreen_test.go

@@ -1073,8 +1073,9 @@ func TestAcceptResponseHandler(t *testing.T) {
 		cfg, err := testConfig("test-local-srv")
 
 		// set a custom AcceptResponseHandler that will set a header on every reject response
-		cfg.AcceptResponseHandler = func(resp *http.Response) {
+		cfg.AcceptResponseHandler = func(_ *smokescreenContext, resp *http.Response) error {
 			resp.Header.Set(testHeader, "This header is added by the AcceptResponseHandler")
+			return nil
 		}
 		r.NoError(err)
 

vendor/modules.txt

@@ -69,7 +69,7 @@ github.com/sirupsen/logrus/hooks/test
 ## explicit; go 1.13
 github.com/stretchr/testify/assert
 github.com/stretchr/testify/require
-# github.com/stripe/goproxy v0.0.0-20220308202309-3f1dfba6d1a4
+# github.com/stripe/goproxy v0.0.0-20230801191332-fabc3ecb7251
 ## explicit; go 1.13
 github.com/stripe/goproxy
 # golang.org/x/mod v0.8.0