package token import ( "context" "crypto/md5" "encoding/hex" "io" "strconv" "time" "wechat-api/ent" "wechat-api/ent/agent" "wechat-api/ent/token" "wechat-api/internal/utils/dberrorhandler" "wechat-api/internal/svc" "wechat-api/internal/types" "github.com/zeromicro/go-zero/core/logx" ) type CheckTokenLogic struct { logx.Logger ctx context.Context svcCtx *svc.ServiceContext } func NewCheckTokenLogic(ctx context.Context, svcCtx *svc.ServiceContext) *CheckTokenLogic { return &CheckTokenLogic{ Logger: logx.WithContext(ctx), ctx: ctx, svcCtx: svcCtx} } func (l *CheckTokenLogic) CheckToken(req *types.CheckTokenReq) (resp *types.CheckTokenResp, err error) { var ( valid bool sign string ) timestamp := time.Now().Unix() tokenItem, err := l.svcCtx.DB.Token.Query().Where(token.TokenEQ(*req.Token)).Limit(1).Only(l.ctx) if err != nil && !ent.IsNotFound(err) { return nil, dberrorhandler.DefaultEntError(l.Logger, err, req) } var agentInfo types.AgentInfo var customAgentBase string var customAgentKey string var openaiBase string var openaiKey string if tokenItem == nil { // 判断Token是否存在 valid = false } else if tokenItem.ExpireAt.Unix() > timestamp { // 判断Token是否过期 if tokenItem.MAC == "" { // 判断MAC是否存在,如果不存咋则说明是第一次激活该Token valid = true err = l.svcCtx.DB.Token.UpdateOneID(tokenItem.ID).SetNotNilMAC(req.Mac).Exec(l.ctx) if err != nil { return nil, dberrorhandler.DefaultEntError(l.Logger, err, req) } } else if tokenItem.MAC != *req.Mac { // 如果Mac地址不匹配说明该请求不合法 valid = false } else { // 这是正常验证,返回OK valid = true } if valid { sign = generateSign(*req.Token, *req.Mac, strconv.Itoa(int(timestamp))) } } if valid { customAgentBase = tokenItem.CustomAgentBase customAgentKey = tokenItem.CustomAgentKey openaiBase = tokenItem.OpenaiBase openaiKey = tokenItem.OpenaiKey if tokenItem.AgentID == 0 { agentInfo = types.AgentInfo{ BaseIDInfo: types.BaseIDInfo{ Id: &tokenItem.AgentID, }, } } else { data, _ := l.svcCtx.DB.Agent.Query().Where(agent.ID(tokenItem.AgentID)).First(l.ctx) if data != nil { agentInfo = types.AgentInfo{ BaseIDInfo: types.BaseIDInfo{ Id: &data.ID, }, Name: &data.Name, Role: &data.Role, Background: &data.Background, Examples: &data.Examples, DatasetId: &data.DatasetID, CollectionId: &data.CollectionID, } } } } datasetBase := "https://fastgpt.gkscrm.com/api/core/dataset/searchTest" datasetKey := "fastgpt-czbAiqKKse65hwwviZhwkgvyDEKE3aeB31Fx18oUsAGojyWQ672HdsWZi1E1C" return &types.CheckTokenResp{ Valid: &valid, Sign: &sign, Timestamp: ×tamp, AgentInfo: &agentInfo, CustomAgentBase: &customAgentBase, CustomAgentKey: &customAgentKey, OpenaiBase: &openaiBase, OpenaiKey: &openaiKey, DatasetBase: &datasetBase, DatasetKey: &datasetKey, }, nil } // 内部Md5方法 func generateSign(token, mac, timestamp string) string { md5Token := easyMd5(token) md5Mac := easyMd5(mac) md5Ts := easyMd5(timestamp) return easyMd5(md5Token + md5Mac + md5Ts) } // 简易Md5用法 func easyMd5(s string) string { hash := md5.New() _, _ = io.WriteString(hash, s) return hex.EncodeToString(hash.Sum(nil)) }