github.com/lastbackend/toolkit@v0.0.0-20241020043710-cafa37b95aad/examples/gateway/middleware/request.go (about)

     1  package middleware
     2  
     3  import (
     4  	"context"
     5  	"crypto/rand"
     6  	"encoding/base64"
     7  	"fmt"
     8  	"net/http"
     9  	"os"
    10  	"strings"
    11  	"sync/atomic"
    12  )
    13  
    14  // KeyPEMBlock to use when setting the request UserID.
    15  type ctxKeyRequestID int
    16  
    17  // RequestIDKey is the key that holds th unique request UserID in a request context.
    18  const RequestIDKey ctxKeyRequestID = 0
    19  
    20  var (
    21  	// prefix is const prefix for request UserID
    22  	prefix string
    23  
    24  	// reqID is counter for request UserID
    25  	reqID uint64
    26  )
    27  
    28  // init Initializes constant part of request UserID
    29  func init() {
    30  	hostname, err := os.Hostname()
    31  	if hostname == "" || err != nil {
    32  		hostname = "localhost"
    33  	}
    34  	var buf [12]byte
    35  	var b64 string
    36  	for len(b64) < 10 {
    37  		_, _ = rand.Read(buf[:])
    38  		b64 = base64.StdEncoding.EncodeToString(buf[:])
    39  		b64 = strings.NewReplacer("+", "", "/", "").Replace(b64)
    40  	}
    41  
    42  	prefix = fmt.Sprintf("%s/%s", hostname, b64[0:10])
    43  }
    44  
    45  // RequestID is a middleware that injects a request UserID into the context of each
    46  // request. A request UserID is a string of the form "host.example.com/random-0001",
    47  // where "random" is a base62 random string that uniquely identifies this go
    48  // process, and where the last number is an atomically incremented request
    49  // counter.
    50  func RequestID(h http.Handler) http.Handler {
    51  	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    52  		id := atomic.AddUint64(&reqID, 1)
    53  		ctx := r.Context()
    54  		ctx = context.WithValue(ctx, RequestIDKey, fmt.Sprintf("%s-%06d", prefix, id))
    55  		h.ServeHTTP(w, r.WithContext(ctx))
    56  	})
    57  }
    58  
    59  // GetReqID returns a request UserID from the given context if one is present.
    60  // Returns the empty string if a request UserID cannot be found.
    61  func GetReqID(ctx context.Context) string {
    62  	if ctx == nil {
    63  		return ""
    64  	}
    65  	if reqID, ok := ctx.Value(RequestIDKey).(string); ok {
    66  		return reqID
    67  	}
    68  	return ""
    69  }