x @ da217b4426264b3ac4861ce9411bcff61cb3426b

···
        5
        5
         

      
        6
        6
         - is - simple assertion library for tests

      
        7
        7
         - envy - type-safe os.Getenv with default values

      
        
        8
        +- reqid - http middleware that add request id header and sets it in context

      
        8
        9
         

      
        9
        10
         This repo is MIT Licensed

      

···
        
        1
        +package reqid

      
        
        2
        +

      
        
        3
        +import (

      
        
        4
        +	"context"

      
        
        5
        +	"crypto/rand"

      
        
        6
        +	"encoding/hex"

      
        
        7
        +	"net/http"

      
        
        8
        +)

      
        
        9
        +

      
        
        10
        +type requestIDKey string

      
        
        11
        +

      
        
        12
        +const (

      
        
        13
        +	RequestID requestIDKey = "request-id"

      
        
        14
        +

      
        
        15
        +	Header = "X-Request-ID"

      
        
        16
        +)

      
        
        17
        +

      
        
        18
        +// Middleware http middleware that sets random generated request id to each

      
        
        19
        +// request it get fed.

      
        
        20
        +func Middleware(next http.Handler) http.Handler {

      
        
        21
        +	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

      
        
        22
        +		rid := r.Header.Get(Header)

      
        
        23
        +		if rid == "" {

      
        
        24
        +			rid = generateRequestID()

      
        
        25
        +			r.Header.Set(Header, rid)

      
        
        26
        +		}

      
        
        27
        +

      
        
        28
        +		ctx := SetContext(r.Context(), rid)

      
        
        29
        +		r = r.WithContext(ctx)

      
        
        30
        +

      
        
        31
        +		w.Header().Add(Header, rid)

      
        
        32
        +		next.ServeHTTP(w, r)

      
        
        33
        +	})

      
        
        34
        +}

      
        
        35
        +

      
        
        36
        +// Get returns the request ID of http request (looks up from headers)

      
        
        37
        +func Get(r *http.Request) string {

      
        
        38
        +	return r.Header.Get(Header)

      
        
        39
        +}

      
        
        40
        +

      
        
        41
        +// GetContext returns the request ID from context.

      
        
        42
        +func GetContext(ctx context.Context) string {

      
        
        43
        +	rid, ok := ctx.Value(RequestID).(string)

      
        
        44
        +	if !ok {

      
        
        45
        +		return ""

      
        
        46
        +	}

      
        
        47
        +	return rid

      
        
        48
        +}

      
        
        49
        +

      
        
        50
        +// SetContext gets a parent context and returns a child context with the set

      
        
        51
        +// provided request ID

      
        
        52
        +func SetContext(ctx context.Context, reqID string) context.Context {

      
        
        53
        +	return context.WithValue(ctx, RequestID, reqID)

      
        
        54
        +}

      
        
        55
        +

      
        
        56
        +func generateRequestID() string {

      
        
        57
        +	b := make([]byte, 13)

      
        
        58
        +	_, err := rand.Read(b)

      
        
        59
        +	if err != nil {

      
        
        60
        +		return "unknown"

      
        
        61
        +	}

      
        
        62
        +	return hex.EncodeToString(b)

      
        
        63
        +}

      

···
        
        1
        +package reqid

      
        
        2
        +

      
        
        3
        +import (

      
        
        4
        +	"net/http"

      
        
        5
        +	"net/http/httptest"

      
        
        6
        +	"testing"

      
        
        7
        +

      
        
        8
        +	"olexsmir.xyz/x/is"

      
        
        9
        +)

      
        
        10
        +

      
        
        11
        +func TestMiddleware(t *testing.T) {

      
        
        12
        +	mux := http.NewServeMux()

      
        
        13
        +	mux.HandleFunc("GET /", testHandler)

      
        
        14
        +	hand := Middleware(mux)

      
        
        15
        +

      
        
        16
        +	w := httptest.NewRecorder()

      
        
        17
        +	req, _ := http.NewRequest(http.MethodGet, "/", nil)

      
        
        18
        +	hand.ServeHTTP(w, req)

      
        
        19
        +

      
        
        20
        +	is.Equal(t, http.StatusOK, w.Code)

      
        
        21
        +	is.NotEqual(t, w.Header().Get(Header), "")

      
        
        22
        +}

      
        
        23
        +

      
        
        24
        +func BenchmarkMiddleware(b *testing.B) {

      
        
        25
        +	mux := http.NewServeMux()

      
        
        26
        +	mux.HandleFunc("GET /", testHandler)

      
        
        27
        +	hand := Middleware(mux)

      
        
        28
        +

      
        
        29
        +	w := httptest.NewRecorder()

      
        
        30
        +	req, _ := http.NewRequest(http.MethodGet, "/", nil)

      
        
        31
        +

      
        
        32
        +	for b.Loop() {

      
        
        33
        +		hand.ServeHTTP(w, req)

      
        
        34
        +	}

      
        
        35
        +}

      
        
        36
        +

      
        
        37
        +func testHandler(w http.ResponseWriter, _ *http.Request) {

      
        
        38
        +	w.WriteHeader(http.StatusOK)

      
        
        39
        +}