package util import ( "context" "gitea.com/red-future/common/utils" "github.com/gogf/gf/v2/frame/g" "github.com/gogf/gf/v2/util/gconv" ) // AsyncCtx 固化异步上下文中的 token 和用户信息,避免请求结束后丢失 func AsyncCtx(ctx context.Context) context.Context { asyncCtx := context.WithoutCancel(ctx) if r := g.RequestFromCtx(ctx); r != nil { if token := r.Header.Get("Authorization"); token != "" { asyncCtx = context.WithValue(asyncCtx, "token", token) } if userInfo := r.Header.Get("X-User-Info"); userInfo != "" { asyncCtx = context.WithValue(asyncCtx, "xUserInfo", userInfo) } } if user, err := utils.GetUserInfo(ctx); err == nil && user != nil { asyncCtx = context.WithValue(asyncCtx, "user", user) } return asyncCtx } // ForwardHeaders 透传调用链路的头信息,优先使用 ctx 中的固化值 func ForwardHeaders(ctx context.Context) map[string]string { headers := make(map[string]string) SetHeaderFromContext(headers, ctx, "Authorization", "token") SetHeaderFromContext(headers, ctx, "X-User-Info", "xUserInfo") FallbackToRequestHeaders(headers, ctx) return headers } // SetHeaderFromContext 从上下文中设置 header func SetHeaderFromContext(headers map[string]string, ctx context.Context, headerKey, ctxKey string) { if value, ok := ctx.Value(ctxKey).(string); ok && value != "" { headers[headerKey] = value } } // FallbackToRequestHeaders 从请求头中获取作为兜底 func FallbackToRequestHeaders(headers map[string]string, ctx context.Context) { r := g.RequestFromCtx(ctx) if r == nil { return } if headers["Authorization"] == "" { if token := r.Header.Get("Authorization"); token != "" { headers["Authorization"] = token } } if headers["X-User-Info"] == "" { if userInfo := r.Header.Get("X-User-Info"); userInfo != "" { headers["X-User-Info"] = userInfo } } } // SetTaskHeadersToCtx 把任务入库时保存的 header 信息注入 ctx,给 worker 调 OSS 用 func SetTaskHeadersToCtx(ctx context.Context, headers map[string]string) context.Context { if headers == nil { return ctx } if v := gconv.String(headers["Authorization"]); v != "" { ctx = context.WithValue(ctx, "token", v) } if v := gconv.String(headers["X-User-Info"]); v != "" { ctx = context.WithValue(ctx, "xUserInfo", v) } return ctx }