smallstep-certificates/internal/requestid/requestid.go

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 != ""
}