101 lines
2.8 KiB
Go
101 lines
2.8 KiB
Go
|
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
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// ParseStoredPayload 解析入库的 request_payload，拆出模型调用 payload 与透传 headers
|
|||
|
|
// 入库格式：{"payload": <any>, "headers": {"Authorization": "...", "X-User-Info":"..."}}
|
|||
|
|
func ParseStoredPayload(v any) (payload any, headers map[string]string) {
|
|||
|
|
if v == nil {
|
|||
|
|
return nil, nil
|
|||
|
|
}
|
|||
|
|
m := gconv.Map(v)
|
|||
|
|
if len(m) == 0 {
|
|||
|
|
return v, nil
|
|||
|
|
}
|
|||
|
|
if h, ok := m["headers"]; ok {
|
|||
|
|
headers = gconv.MapStrStr(h)
|
|||
|
|
}
|
|||
|
|
if p, ok := m["payload"]; ok {
|
|||
|
|
payload = p
|
|||
|
|
} else {
|
|||
|
|
payload = v
|
|||
|
|
}
|
|||
|
|
return
|
|||
|
|
}
|