Add Support for Reject Handler with Context

pkg/smokescreen/config.go

@@ -87,6 +87,9 @@ type Config struct {
 	// Custom handler to allow clients to modify reject responses
 	RejectResponseHandler func(*http.Response)
 
+	// Custom handler to allow clients to modify reject responses
+	RejectResponseHandlerWithCtx func(*SmokescreenContext, *http.Response)
+
 	// Custom handler to allow clients to modify successful CONNECT responses
 	AcceptResponseHandler func(*SmokescreenContext, *http.Response) error
 

pkg/smokescreen/smokescreen.go

@@ -404,6 +404,9 @@ func rejectResponse(pctx *goproxy.ProxyCtx, err error) *http.Response {
 	if sctx.cfg.RejectResponseHandler != nil {
 		sctx.cfg.RejectResponseHandler(resp)
 	}
+	if sctx.cfg.RejectResponseHandlerWithCtx != nil {
+		sctx.cfg.RejectResponseHandlerWithCtx(sctx, resp)
+	}
 	return resp
 }
 

pkg/smokescreen/smokescreen_test.go

@@ -1067,6 +1067,47 @@ func TestRejectResponseHandler(t *testing.T) {
 	})
 }
 
+func TestRejectResponseHandlerWithCtx(t *testing.T) {
+	r := require.New(t)
+	testHeader := "TestRejectResponseHandlerWithCtxHeader"
+	t.Run("Testing custom reject response handler", func(t *testing.T) {
+		cfg, err := testConfig("test-local-srv")
+
+		// set a custom RejectResponseHandler that will set a header on every reject response
+		cfg.RejectResponseHandlerWithCtx = func(_ *SmokescreenContext, resp *http.Response) {
+			resp.Header.Set(testHeader, "This header is added by the RejectResponseHandlerWithCtx")
+		}
+		r.NoError(err)
+
+		proxySrv := proxyServer(cfg)
+		r.NoError(err)
+		defer proxySrv.Close()
+
+		// Create a http.Client that uses our proxy
+		client, err := proxyClient(proxySrv.URL)
+		r.NoError(err)
+
+		// Send a request that should be blocked
+		resp, err := client.Get("http://127.0.0.1")
+		r.NoError(err)
+
+		// The RejectResponseHandler should set our custom header
+		h := resp.Header.Get(testHeader)
+		if h == "" {
+			t.Errorf("Expecting header %s to be set by RejectResponseHandler", testHeader)
+		}
+		// Send a request that should be allowed
+		resp, err = client.Get("http://example.com")
+		r.NoError(err)
+
+		// The header set by our custom reject response handler should not be set
+		h = resp.Header.Get(testHeader)
+		if h != "" {
+			t.Errorf("Expecting header %s to not be set by RejectResponseHandler", testHeader)
+		}
+	})
+}
+
 // Test that Smokescreen calls the custom accept response handler (if defined in the Config struct)
 // after every accepted request
 func TestAcceptResponseHandler(t *testing.T) {