geekai-Go资源

package handler // * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ // * Copyright 2023 The Geek-AI Authors. All rights reserved. // * Use of this source code is governed by a Apache-2.0 license // * that can be found in the LICENSE file. // * @Author [email protected] // * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ import ( "fmt" "geekai/core" "geekai/core/types" "geekai/service" "geekai/store/model" "geekai/store/vo" "geekai/utils" "geekai/utils/resp" "github.com/imroc/req/v3" "strings" "time" "github.com/go-redis/redis/v8" "github.com/golang-jwt/jwt/v5" "github.com/gin-gonic/gin" "github.com/lionsoul2014/ip2region/binding/golang/xdb" "gorm.io/gorm" ) type UserHandler struct { BaseHandler searcher *xdb.Searcher redis *redis.Client licenseService *service.LicenseService captcha *service.CaptchaService userService *service.UserService } func NewUserHandler( app *core.AppServer, db *gorm.DB, searcher *xdb.Searcher, client *redis.Client, captcha *service.CaptchaService, userService *service.UserService, licenseService *service.LicenseService) *UserHandler { return &UserHandler{ BaseHandler: BaseHandler{DB: db, App: app}, searcher: searcher, redis: client, captcha: captcha, licenseService: licenseService, userService: userService, } } // Register user register func (h *UserHandler) Register(c *gin.Context) { // parameters process var data struct { RegWay string `json:"reg_way"` Username string `json:"username"` Mobile string `json:"mobile"` Email string `json:"email"` Password string `json:"password"` Code string `json:"code"` InviteCode string `json:"invite_code"` Key string `json:"key,omitempty"` Dots string `json:"dots,omitempty"` X int `json:"x,omitempty"` } if err := c.ShouldBindJSON(&data); err != nil { resp.ERROR(c, types.InvalidArgs) return } if h.App.SysConfig.EnabledVerify && data.RegWay == "username" { var check bool if data.X != 0 { check = h.captcha.SlideCheck(data) } else { check = h.captcha.Check(data) } if !check { resp.ERROR(c, "请先完人机验证") return } } data.Password = strings.TrimSpace(data.Password) if len(data.Password) < 8 { resp.ERROR(c, "密码长度不能少于8个字符") return } // 检测最大注册人数 var totalUser int64 h.DB.Model(&model.User{}).Count(&totalUser) if h.licenseService.GetLicense().Configs.UserNum > 0 && int(totalUser) >= h.licenseService.GetLicense().Configs.UserNum { resp.ERROR(c, "当前注册用户数已达上限，请请升级 License") return } // 检查验证码 var key string if data.RegWay == "email" { key = CodeStorePrefix + data.Email code, err := h.redis.Get(c, key).Result() if err != nil || code != data.Code { resp.ERROR(c, "验证码错误") return } } else if data.RegWay == "mobile" { key = CodeStorePrefix + data.Mobile code, err := h.redis.Get(c, key).Result() if err != nil || code != data.Code { resp.ERROR(c, "验证码错误") return } } // 验证邀请码 inviteCode := model.InviteCode{} if data.InviteCode != "" { res := h.DB.Where("code = ?", data.InviteCode).First(&inviteCode) if res.Error != nil { resp.ERROR(c, "无效的邀请码") return } } salt := utils.RandString(8) user := model.User{ Username: data.Username, Password: utils.GenPassword(data.Password, salt), Avatar: "/images/avatar/user.png", Salt: salt, Status: true, ChatRoles: utils.JsonEncode([]string{"gpt"}), // 默认只订阅通用助手角色 Power: h.App.SysConfig.InitPower, } // check if the username is existing var item model.User session := h.DB.Session(&gorm.Session{}) if data.Mobile != "" { session = session.Where("mobile = ?", data.Mobile) user.Username = data.Mobile user.Mobile = data.Mobile } else if data.Email != "" { session = session.Where("email = ?", data.Email) user.Username = data.Email user.Email = data.Email } else if data.Username != "" { session = session.Where("username = ?", data.Username) } session.First(&item) if item.Id > 0 { resp.ERROR(c, "该用户名已经被注册") return } // 被邀请人也获得赠送算力 if data.InviteCode != "" { user.Power += h.App.SysConfig.InvitePower } if h.licenseService.GetLicense().Configs.DeCopy { user.Nickname = fmt.Sprintf("用户@%d", utils.RandomNumber(6)) } else { user.Nickname = fmt.Sprintf("极客学长@%d", utils.RandomNumber(6)) } tx := h.DB.Begin() if err := tx.Create(&user).Error; err != nil { resp.ERROR(c, err.Error()) return } // 记录邀请关系 if data.InviteCode != "" { // 增加邀请数量 h.DB.Model(&model.InviteCode{}).Where("code = ?", data.InviteCode).UpdateColumn("reg_num", gorm.Expr("reg_num + ?", 1)) if h.App.SysConfig.InvitePower > 0 { err := h.userService.IncreasePower(int(inviteCode.UserId), h.App.SysConfig.InvitePower, model.PowerLog{ Type: types.PowerInvite, Model: "Invite", Remark: fmt.Sprintf("邀请用户注册奖励，金额：%d，邀请码：%s，新用户：%s", h.App.SysConfig.InvitePower, inviteCode.Code, user.Username), }) if err != nil { tx.Rollback() resp.ERROR(c, err.Error()) return } } // 添加邀请记录 err := tx.Create(&model.InviteLog{ InviterId: inviteCode.UserId, UserId: user.Id, Username: user.Username, InviteCode: inviteCode.Code, Remark: fmt.Sprintf("奖励 %d 算力", h.App.SysConfig.InvitePower), }).Error if err != nil { tx.Rollback() resp.ERROR(c, err.Error()) return } } tx.Commit() _ = h.redis.Del(c, key) // 注册成功，删除短信验证码 // 自动登录创建 token token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{ "user_id": user.Id, "expired": time.Now().Add(time.Second * time.Duration(h.App.Config.Session.MaxAge)).Unix(), }) tokenString, err := token.SignedString([]byte(h.App.Config.Session.SecretKey)) if err != nil { resp.ERROR(c, "Failed to generate token, "+err.Error()) return } // 保存到 redis key = fmt.Sprintf("users/%d", user.Id) if _, err := h.redis.Set(c, key, tokenString, 0).Result(); err != nil { resp.ERROR(c, "error with save token: "+err.Error()) return } resp.SUCCESS(c, gin.H{"token": tokenString, "user_id": user.Id, "username": user.Username}) } // Login 用户登录 func (h *UserHandler) Login(c *gin.Context) { var data struct { Username string `json:"username"` Password string `json:"password"` Key string `json:"key,omitempty"` Dots string `json:"dots,omitempty"` X int `json:"x,omitempty"` } if err := c.ShouldBindJSON(&data); err != nil { resp.ERROR(c, types.InvalidArgs) return } verifyKey := fmt.Sprintf("users/verify/%s", data.Username) needVerify, err := h.redis.Get(c, verifyKey).Bool() if h.App.SysConfig.EnabledVerify && needVerify { var check bool if data.X != 0 { check = h.captcha.SlideCheck(data) } else { check = h.captcha.Check(data) } if !check { resp.ERROR(c, "请先完人机验证") return } } var user model.User res := h.DB.Where("username = ?", data.Username).First(&user) if res.Error != nil { h.redis.Set(c, verifyKey, true, 0) resp.ERROR(c, "用户名不存在") return } password := utils.GenPassword(data.Password, user.Salt) if password != user.Password { h.redis.Set(c, verifyKey, true, 0) resp.ERROR(c, "用户名或密码错误") return } if user.Status == false { resp.ERROR(c, "该用户已被禁止登录，请联系管理员") return } // 更新最后登录时间和IP user.LastLoginIp = c.ClientIP() user.LastLoginAt = time.Now().Unix() h.DB.Model(&user).Updates(user) h.DB.Create(&model.UserLoginLog{ UserId: user.Id, Username: user.Username, LoginIp: c.ClientIP(), LoginAddress: utils.Ip2Region(h.search