144 lines
3.4 KiB
Go
144 lines
3.4 KiB
Go
package logging
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"sync"
|
|
"sync/atomic"
|
|
)
|
|
|
|
type endpointKey struct{}
|
|
type responseStatusKey struct{}
|
|
type responseHeadersKey struct{}
|
|
type clientRequestMetadataKey struct{}
|
|
|
|
// ClientRequestMetadata stores immutable downstream request metadata for asynchronous consumers.
|
|
type ClientRequestMetadata struct {
|
|
ClientIP string
|
|
XForwardedFor string
|
|
UserAgent string
|
|
}
|
|
|
|
type responseStatusHolder struct {
|
|
status atomic.Int32
|
|
}
|
|
|
|
type responseHeadersHolder struct {
|
|
mu sync.RWMutex
|
|
headers http.Header
|
|
}
|
|
|
|
func WithEndpoint(ctx context.Context, endpoint string) context.Context {
|
|
if ctx == nil {
|
|
ctx = context.Background()
|
|
}
|
|
return context.WithValue(ctx, endpointKey{}, endpoint)
|
|
}
|
|
|
|
func GetEndpoint(ctx context.Context) string {
|
|
if ctx == nil {
|
|
return ""
|
|
}
|
|
if endpoint, ok := ctx.Value(endpointKey{}).(string); ok {
|
|
return endpoint
|
|
}
|
|
return ""
|
|
}
|
|
|
|
// WithClientRequestMetadata stores a snapshot of downstream request metadata in ctx.
|
|
func WithClientRequestMetadata(ctx context.Context, metadata ClientRequestMetadata) context.Context {
|
|
if ctx == nil {
|
|
ctx = context.Background()
|
|
}
|
|
return context.WithValue(ctx, clientRequestMetadataKey{}, metadata)
|
|
}
|
|
|
|
// GetClientRequestMetadata returns downstream request metadata stored in ctx.
|
|
func GetClientRequestMetadata(ctx context.Context) ClientRequestMetadata {
|
|
if ctx == nil {
|
|
return ClientRequestMetadata{}
|
|
}
|
|
if metadata, ok := ctx.Value(clientRequestMetadataKey{}).(ClientRequestMetadata); ok {
|
|
return metadata
|
|
}
|
|
return ClientRequestMetadata{}
|
|
}
|
|
|
|
func WithResponseStatusHolder(ctx context.Context) context.Context {
|
|
if ctx == nil {
|
|
ctx = context.Background()
|
|
}
|
|
if holder, ok := ctx.Value(responseStatusKey{}).(*responseStatusHolder); ok && holder != nil {
|
|
return ctx
|
|
}
|
|
return context.WithValue(ctx, responseStatusKey{}, &responseStatusHolder{})
|
|
}
|
|
|
|
func WithResponseHeadersHolder(ctx context.Context) context.Context {
|
|
if ctx == nil {
|
|
ctx = context.Background()
|
|
}
|
|
if holder, ok := ctx.Value(responseHeadersKey{}).(*responseHeadersHolder); ok && holder != nil {
|
|
return ctx
|
|
}
|
|
return context.WithValue(ctx, responseHeadersKey{}, &responseHeadersHolder{})
|
|
}
|
|
|
|
func SetResponseStatus(ctx context.Context, status int) {
|
|
if ctx == nil || status <= 0 {
|
|
return
|
|
}
|
|
holder, ok := ctx.Value(responseStatusKey{}).(*responseStatusHolder)
|
|
if !ok || holder == nil {
|
|
return
|
|
}
|
|
holder.status.Store(int32(status))
|
|
}
|
|
|
|
func SetResponseHeaders(ctx context.Context, headers http.Header) {
|
|
if ctx == nil {
|
|
return
|
|
}
|
|
holder, ok := ctx.Value(responseHeadersKey{}).(*responseHeadersHolder)
|
|
if !ok || holder == nil {
|
|
return
|
|
}
|
|
holder.mu.Lock()
|
|
defer holder.mu.Unlock()
|
|
holder.headers = cloneHTTPHeader(headers)
|
|
}
|
|
|
|
func GetResponseStatus(ctx context.Context) int {
|
|
if ctx == nil {
|
|
return 0
|
|
}
|
|
holder, ok := ctx.Value(responseStatusKey{}).(*responseStatusHolder)
|
|
if !ok || holder == nil {
|
|
return 0
|
|
}
|
|
return int(holder.status.Load())
|
|
}
|
|
|
|
func GetResponseHeaders(ctx context.Context) http.Header {
|
|
if ctx == nil {
|
|
return nil
|
|
}
|
|
holder, ok := ctx.Value(responseHeadersKey{}).(*responseHeadersHolder)
|
|
if !ok || holder == nil {
|
|
return nil
|
|
}
|
|
holder.mu.RLock()
|
|
defer holder.mu.RUnlock()
|
|
return cloneHTTPHeader(holder.headers)
|
|
}
|
|
|
|
func cloneHTTPHeader(src http.Header) http.Header {
|
|
if len(src) == 0 {
|
|
return nil
|
|
}
|
|
dst := make(http.Header, len(src))
|
|
for key, values := range src {
|
|
dst[key] = append([]string(nil), values...)
|
|
}
|
|
return dst
|
|
}
|