Add intercept module

pkg/proxy/proxy.go

@@ -7,16 +7,24 @@ import (
 	"crypto/x509"
 	"errors"
 	"fmt"
+	"math/rand"
 	"net"
 	"net/http"
 	"net/http/httputil"
+	"strings"
+	"time"
+
+	"github.com/oklog/ulid"
 
 	"github.com/dstotijn/hetty/pkg/log"
 )
 
+//nolint:gosec
+var ulidEntropy = rand.New(rand.NewSource(time.Now().UnixNano()))
+
 type contextKey int
 
-const ReqLogIDKey contextKey = 0
+const reqIDKey contextKey = 0
 
 // Proxy implements http.Handler and offers MITM behaviour for modifying
 // HTTP requests and responses.
@@ -54,7 +62,25 @@ func NewProxy(cfg Config) (*Proxy, error) {
 		p.logger = log.NewNopLogger()
 	}
 
+	transport := &http.Transport{
+		// Values taken from `http.DefaultTransport`.
+		Proxy: http.ProxyFromEnvironment,
+		DialContext: (&net.Dialer{
+			Timeout:   30 * time.Second,
+			KeepAlive: 30 * time.Second,
+		}).DialContext,
+		ForceAttemptHTTP2:     true,
+		MaxIdleConns:          100,
+		IdleConnTimeout:       90 * time.Second,
+		TLSHandshakeTimeout:   10 * time.Second,
+		ExpectContinueTimeout: 1 * time.Second,
+
+		// Non-default transport values.
+		DisableCompression: true,
+	}
+
 	p.handler = &httputil.ReverseProxy{
+		Transport:      transport,
 		Director:       p.modifyRequest,
 		ModifyResponse: p.modifyResponse,
 		ErrorHandler:   p.errorHandler,
@@ -69,6 +95,10 @@ func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	reqID := ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy)
+	ctx := context.WithValue(r.Context(), reqIDKey, reqID)
+	*r = *r.WithContext(ctx)
+
 	p.handler.ServeHTTP(w, r)
 }
 
@@ -91,6 +121,25 @@ func (p *Proxy) modifyRequest(r *http.Request) {
 	// set this header.
 	r.Header["X-Forwarded-For"] = nil
 
+	// Strip unsupported encodings.
+	if acceptEncs := r.Header.Get("Accept-Encoding"); acceptEncs != "" {
+		directives := strings.Split(acceptEncs, ",")
+		updated := make([]string, 0, len(directives))
+
+		for _, directive := range directives {
+			stripped := strings.TrimSpace(directive)
+			if strings.HasPrefix(stripped, "*") || strings.HasPrefix(stripped, "gzip") {
+				updated = append(updated, stripped)
+			}
+		}
+
+		if len(updated) == 0 {
+			r.Header.Del("Accept-Encoding")
+		} else {
+			r.Header.Set("Accept-Encoding", strings.Join(updated, ", "))
+		}
+	}
+
 	fn := nopReqModifier
 
 	for i := len(p.reqModifiers) - 1; i >= 0; i-- {
@@ -103,6 +152,11 @@ func (p *Proxy) modifyRequest(r *http.Request) {
 func (p *Proxy) modifyResponse(res *http.Response) error {
 	fn := nopResModifier
 
+	// TODO: Make decompressing gzip formatted response bodies a configurable project setting.
+	if err := gunzipResponseBody(res); err != nil {
+		return fmt.Errorf("proxy: failed to gunzip response body: %w", err)
+	}
+
 	for i := len(p.resModifiers) - 1; i >= 0; i-- {
 		fn = p.resModifiers[i](fn)
 	}
@@ -110,6 +164,15 @@ func (p *Proxy) modifyResponse(res *http.Response) error {
 	return fn(res)
 }
 
+func WithRequestID(ctx context.Context, id ulid.ULID) context.Context {
+	return context.WithValue(ctx, reqIDKey, id)
+}
+
+func RequestIDFromContext(ctx context.Context) (ulid.ULID, bool) {
+	id, ok := ctx.Value(reqIDKey).(ulid.ULID)
+	return id, ok
+}
+
 // 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.
@@ -170,13 +233,14 @@ func (p *Proxy) clientTLSConn(conn net.Conn) (*tls.Conn, error) {
 }
 
 func (p *Proxy) errorHandler(w http.ResponseWriter, r *http.Request, err error) {
-	if errors.Is(err, context.Canceled) {
-		return
+	switch {
+	case !errors.Is(err, context.Canceled):
+		p.logger.Errorw("Failed to proxy request.",
+			"error", err)
+	case errors.Is(err, context.Canceled):
+		p.logger.Debugw("Proxy request was cancelled.")
 	}
 
-	p.logger.Errorw("Failed to proxy request.",
-		"error", err)
-
 	w.WriteHeader(http.StatusBadGateway)
 }