package reqid
import (
"context"
"crypto/rand"
"encoding/hex"
"net/http"
)
type requestIDKey string
const (
RequestID requestIDKey = "request-id"
Header = "X-Request-ID"
)
// Middleware http middleware that sets random generated request id to each
// request it get fed.
func Middleware(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
rid := r.Header.Get(Header)
if rid == "" {
rid = generateRequestID()
r.Header.Set(Header, rid)
}
ctx := SetContext(r.Context(), rid)
r = r.WithContext(ctx)
w.Header().Add(Header, rid)
next.ServeHTTP(w, r)
})
}
// Get returns the request ID of http request (looks up from headers)
func Get(r *http.Request) string {
return r.Header.Get(Header)
}
// GetContext returns the request ID from context.
func GetContext(ctx context.Context) string {
rid, ok := ctx.Value(RequestID).(string)
if !ok {
return ""
}
return rid
}
// SetContext gets a parent context and returns a child context with the set
// provided request ID
func SetContext(ctx context.Context, reqID string) context.Context {
return context.WithValue(ctx, RequestID, reqID)
}
func generateRequestID() string {
b := make([]byte, 13)
_, err := rand.Read(b)
if err != nil {
return "unknown"
}
return hex.EncodeToString(b)
}
|