github.com/suchongming/fabric@v2.1.1+incompatible/core/middleware/request_id.go (about)

     1  /*
     2  Copyright IBM Corp. All Rights Reserved.
     3  
     4  SPDX-License-Identifier: Apache-2.0
     5  */
     6  
     7  package middleware
     8  
     9  import (
    10  	"context"
    11  	"net/http"
    12  )
    13  
    14  var requestIDKey = requestIDKeyType{}
    15  
    16  type requestIDKeyType struct{}
    17  
    18  func RequestID(ctx context.Context) string {
    19  	if reqID, ok := ctx.Value(requestIDKey).(string); ok {
    20  		return reqID
    21  	}
    22  	return "unknown"
    23  }
    24  
    25  type GenerateIDFunc func() string
    26  
    27  type requestID struct {
    28  	generateID GenerateIDFunc
    29  	next       http.Handler
    30  }
    31  
    32  func WithRequestID(generator GenerateIDFunc) Middleware {
    33  	return func(next http.Handler) http.Handler {
    34  		return &requestID{next: next, generateID: generator}
    35  	}
    36  }
    37  
    38  func (r *requestID) ServeHTTP(w http.ResponseWriter, req *http.Request) {
    39  	reqID := req.Header.Get("X-Request-Id")
    40  	if reqID == "" {
    41  		reqID = r.generateID()
    42  		req.Header.Set("X-Request-Id", reqID)
    43  	}
    44  
    45  	ctx := context.WithValue(req.Context(), requestIDKey, reqID)
    46  	req = req.WithContext(ctx)
    47  
    48  	w.Header().Add("X-Request-Id", reqID)
    49  
    50  	r.next.ServeHTTP(w, req)
    51  }