github.com/searKing/golang/go@v1.2.117/net/http/interceptors.x-request-id.go (about)

     1  // Copyright 2022 The searKing Author. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package http
     6  
     7  import (
     8  	"context"
     9  	"net/http"
    10  )
    11  
    12  // DefaultXRequestIDKey is metadata key name for request ID
    13  var DefaultXRequestIDKey = "X-Request-ID"
    14  
    15  // SetInOutMetadata injects requestIDs in req|resp's Header and context
    16  // 将request-id追加注入请求|响应头及context中
    17  func SetInOutMetadata(ctx context.Context, w http.ResponseWriter, r *http.Request, requestID string) context.Context {
    18  	if r != nil {
    19  		r.Header.Set(DefaultXRequestIDKey, requestID)
    20  	}
    21  	if w != nil {
    22  		w.Header().Set(DefaultXRequestIDKey, requestID)
    23  	}
    24  	return context.WithValue(ctx, DefaultXRequestIDKey, requestID)
    25  }
    26  
    27  // FromHTTPContext parse request id from gin.Context
    28  // query | header | post form | context
    29  // 从请求中提取request-id
    30  func FromHTTPContext(r *http.Request, keys ...any) string {
    31  	key := DefaultXRequestIDKey
    32  	if requestID := r.Header.Get(key); requestID != "" {
    33  		return requestID
    34  	}
    35  	if requestID := r.URL.Query().Get(key); requestID != "" {
    36  		return requestID
    37  	}
    38  	if requestID := r.FormValue(key); requestID != "" {
    39  		return requestID
    40  	}
    41  	if requestID := r.PostFormValue(key); requestID != "" {
    42  		return requestID
    43  	}
    44  
    45  	switch requestIDs := r.Context().Value(key).(type) {
    46  	case string:
    47  		if requestIDs != "" {
    48  			return requestIDs
    49  		}
    50  	case []string:
    51  		if len(requestIDs) > 0 {
    52  			return requestIDs[0]
    53  		}
    54  	}
    55  
    56  	return fromContext(r.Context(), keys...)
    57  }
    58  
    59  // fromContext takes out first value from context by keys
    60  func fromContext(ctx context.Context, keys ...any) string {
    61  	for _, key := range keys {
    62  		val := ctx.Value(key)
    63  		switch val := val.(type) {
    64  		case string:
    65  			if val != "" {
    66  				return val
    67  			}
    68  		case []string:
    69  			if len(val) > 0 {
    70  				if val[0] != "" {
    71  					return val[0]
    72  				}
    73  			}
    74  		}
    75  	}
    76  	return ""
    77  }
    78  
    79  func RequestIDFromContext(ctx context.Context) string {
    80  	switch requestIDs := ctx.Value(DefaultXRequestIDKey).(type) {
    81  	case string:
    82  		return requestIDs
    83  	case []string:
    84  		if len(requestIDs) > 0 {
    85  			return requestIDs[0]
    86  		}
    87  	}
    88  	return ""
    89  }