gitlab.com/gitlab-org/labkit@v1.21.0/correlation/inbound_http.go (about)

     1  package correlation
     2  
     3  import (
     4  	"net/http"
     5  )
     6  
     7  // InjectCorrelationID is an HTTP middleware to generate an Correlation-ID for the incoming request,
     8  // or extract the existing Correlation-ID from the incoming request. By default, any upstream Correlation-ID,
     9  // passed in via the `X-Request-ID` header will be ignored. To enable this behaviour, the `WithPropagation`
    10  // option should be passed into the options.
    11  // Whether the Correlation-ID is generated or propagated, once inside this handler the request context
    12  // will have a Correlation-ID associated with it.
    13  func InjectCorrelationID(h http.Handler, opts ...InboundHandlerOption) http.Handler {
    14  	config := applyInboundHandlerOptions(opts)
    15  
    16  	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    17  		parent := r.Context()
    18  
    19  		correlationID := ""
    20  		clientName := ""
    21  		if config.shouldPropagate(r) {
    22  			correlationID = r.Header.Get(propagationHeader)
    23  			clientName = r.Header.Get(clientNameHeader)
    24  		}
    25  
    26  		if correlationID == "" {
    27  			correlationID = SafeRandomID()
    28  		}
    29  
    30  		ctx := ContextWithCorrelation(parent, correlationID)
    31  
    32  		if clientName != "" {
    33  			ctx = ContextWithClientName(ctx, clientName)
    34  		}
    35  		if config.sendResponseHeader {
    36  			// Set the response header.
    37  			w.Header().Set(propagationHeader, correlationID)
    38  		}
    39  		h.ServeHTTP(w, r.WithContext(ctx))
    40  	})
    41  }