Add MITM proxy for HTTPS requests

proxy.go

@@ -1,79 +1,115 @@
 package main
 
 import (
-	"io"
+	"context"
+	"crypto"
+	"crypto/tls"
+	"crypto/x509"
+	"fmt"
 	"log"
 	"net"
 	"net/http"
 	"net/http/httputil"
 )
 
+var httpHandler = &httputil.ReverseProxy{
+	Director: func(r *http.Request) {
+		r.URL.Host = r.Host
+		r.URL.Scheme = "http"
+	},
+	ErrorHandler: proxyErrorHandler,
+}
+
+var httpsHandler = &httputil.ReverseProxy{
+	Director: func(r *http.Request) {
+		r.URL.Host = r.Host
+		r.URL.Scheme = "https"
+	},
+	ErrorHandler: proxyErrorHandler,
+}
+
+func proxyErrorHandler(w http.ResponseWriter, r *http.Request, err error) {
+	if err == context.Canceled {
+		return
+	}
+	log.Printf("[ERROR]: Proxy error: %v", err)
+	w.WriteHeader(http.StatusBadGateway)
+}
+
 // Proxy is used to forward HTTP requests.
 type Proxy struct {
-	rp httputil.ReverseProxy
+	certConfig *CertConfig
 }
 
 // NewProxy returns a new Proxy.
-func NewProxy() *Proxy {
-	return &Proxy{
-		rp: httputil.ReverseProxy{
-			Director: func(r *http.Request) {
-				log.Printf("Director handled URL: %v", r.URL)
-			},
-		},
+func NewProxy(ca *x509.Certificate, key crypto.PrivateKey) (*Proxy, error) {
+	certConfig, err := NewCertConfig(ca, key)
+	if err != nil {
+		return nil, err
 	}
+
+	return &Proxy{
+		certConfig: certConfig,
+	}, nil
 }
 
 func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	log.Printf("ServeHTTP: Received request (host: %v, url: %v", r.Host, r.URL)
-
 	if r.Method == http.MethodConnect {
 		p.handleConnect(w, r)
 		return
 	}
 
-	p.rp.ServeHTTP(w, r)
-	log.Printf("ServeHTTP: Finished (host: %v, url: %v", r.Host, r.URL)
+	httpHandler.ServeHTTP(w, r)
 }
 
+// handleConnect hijacks the incoming HTTP request and sets up an HTTP tunnel.
+// During the TLS handshake with the client, we use the proxy's CA config to
+// create a certificate on-the-fly.
 func (p *Proxy) handleConnect(w http.ResponseWriter, r *http.Request) {
 	hj, ok := w.(http.Hijacker)
 	if !ok {
-		log.Printf("handleConnect: ResponseWriter is not a http.Hijacker (type: %T)", w)
+		log.Printf("[ERROR] handleConnect: ResponseWriter is not a http.Hijacker (type: %T)", w)
 		writeError(w, r, http.StatusServiceUnavailable)
 		return
 	}
 
-	// destConn is the TCP connection to the destination web server of the
-	// proxied HTTP request.
-	destConn, err := net.Dial("tcp", r.Host)
-	if err != nil {
-		log.Printf("handleConnect: Connect to destination host failed: %v", err)
-		writeError(w, r, http.StatusBadGateway)
-		return
-	}
-	defer destConn.Close()
-
 	w.WriteHeader(http.StatusOK)
 
-	// clientConn is the TCP connection to the client.
 	clientConn, _, err := hj.Hijack()
 	if err != nil {
-		log.Printf("handleConnect: Hijack failed: %v", err)
+		log.Printf("[ERROR] Hijacking client connection failed: %v", err)
 		writeError(w, r, http.StatusServiceUnavailable)
 		return
 	}
 	defer clientConn.Close()
 
-	errc := make(chan error, 1)
-	go tunnelData(destConn, clientConn, errc)
-	go tunnelData(clientConn, destConn, errc)
-	<-errc
+	// Secure connection to client.
+	clientConn, err = p.clientTLSConn(clientConn)
+	if err != nil {
+		log.Printf("[ERROR] Securing client connection failed: %v", err)
+		return
+	}
+	clientConnNotify := ConnNotify{clientConn, make(chan struct{})}
+
+	l := &OnceAcceptListener{clientConnNotify.Conn}
+
+	err = http.Serve(l, httpsHandler)
+	if err != nil && err != ErrAlreadyAccepted {
+		log.Printf("[ERROR] Serving HTTP request failed: %v", err)
+	}
+	<-clientConnNotify.closed
 }
 
-func tunnelData(dst, src io.ReadWriter, errc chan<- error) {
-	_, err := io.Copy(dst, src)
-	errc <- err
+func (p *Proxy) clientTLSConn(conn net.Conn) (*tls.Conn, error) {
+	tlsConfig := p.certConfig.TLSConfig()
+
+	tlsConn := tls.Server(conn, tlsConfig)
+	if err := tlsConn.Handshake(); err != nil {
+		tlsConn.Close()
+		return nil, fmt.Errorf("handshake error: %v", err)
+	}
+
+	return tlsConn, nil
 }
 
 func writeError(w http.ResponseWriter, r *http.Request, code int) {