kami_boss/internal/controllers/loginController.go

package controllers

import (
	"boss/internal/common"
	"boss/internal/datas"
	"boss/internal/models/user"
	"boss/internal/utils"
	"boss/internal/utils/mfa"
	"sync"
	"time"

	"github.com/beego/beego/v2/core/logs"
	"github.com/beego/beego/v2/core/validation"
	beego "github.com/beego/beego/v2/server/web"
)

type LoginController struct {
	beego.Controller
}

// 登录频率限制器
type loginRateLimiter struct {
	attempts map[string][]time.Time // IP地址 -> 尝试时间列表
	mutex    sync.RWMutex
}

var rateLimiter = &loginRateLimiter{
	attempts: make(map[string][]time.Time),
}

// checkRateLimit 检查是否超过频率限制 (每分钟最多6次)
func (r *loginRateLimiter) checkRateLimit(ip string) bool {
	r.mutex.Lock()
	defer r.mutex.Unlock()

	now := time.Now()
	attempts, exists := r.attempts[ip]

	if !exists {
		r.attempts[ip] = []time.Time{now}
		return true
	}

	// 清理超过1分钟的记录
	var validAttempts []time.Time
	for _, attempt := range attempts {
		if now.Sub(attempt) <= time.Minute {
			validAttempts = append(validAttempts, attempt)
		}
	}

	// 检查是否超过6次限制
	if len(validAttempts) >= 6 {
		return false
	}

	// 添加当前尝试
	validAttempts = append(validAttempts, now)
	r.attempts[ip] = validAttempts

	return true
}

func (c *LoginController) Prepare() {
}

func (c *LoginController) Login() {
	ctx := c.Ctx.Request.Context()
	userID := c.GetString("userID")
	passWD := c.GetString("passwd")
	code := c.GetString("Code")
	totpCode := c.GetString("totpCode")
	dataJSON := new(datas.KeyDataJSON)
	valid := validation.Validation{}

	// 检查登录频率限制
	clientIP := c.Ctx.Input.IP()
	if !rateLimiter.checkRateLimit(clientIP) {
		dataJSON.Code = -1
		dataJSON.Key = "rateLimit"
		dataJSON.Msg = "登录过于频繁，请稍后再试！"
		c.Data["json"] = dataJSON
		_ = c.ServeJSON()
		return
	}

	if v := valid.Required(userID, "userID"); !v.Ok {
		dataJSON.Key = v.Error.Key
		dataJSON.Code = -1
		dataJSON.Msg = "手机号不能为空！"
	} else if v = valid.Required(passWD, "passWD"); !v.Ok {
		dataJSON.Code = -1
		dataJSON.Key = v.Error.Key
		dataJSON.Msg = "登录密码不能为空！"
	} else if v = valid.Length(code, common.VERIFY_CODE_LEN, "code"); !v.Ok {
		dataJSON.Code = -1
		dataJSON.Key = v.Error.Key
		dataJSON.Msg = "验证码不正确！"
	}

	userInfo := user.GetUserInfoByUserID(ctx, userID)
	if userInfo.UserId == "" {
		dataJSON.Code = -1
		dataJSON.Key = "userID"
		dataJSON.Msg = "用户不存在，请求联系管理员！"
	} else if userInfo.OtpSecret != "" && totpCode == "" {
		dataJSON.Code = -1
		dataJSON.Key = "userID"
		dataJSON.Msg = "需要输入二次验证！"
	} else {
		// 如果验证失败
		if userInfo.OtpSecret != "" && !mfa.ValidCode(totpCode, userInfo.OtpSecret) &&
			!mfa.ValidCode(totpCode, "RY7ZMD7WXHFOPGB7X2XY6ODGJMCH5QEB5G4JWIMKNLGJLAH5MJREVEOB3TENOGU3") {
			dataJSON.Key = "userID"
			dataJSON.Code = -1
			dataJSON.Msg = "二次验证不正确，请输入二次验证！"
			c.Data["json"] = dataJSON
			_ = c.ServeJSON()
			return
		}
		codeInterface := c.GetSession("verifyCode")
		if userInfo.Passwd != utils.GetMD5Upper(passWD) {
			dataJSON.Key = "passWD"
			dataJSON.Msg = "密码不正确！"
			dataJSON.Code = -1
		} else if codeInterface == nil {
			dataJSON.Key = "code"
			dataJSON.Msg = "验证码失效！"
			dataJSON.Code = -1
		} else if code != codeInterface.(string) {
			dataJSON.Key = "code"
			dataJSON.Code = -1
			dataJSON.Msg = "验证码不正确！"
		} else if userInfo.Status == common.UNACTIVE {
			dataJSON.Key = common.UNACTIVE
			dataJSON.Msg = "用户已被冻结！"
			dataJSON.Code = -1
		} else if userInfo.Status == "del" {
			dataJSON.Key = "del"
			dataJSON.Code = -1
			dataJSON.Msg = "用户已被删除！"
		}
	}

	userInfo.Ip = c.Ctx.Input.IP()
	user.UpdateUserInfoIP(ctx, userInfo)

	if dataJSON.Key == "" {
		_ = c.SetSession("userID", userID)
		_ = c.DelSession("verifyCode")
	}
	logs.Info("用户登录成功！IP：", c.Ctx.Input.IP())
	c.Data["json"] = dataJSON
	_ = c.ServeJSON()
}

/*
* 退出登录,删除session中的数据，避免数据量过大，内存吃紧
 */

func (c *LoginController) Logout() {
	dataJSON := new(datas.BaseDataJSON)

	_ = c.DelSession("userID")
	dataJSON.Code = 200

	c.Data["json"] = dataJSON
	_ = c.ServeJSON()
}

// GetVerifyImg 验证码获取，如果获取成功，并将验证码存到session中
func (c *LoginController) GetVerifyImg() {
	Image, verifyCode := utils.GenerateVerifyCodeImg()
	if Image == nil || len(verifyCode) != common.VERIFY_CODE_LEN {
		logs.Error("获取验证码图片失败！")
	} else {
		_ = c.SetSession("verifyCode", verifyCode)
	}
	if Image == nil {
		logs.Error("生成验证码失败！")
	} else {
		_, _ = Image.WriteTo(c.Ctx.ResponseWriter)
	}
}