26 lines
551 B
Go
26 lines
551 B
Go
package middleware
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net/http"
|
|
"sync/atomic"
|
|
)
|
|
|
|
type ctxKeyRequestID int
|
|
|
|
const RequestIDKey ctxKeyRequestID = 0
|
|
|
|
func (m *Middleware) RequestID(next http.HandlerFunc) http.HandlerFunc {
|
|
return func(w http.ResponseWriter, r *http.Request) {
|
|
ctx := r.Context()
|
|
requestID := r.Header.Get(m.cfg.RequestIDHeader)
|
|
if requestID == "" {
|
|
myID := atomic.AddUint64(&m.reqID, 1)
|
|
requestID = fmt.Sprintf("%06d", myID)
|
|
}
|
|
ctx = context.WithValue(ctx, RequestIDKey, requestID)
|
|
r = r.WithContext(ctx)
|
|
next(w, r)
|
|
}
|
|
}
|