ifms_go/pkg/util/context.go
2025-06-10 17:50:46 +08:00

141 lines
2.9 KiB
Go

package util
import (
"context"
"crypto/md5"
"encoding/hex"
"ifms/internal/config"
"ifms/pkg/convert"
"gorm.io/gorm"
"ifms/pkg/encoding/json"
)
type (
traceIDCtx struct{}
transCtx struct{}
rowLockCtx struct{}
userIDCtx struct{}
userTokenCtx struct{}
isRootUserCtx struct{}
userCacheCtx struct{}
loginUserCtx struct{}
)
func NewTraceID(ctx context.Context, traceID string) context.Context {
return context.WithValue(ctx, traceIDCtx{}, traceID)
}
func FromTraceID(ctx context.Context) string {
v := ctx.Value(traceIDCtx{})
if v != nil {
return v.(string)
}
return ""
}
func NewTrans(ctx context.Context, db *gorm.DB) context.Context {
return context.WithValue(ctx, transCtx{}, db)
}
func FromTrans(ctx context.Context) (*gorm.DB, bool) {
v := ctx.Value(transCtx{})
if v != nil {
return v.(*gorm.DB), true
}
return nil, false
}
func NewRowLock(ctx context.Context) context.Context {
return context.WithValue(ctx, rowLockCtx{}, true)
}
func FromRowLock(ctx context.Context) bool {
v := ctx.Value(rowLockCtx{})
return v != nil && v.(bool)
}
func NewLoginUser(ctx context.Context, loginUser *config.LoginUser) context.Context {
return context.WithValue(ctx, loginUserCtx{}, loginUser)
}
func FromLoginUser(ctx context.Context) *config.LoginUser {
v := ctx.Value(loginUserCtx{})
if v != nil {
return v.(*config.LoginUser)
}
return &config.LoginUser{}
}
func NewUserID(ctx context.Context, userID string) context.Context {
return context.WithValue(ctx, userIDCtx{}, userID)
}
func FromUserID(ctx context.Context) int64 {
v := ctx.Value(userIDCtx{})
if v != nil {
return convert.ToInt64(v)
}
return 0
}
func NewUserToken(ctx context.Context, userToken string) context.Context {
return context.WithValue(ctx, userTokenCtx{}, userToken)
}
func FromUserToken(ctx context.Context) string {
v := ctx.Value(userTokenCtx{})
if v != nil {
return v.(string)
}
return ""
}
func NewIsRootUser(ctx context.Context) context.Context {
return context.WithValue(ctx, isRootUserCtx{}, true)
}
func FromIsRootUser(ctx context.Context) bool {
v := ctx.Value(isRootUserCtx{})
return v != nil && v.(bool)
}
// Set user cache object
type UserCache struct {
RoleIDs []int64 `json:"rids"`
}
func ParseUserCache(s string) UserCache {
var a UserCache
if s == "" {
return a
}
_ = json.Unmarshal([]byte(s), &a)
return a
}
func (a UserCache) String() string {
return json.MarshalToString(a)
}
func NewUserCache(ctx context.Context, userCache UserCache) context.Context {
return context.WithValue(ctx, userCacheCtx{}, userCache)
}
func FromUserCache(ctx context.Context) UserCache {
v := ctx.Value(userCacheCtx{})
if v != nil {
return v.(UserCache)
}
return UserCache{}
}
func EncryptToMD5(s string) string {
// 创建一个 MD5 哈希对象
h := md5.New()
// 将字符串写入哈希对象
h.Write([]byte(s))
// 将字节数组转换为十六进制字符串
return hex.EncodeToString(h.Sum(nil))
}