From ef63de54ea52e4a60aa9c5322db299476fc4c4de Mon Sep 17 00:00:00 2001
From: Hubert Grochowski <hubert.grochowski@saucelabs.com>
Date: Mon, 11 Sep 2023 11:29:32 +0200
Subject: [PATCH] martian: log request id from header

If specified and the header is present in the request, the proxy will associate the value with the request in the logs.
---
 internal/martian/handler.go |  2 +-
 internal/martian/proxy.go   | 20 +++++++++++++++++++-
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/internal/martian/handler.go b/internal/martian/handler.go
index d705165d..00da7095 100644
--- a/internal/martian/handler.go
+++ b/internal/martian/handler.go
@@ -89,7 +89,7 @@ func (p proxyHandler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	}
 	ctx := withSession(session)
 
-	outreq := req.Clone(ctx.addToContext(req.Context()))
+	outreq := req.Clone(p.requestContext(ctx, req))
 	if req.ContentLength == 0 {
 		outreq.Body = http.NoBody
 	}
diff --git a/internal/martian/proxy.go b/internal/martian/proxy.go
index 22dc6f5b..7d31685b 100644
--- a/internal/martian/proxy.go
+++ b/internal/martian/proxy.go
@@ -62,6 +62,11 @@ type Proxy struct {
 	// AllowHTTP disables automatic HTTP to HTTPS upgrades when the listener is TLS.
 	AllowHTTP bool
 
+	// RequestIDHeader specifies a special header name that the proxy will use to identify requests.
+	// If the header is present in the request, the proxy will associate the value with the request in the logs.
+	// If empty, no action is taken, and the proxy will generate a new request ID.
+	RequestIDHeader string
+
 	// ConnectPassthrough passes CONNECT requests to the RoundTripper,
 	// and uses the response body as the connection.
 	ConnectPassthrough bool
@@ -368,7 +373,7 @@ func (p *Proxy) readRequest(ctx *Context, conn net.Conn, brw *bufio.ReadWriter)
 		default:
 		}
 
-		req = req.WithContext(ctx.addToContext(req.Context()))
+		req = req.WithContext(p.requestContext(ctx, req))
 	}
 
 	// Adjust the read deadline if necessary.
@@ -381,6 +386,19 @@ func (p *Proxy) readRequest(ctx *Context, conn net.Conn, brw *bufio.ReadWriter)
 	return req, err
 }
 
+func (p *Proxy) requestContext(mctx *Context, req *http.Request) context.Context {
+	ctx := req.Context()
+	ctx = mctx.addToContext(ctx)
+
+	if h := p.RequestIDHeader; h != "" {
+		if id := req.Header.Get(h); id != "" {
+			ctx = context.WithValue(ctx, log.TraceContextKey, id)
+		}
+	}
+
+	return ctx
+}
+
 func (p *Proxy) handleConnectRequest(ctx *Context, req *http.Request, session *Session, brw *bufio.ReadWriter, conn net.Conn) error {
 	if err := p.reqmod.ModifyRequest(req); err != nil {
 		log.Errorf(req.Context(), "error modifying CONNECT request: %v", err)