46 lines
1.0 KiB
Go
46 lines
1.0 KiB
Go
package middleware
|
||
|
||
import (
|
||
"crypto/rand"
|
||
"encoding/hex"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
)
|
||
|
||
// RequestIDHeader 请求 ID 的响应头名称
|
||
const RequestIDHeader = "X-Request-ID"
|
||
|
||
// RequestIDContextKey 上下文中存储请求 ID 的键
|
||
const RequestIDContextKey = "request_id"
|
||
|
||
// RequestID 为每个请求生成唯一 ID，写入上下文和响应头
|
||
func RequestID() gin.HandlerFunc {
|
||
return func(c *gin.Context) {
|
||
// 优先从请求头读取上游传入的 request_id
|
||
rid := c.GetHeader(RequestIDHeader)
|
||
if rid == "" {
|
||
rid = generateRequestID()
|
||
}
|
||
c.Set(RequestIDContextKey, rid)
|
||
c.Header(RequestIDHeader, rid)
|
||
c.Next()
|
||
}
|
||
}
|
||
|
||
// GetRequestID 从 gin.Context 中获取请求 ID
|
||
func GetRequestID(c *gin.Context) string {
|
||
if rid, exists := c.Get(RequestIDContextKey); exists {
|
||
if s, ok := rid.(string); ok {
|
||
return s
|
||
}
|
||
}
|
||
return ""
|
||
}
|
||
|
||
// generateRequestID 生成 16 字节随机 hex 请求 ID
|
||
func generateRequestID() string {
|
||
b := make([]byte, 16)
|
||
_, _ = rand.Read(b)
|
||
return hex.EncodeToString(b)
|
||
}
|