1 Star 0 Fork 0

yzsunjianguo/common_pkg

加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
文件
克隆/下载
requstid.go 4.71 KB
一键复制 编辑 原始数据 按行查看 历史
yzsunjianguo 提交于 2024-02-08 11:08 . fix
package middleware
import (
"context"
"net/http"
"gitee.com/yzsunjianguo/common_pkg/krand"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
)
var (
// ContextRequestIDKey request id for context
ContextRequestIDKey = "request_id"
// HeaderXRequestIDKey header request id key
HeaderXRequestIDKey = "X-Request-Id"
)
// RequestIDOption set the request id options.
type RequestIDOption func(*requestIDOptions)
type requestIDOptions struct {
contextRequestIDKey string
headerXRequestIDKey string
}
func defaultRequestIDOptions() *requestIDOptions {
return &requestIDOptions{
contextRequestIDKey: ContextRequestIDKey,
headerXRequestIDKey: HeaderXRequestIDKey,
}
}
func (o *requestIDOptions) apply(opts ...RequestIDOption) {
for _, opt := range opts {
opt(o)
}
}
func (o *requestIDOptions) setRequestIDKey() {
if o.contextRequestIDKey != ContextRequestIDKey {
ContextRequestIDKey = o.contextRequestIDKey
}
if o.headerXRequestIDKey != HeaderXRequestIDKey {
HeaderXRequestIDKey = o.headerXRequestIDKey
}
}
// WithContextRequestIDKey set context request id key, minimum length of 4
func WithContextRequestIDKey(key string) RequestIDOption {
return func(o *requestIDOptions) {
if len(key) < 4 {
return
}
o.contextRequestIDKey = key
}
}
// WithHeaderRequestIDKey set header request id key, minimum length of 4
func WithHeaderRequestIDKey(key string) RequestIDOption {
return func(o *requestIDOptions) {
if len(key) < 4 {
return
}
o.headerXRequestIDKey = key
}
}
// CtxKeyString for context.WithValue key type
type CtxKeyString string
// RequestIDKey request_id
var RequestIDKey = CtxKeyString(ContextRequestIDKey)
// -------------------------------------------------------------------------------------------
// RequestID is an interceptor that injects a 'request id' into the context and request/response header of each request.
func RequestID(opts ...RequestIDOption) gin.HandlerFunc {
// customized request id key
o := defaultRequestIDOptions()
o.apply(opts...)
o.setRequestIDKey()
return func(c *gin.Context) {
// Check for incoming header, use it if exists
requestID := c.Request.Header.Get(HeaderXRequestIDKey)
// Create request id
if requestID == "" {
requestID = krand.String(krand.R_All, 32)
c.Request.Header.Set(HeaderXRequestIDKey, requestID)
}
// Expose it for use in the application
c.Set(ContextRequestIDKey, requestID)
// Set X-Request-Id header
c.Writer.Header().Set(HeaderXRequestIDKey, requestID)
c.Next()
}
}
// GCtxRequestID get request id from gin.Context
func GCtxRequestID(c *gin.Context) string {
if v, isExist := c.Get(ContextRequestIDKey); isExist {
if requestID, ok := v.(string); ok {
return requestID
}
}
return ""
}
// GCtxRequestIDField get request id field from gin.Context
func GCtxRequestIDField(c *gin.Context) zap.Field {
return zap.String(ContextRequestIDKey, GCtxRequestID(c))
}
// HeaderRequestID get request id from the header
func HeaderRequestID(c *gin.Context) string {
return c.Request.Header.Get(HeaderXRequestIDKey)
}
// HeaderRequestIDField get request id field from header
func HeaderRequestIDField(c *gin.Context) zap.Field {
return zap.String(HeaderXRequestIDKey, HeaderRequestID(c))
}
// -------------------------------------------------------------------------------------------
// RequestHeaderKey request header key
var RequestHeaderKey = "request_header_key"
// WrapCtx wrap context, put the Keys and Header of gin.Context into context
func WrapCtx(c *gin.Context) context.Context {
ctx := context.WithValue(c.Request.Context(), ContextRequestIDKey, c.GetString(ContextRequestIDKey)) //nolint
//保存请求的ip
clientIP := c.ClientIP()
if clientIP == "::1" {
clientIP = "127.0.0.1"
}
ctx = context.WithValue(ctx, "clientIP", clientIP)
return context.WithValue(ctx, RequestHeaderKey, c.Request.Header) //nolint
}
// GetFromCtx get value from context
func GetFromCtx(ctx context.Context, key string) interface{} {
return ctx.Value(key)
}
// CtxRequestID get request id from context.Context
func CtxRequestID(ctx context.Context) string {
v := ctx.Value(ContextRequestIDKey)
if str, ok := v.(string); ok {
return str
}
return ""
}
// CtxRequestIDField get request id field from context.Context
func CtxRequestIDField(ctx context.Context) zap.Field {
return zap.String(ContextRequestIDKey, CtxRequestID(ctx))
}
// GetFromHeader get value from header
func GetFromHeader(ctx context.Context, key string) string {
header, ok := ctx.Value(RequestHeaderKey).(http.Header)
if !ok {
return ""
}
return header.Get(key)
}
// GetFromHeaders get values from header
func GetFromHeaders(ctx context.Context, key string) []string {
header, ok := ctx.Value(RequestHeaderKey).(http.Header)
if !ok {
return []string{}
}
return header.Values(key)
}
Loading...
马建仓 AI 助手
尝试更多
代码解读
代码找茬
代码优化
Go
1
https://gitee.com/yzsunjianguo/common_pkg.git
git@gitee.com:yzsunjianguo/common_pkg.git
yzsunjianguo
common_pkg
common_pkg
v1.0.0

搜索帮助