Skip to content

Commit

Permalink
martian: log request id from header
Browse files Browse the repository at this point in the history
If specified and the header is present in the request, the proxy will associate the value with the request in the logs.
  • Loading branch information
Choraden committed Sep 12, 2023
1 parent 52196e3 commit ef63de5
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 2 deletions.
2 changes: 1 addition & 1 deletion internal/martian/handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down
20 changes: 19 additions & 1 deletion internal/martian/proxy.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand All @@ -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)
Expand Down

0 comments on commit ef63de5

Please sign in to comment.