57 lines
1.3 KiB
Go
57 lines
1.3 KiB
Go
|
package middleware
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"crypto/rand"
|
|||
|
|
"encoding/hex"
|
|||
|
|
"fmt"
|
|||
|
|
"time"
|
|||
|
|
|
|||
|
|
"github.com/gin-gonic/gin"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
const (
|
|||
|
|
// TraceIDHeader 追踪 ID 的 HTTP 响应头名称
|
|||
|
|
TraceIDHeader = "X-Trace-ID"
|
|||
|
|
// TraceIDKey gin.Context 中的 key
|
|||
|
|
TraceIDKey = "trace_id"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
// TraceID 中间件：为每个请求生成唯一追踪 ID
|
|||
|
|
// 追踪 ID 写入 gin.Context 和响应头，供日志和下游服务关联
|
|||
|
|
func TraceID() gin.HandlerFunc {
|
|||
|
|
return func(c *gin.Context) {
|
|||
|
|
// 优先复用上游传入的 Trace ID（如 API 网关、前端）
|
|||
|
|
traceID := c.GetHeader(TraceIDHeader)
|
|||
|
|
if traceID == "" {
|
|||
|
|
traceID = generateTraceID()
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
c.Set(TraceIDKey, traceID)
|
|||
|
|
c.Header(TraceIDHeader, traceID)
|
|||
|
|
|
|||
|
|
c.Next()
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// generateTraceID 生成 16 字节随机 hex 字符串，格式：时间前缀+随机后缀
|
|||
|
|
// 例：20260405-a1b2c3d4e5f60718
|
|||
|
|
func generateTraceID() string {
|
|||
|
|
b := make([]byte, 8)
|
|||
|
|
_, err := rand.Read(b)
|
|||
|
|
if err != nil {
|
|||
|
|
// 降级：使用时间戳
|
|||
|
|
return fmt.Sprintf("%d", time.Now().UnixNano())
|
|||
|
|
}
|
|||
|
|
return fmt.Sprintf("%s-%s", time.Now().Format("20060102"), hex.EncodeToString(b))
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// GetTraceID 从 gin.Context 获取 trace ID（供 handler 使用）
|
|||
|
|
func GetTraceID(c *gin.Context) string {
|
|||
|
|
if v, exists := c.Get(TraceIDKey); exists {
|
|||
|
|
if id, ok := v.(string); ok {
|
|||
|
|
return id
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return ""
|
|||
|
|
}
|