Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Decouple request ID middleware from logging middleware
- Loading branch information
Showing
9 changed files
with
155 additions
and
124 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,82 @@ | ||
package requestid | ||
|
||
import ( | ||
"context" | ||
"net/http" | ||
|
||
"github.com/rs/xid" | ||
) | ||
|
||
const ( | ||
// requestIDHeader is the header name used for propagating request IDs. If | ||
// available in an HTTP request, it'll be used instead of the X-Smallstep-Id | ||
// header. It'll always be used in response and set to the request ID. | ||
requestIDHeader = "X-Request-Id" | ||
|
||
// defaultTraceHeader is the default Smallstep tracing header that's currently | ||
// in use. It is used as a fallback to retrieve a request ID from, if the | ||
// "X-Request-Id" request header is not set. | ||
defaultTraceHeader = "X-Smallstep-Id" | ||
) | ||
|
||
type Handler struct { | ||
legacyTraceHeader string | ||
} | ||
|
||
// New creates a new request ID [handler]. It takes a trace header, | ||
// which is used keep the legacy behavior intact, which relies on the | ||
// X-Smallstep-Id header instead of X-Request-Id. | ||
func New(legacyTraceHeader string) *Handler { | ||
if legacyTraceHeader == "" { | ||
legacyTraceHeader = defaultTraceHeader | ||
} | ||
|
||
return &Handler{legacyTraceHeader: legacyTraceHeader} | ||
} | ||
|
||
// Middleware wraps an [http.Handler] with request ID extraction | ||
// from the X-Reqeust-Id header by default, or from the X-Smallstep-Id | ||
// header if not set. If both are not set, a new request ID is generated. | ||
// In all cases, the request ID is added to the request context, and | ||
// set to be reflected in the response. | ||
func (h *Handler) Middleware(next http.Handler) http.Handler { | ||
fn := func(w http.ResponseWriter, req *http.Request) { | ||
requestID := req.Header.Get(requestIDHeader) | ||
if requestID == "" { | ||
requestID = req.Header.Get(h.legacyTraceHeader) | ||
} | ||
|
||
if requestID == "" { | ||
requestID = newRequestID() | ||
req.Header.Set(h.legacyTraceHeader, requestID) // legacy behavior | ||
} | ||
|
||
// immediately set the request ID to be reflected in the response | ||
w.Header().Set(requestIDHeader, requestID) | ||
|
||
// continue down the handler chain | ||
ctx := NewContext(req.Context(), requestID) | ||
next.ServeHTTP(w, req.WithContext(ctx)) | ||
} | ||
return http.HandlerFunc(fn) | ||
} | ||
|
||
// newRequestID creates a new request ID using github.com/rs/xid. | ||
func newRequestID() string { | ||
return xid.New().String() | ||
} | ||
|
||
type requestIDKey struct{} | ||
|
||
// NewContext returns a new context with the given request ID added to the | ||
// context. | ||
func NewContext(ctx context.Context, requestID string) context.Context { | ||
return context.WithValue(ctx, requestIDKey{}, requestID) | ||
} | ||
|
||
// FromContext returns the request ID from the context if it exists and | ||
// is not the empty value. | ||
func FromContext(ctx context.Context) (string, bool) { | ||
v, ok := ctx.Value(requestIDKey{}).(string) | ||
return v, ok && v != "" | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.