151 lines
3.6 KiB
Go
151 lines
3.6 KiB
Go
package logging
|
|
|
|
import (
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
// CPATraceIDHeader is the downstream response header used to correlate requests with selected credentials.
|
|
const CPATraceIDHeader = "X-CPA-TRACE-ID"
|
|
|
|
const ginCPATraceStateKey = "__cpa_trace_state__"
|
|
|
|
// FormatCPATraceID builds a CPA trace ID from the selection time, auth index, and request ID.
|
|
func FormatCPATraceID(selectedAt time.Time, authIndex, requestID string) string {
|
|
authIndex = strings.TrimSpace(authIndex)
|
|
requestID = strings.TrimSpace(requestID)
|
|
if selectedAt.IsZero() || authIndex == "" || requestID == "" {
|
|
return ""
|
|
}
|
|
return selectedAt.Format("20060102150405") + "-" + authIndex + "-" + requestID
|
|
}
|
|
|
|
type cpaTraceState struct {
|
|
mu sync.RWMutex
|
|
traceID string
|
|
}
|
|
|
|
func (s *cpaTraceState) set(traceID string) {
|
|
if s == nil {
|
|
return
|
|
}
|
|
s.mu.Lock()
|
|
s.traceID = strings.TrimSpace(traceID)
|
|
s.mu.Unlock()
|
|
}
|
|
|
|
func (s *cpaTraceState) get() string {
|
|
if s == nil {
|
|
return ""
|
|
}
|
|
s.mu.RLock()
|
|
traceID := s.traceID
|
|
s.mu.RUnlock()
|
|
return traceID
|
|
}
|
|
|
|
func ginCPATraceState(c *gin.Context) *cpaTraceState {
|
|
if c == nil {
|
|
return nil
|
|
}
|
|
if value, exists := c.Get(ginCPATraceStateKey); exists {
|
|
if state, ok := value.(*cpaTraceState); ok && state != nil {
|
|
return state
|
|
}
|
|
}
|
|
state := &cpaTraceState{}
|
|
c.Set(ginCPATraceStateKey, state)
|
|
return state
|
|
}
|
|
|
|
// GinCPATraceIDCallback returns a callback that is safe to invoke after the Gin context is released.
|
|
func GinCPATraceIDCallback(c *gin.Context) func(string) {
|
|
state := ginCPATraceState(c)
|
|
if state == nil {
|
|
return nil
|
|
}
|
|
requestID := GetGinRequestID(c)
|
|
if requestID == "" && c.Request != nil {
|
|
requestID = GetRequestID(c.Request.Context())
|
|
}
|
|
requestID = strings.TrimSpace(requestID)
|
|
if requestID == "" {
|
|
return nil
|
|
}
|
|
return func(authIndex string) {
|
|
if traceID := FormatCPATraceID(time.Now(), authIndex, requestID); traceID != "" {
|
|
state.set(traceID)
|
|
}
|
|
}
|
|
}
|
|
|
|
// SetGinCPATraceID stores the trace ID until the downstream response headers are committed.
|
|
func SetGinCPATraceID(c *gin.Context, authIndex string) {
|
|
if callback := GinCPATraceIDCallback(c); callback != nil {
|
|
callback(authIndex)
|
|
}
|
|
}
|
|
|
|
// GetGinCPATraceID returns the trace ID stored for the current request.
|
|
func GetGinCPATraceID(c *gin.Context) string {
|
|
if c == nil {
|
|
return ""
|
|
}
|
|
value, exists := c.Get(ginCPATraceStateKey)
|
|
if !exists {
|
|
return ""
|
|
}
|
|
state, _ := value.(*cpaTraceState)
|
|
return state.get()
|
|
}
|
|
|
|
// CPATraceIDMiddleware injects a stored trace ID immediately before response headers are committed.
|
|
func CPATraceIDMiddleware() gin.HandlerFunc {
|
|
return func(c *gin.Context) {
|
|
state := ginCPATraceState(c)
|
|
c.Writer = &cpaTraceResponseWriter{ResponseWriter: c.Writer, state: state}
|
|
c.Next()
|
|
}
|
|
}
|
|
|
|
type cpaTraceResponseWriter struct {
|
|
gin.ResponseWriter
|
|
state *cpaTraceState
|
|
}
|
|
|
|
func (w *cpaTraceResponseWriter) WriteHeader(statusCode int) {
|
|
w.applyTraceHeader()
|
|
w.ResponseWriter.WriteHeader(statusCode)
|
|
}
|
|
|
|
func (w *cpaTraceResponseWriter) WriteHeaderNow() {
|
|
w.applyTraceHeader()
|
|
w.ResponseWriter.WriteHeaderNow()
|
|
}
|
|
|
|
func (w *cpaTraceResponseWriter) Write(data []byte) (int, error) {
|
|
w.applyTraceHeader()
|
|
return w.ResponseWriter.Write(data)
|
|
}
|
|
|
|
func (w *cpaTraceResponseWriter) WriteString(data string) (int, error) {
|
|
w.applyTraceHeader()
|
|
return w.ResponseWriter.WriteString(data)
|
|
}
|
|
|
|
func (w *cpaTraceResponseWriter) Flush() {
|
|
w.applyTraceHeader()
|
|
w.ResponseWriter.Flush()
|
|
}
|
|
|
|
func (w *cpaTraceResponseWriter) applyTraceHeader() {
|
|
if w == nil || w.ResponseWriter == nil || w.ResponseWriter.Written() {
|
|
return
|
|
}
|
|
if traceID := w.state.get(); traceID != "" {
|
|
w.ResponseWriter.Header().Set(CPATraceIDHeader, traceID)
|
|
}
|
|
}
|