From 9d9c46c63f82addf409670be73f329606e92bb52 Mon Sep 17 00:00:00 2001
From: David Stotijn <dstotijn@gmail.com>
Date: Sun, 24 Nov 2019 19:46:38 +0100
Subject: [PATCH] Add request and response modifier support

---
 main.go         |  9 ++++++++
 proxy/modify.go | 11 +++++++++
 proxy/proxy.go  | 60 ++++++++++++++++++++++++++++++++++---------------
 3 files changed, 62 insertions(+), 18 deletions(-)
 create mode 100644 proxy/modify.go

diff --git a/main.go b/main.go
index 7b804f2..50c3ffe 100644
--- a/main.go
+++ b/main.go
@@ -33,6 +33,15 @@ func main() {
 		log.Fatalf("[FATAL] Could not create Proxy: %v", err)
 	}
 
+	proxy.UseRequestModifier(func(req *http.Request) {
+		log.Printf("[DEBUG] Incoming request: %v", req.URL)
+	})
+
+	proxy.UseResponseModifier(func(res *http.Response) error {
+		log.Printf("[DEBUG] Downstream response: %v %v %v", res.Proto, res.StatusCode, http.StatusText(res.StatusCode))
+		return nil
+	})
+
 	s := &http.Server{
 		Addr:         ":8080",
 		Handler:      proxy,
diff --git a/proxy/modify.go b/proxy/modify.go
new file mode 100644
index 0000000..4364083
--- /dev/null
+++ b/proxy/modify.go
@@ -0,0 +1,11 @@
+package proxy
+
+import "net/http"
+
+var (
+	nopReqModifier = func(req *http.Request) {}
+	nopResModifier = func(res *http.Response) error { return nil }
+)
+
+type RequestModifyFunc func(req *http.Request)
+type ResponseModifyFunc func(res *http.Response) error
diff --git a/proxy/proxy.go b/proxy/proxy.go
index c24bac0..4597b6c 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -12,19 +12,6 @@ import (
 	"net/http/httputil"
 )
 
-var httpHandler = &httputil.ReverseProxy{
-	Director:     func(r *http.Request) {},
-	ErrorHandler: errorHandler,
-}
-
-var httpsHandler = &httputil.ReverseProxy{
-	Director: func(r *http.Request) {
-		r.URL.Host = r.Host
-		r.URL.Scheme = "https"
-	},
-	ErrorHandler: errorHandler,
-}
-
 func errorHandler(w http.ResponseWriter, r *http.Request, err error) {
 	if err == context.Canceled {
 		return
@@ -37,6 +24,11 @@ func errorHandler(w http.ResponseWriter, r *http.Request, err error) {
 // HTTP requests and responses.
 type Proxy struct {
 	certConfig *CertConfig
+	handler    http.Handler
+
+	// TODO: Add mutex for modifier funcs.
+	reqModifier RequestModifyFunc
+	resModifier ResponseModifyFunc
 }
 
 // NewProxy returns a new Proxy.
@@ -46,9 +38,19 @@ func NewProxy(ca *x509.Certificate, key crypto.PrivateKey) (*Proxy, error) {
 		return nil, err
 	}
 
-	return &Proxy{
-		certConfig: certConfig,
-	}, nil
+	p := &Proxy{
+		certConfig:  certConfig,
+		reqModifier: nopReqModifier,
+		resModifier: nopResModifier,
+	}
+
+	p.handler = &httputil.ReverseProxy{
+		Director:       p.modifyRequest,
+		ModifyResponse: p.modifyResponse,
+		ErrorHandler:   errorHandler,
+	}
+
+	return p, nil
 }
 
 func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
@@ -57,7 +59,21 @@ func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	httpHandler.ServeHTTP(w, r)
+	p.handler.ServeHTTP(w, r)
+}
+
+func (p *Proxy) modifyRequest(r *http.Request) {
+	// Fix r.URL for HTTPS requests after CONNECT.
+	if r.URL.Scheme == "" {
+		r.URL.Host = r.Host
+		r.URL.Scheme = "https"
+	}
+
+	p.reqModifier(r)
+}
+
+func (p *Proxy) modifyResponse(res *http.Response) error {
+	return p.resModifier(res)
 }
 
 // handleConnect hijacks the incoming HTTP request and sets up an HTTP tunnel.
@@ -91,7 +107,7 @@ func (p *Proxy) handleConnect(w http.ResponseWriter, r *http.Request) {
 
 	l := &OnceAcceptListener{clientConnNotify.Conn}
 
-	err = http.Serve(l, httpsHandler)
+	err = http.Serve(l, p.handler)
 	if err != nil && err != ErrAlreadyAccepted {
 		log.Printf("[ERROR] Serving HTTP request failed: %v", err)
 	}
@@ -110,6 +126,14 @@ func (p *Proxy) clientTLSConn(conn net.Conn) (*tls.Conn, error) {
 	return tlsConn, nil
 }
 
+func (p *Proxy) UseRequestModifier(fn RequestModifyFunc) {
+	p.reqModifier = fn
+}
+
+func (p *Proxy) UseResponseModifier(fn ResponseModifyFunc) {
+	p.resModifier = fn
+}
+
 func writeError(w http.ResponseWriter, r *http.Request, code int) {
 	http.Error(w, http.StatusText(code), code)
 }