git.sr.ht/~pingoo/stdx@v0.0.0-20240218134121-094174641f6e/httpx/middlewarex/request_id.go (about)

     1  package middlewarex
     2  
     3  import (
     4  	"context"
     5  	"net/http"
     6  
     7  	"git.sr.ht/~pingoo/stdx/guid"
     8  )
     9  
    10  type requestIDContextKey struct{}
    11  
    12  // RequestIDCtxKey is the key that holds the unique request ID in a request context.
    13  var RequestIDCtxKey = requestIDContextKey{}
    14  
    15  func RequestID(header string) func(next http.Handler) http.Handler {
    16  	isHeaderEmpty := header == ""
    17  	return func(next http.Handler) http.Handler {
    18  		fn := func(w http.ResponseWriter, r *http.Request) {
    19  			requestID := guid.NewTimeBased()
    20  
    21  			if !isHeaderEmpty {
    22  				w.Header().Set(header, requestID.String())
    23  			}
    24  
    25  			ctx := r.Context()
    26  			ctx = context.WithValue(ctx, RequestIDCtxKey, requestID)
    27  
    28  			next.ServeHTTP(w, r.WithContext(ctx))
    29  		}
    30  		return http.HandlerFunc(fn)
    31  	}
    32  }