x @ v0.2.0

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)
}