64 lines
1.2 KiB
Go
64 lines
1.2 KiB
Go
|
package requestid
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"fmt"
|
||
|
|
"net/http"
|
||
|
|
|
||
|
|
"github.com/google/uuid"
|
||
|
|
)
|
||
|
|
|
||
|
|
const defaultHeaderName = "X-Request-ID"
|
||
|
|
|
||
|
|
// Config plugin configuration
|
||
|
|
type Config struct {
|
||
|
|
HeaderName string `json:"headerName"`
|
||
|
|
}
|
||
|
|
|
||
|
|
// CreateConfig create default plugin configuration
|
||
|
|
func CreateConfig() *Config {
|
||
|
|
return &Config{
|
||
|
|
HeaderName: defaultHeaderName,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// RequestIDHeader
|
||
|
|
type RequestIDHeader struct {
|
||
|
|
headerName string
|
||
|
|
name string
|
||
|
|
next http.Handler
|
||
|
|
}
|
||
|
|
|
||
|
|
// New create new RequestIDHeader
|
||
|
|
func New(ctx context.Context, next http.Handler, config *Config, name string) (http.Handler, error) {
|
||
|
|
hdr := &RequestIDHeader{
|
||
|
|
next: next,
|
||
|
|
name: name,
|
||
|
|
}
|
||
|
|
|
||
|
|
if config == nil {
|
||
|
|
return nil, fmt.Errorf("config can not be nil")
|
||
|
|
}
|
||
|
|
|
||
|
|
if config.HeaderName == "" {
|
||
|
|
hdr.headerName = defaultHeaderName
|
||
|
|
} else {
|
||
|
|
hdr.headerName = config.HeaderName
|
||
|
|
}
|
||
|
|
|
||
|
|
return hdr, nil
|
||
|
|
|
||
|
|
}
|
||
|
|
|
||
|
|
func (r *RequestIDHeader) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
|
||
|
|
headerArr := req.Header[r.headerName]
|
||
|
|
uuid := uuid.New().String()
|
||
|
|
if len(headerArr) == 0 {
|
||
|
|
req.Header.Add(r.headerName, uuid)
|
||
|
|
} else if headerArr[0] == "" {
|
||
|
|
req.Header[r.headerName][0] = uuid
|
||
|
|
}
|
||
|
|
|
||
|
|
r.next.ServeHTTP(rw, req)
|
||
|
|
}
|