RUSTDESK/RUSTDESK-AP-SERVER-SUNLIX
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(login): Captcha upgrade and add the function to ban IP addresses (#250)

Browse Source
This commit is contained in:
lejianwen
2025-05-25 00:06:06 +08:00
parent fe4a115c9d
commit f2ea022965
18 changed files with 787 additions and 160 deletions
Download Patch File Download Diff File Expand all files Collapse all files

191
http/controller/admin/login.go
View File

@@ -11,135 +11,11 @@ import (
adResp "github.com/lejianwen/rustdesk-api/v2/http/response/admin"
"github.com/lejianwen/rustdesk-api/v2/model"
"github.com/lejianwen/rustdesk-api/v2/service"
"github.com/mojocn/base64Captcha"
"sync"
"time"
)
type Login struct {
}
// Captcha 验证码结构
type Captcha struct {
Id string `json:"id"` // 验证码 ID
B64 string `json:"b64"` // base64 验证码
Code string `json:"-"` // 验证码内容
ExpiresAt time.Time `json:"-"` // 过期时间
}
type LoginLimiter struct {
mu sync.RWMutex
failCount map[string]int // 记录每个 IP 的失败次数
timestamp map[string]time.Time // 记录每个 IP 的最后失败时间
captchas map[string]Captcha // 每个 IP 的验证码
threshold int // 失败阈值
expiry time.Duration // 失败记录过期时间
}
func NewLoginLimiter(threshold int, expiry time.Duration) *LoginLimiter {
return &LoginLimiter{
failCount: make(map[string]int),
timestamp: make(map[string]time.Time),
captchas: make(map[string]Captcha),
threshold: threshold,
expiry: expiry,
}
}
// RecordFailure 记录登录失败
func (l *LoginLimiter) RecordFailure(ip string) {
l.mu.Lock()
defer l.mu.Unlock()
// 如果该 IP 的记录已经过期，重置计数
if lastTime, exists := l.timestamp[ip]; exists && time.Since(lastTime) > l.expiry {
l.failCount[ip] = 0
}
// 更新失败次数和时间戳
l.failCount[ip]++
l.timestamp[ip] = time.Now()
}
// NeedsCaptcha 检查是否需要验证码
func (l *LoginLimiter) NeedsCaptcha(ip string) bool {
l.mu.RLock()
defer l.mu.RUnlock()
// 检查记录是否存在且未过期
if lastTime, exists := l.timestamp[ip]; exists && time.Since(lastTime) <= l.expiry {
return l.failCount[ip] >= l.threshold
}
return false
}
// GenerateCaptcha 为指定 IP 生成验证码
func (l *LoginLimiter) GenerateCaptcha(ip string) Captcha {
l.mu.Lock()
defer l.mu.Unlock()
capd := base64Captcha.NewDriverString(50, 150, 5, 10, 4, "1234567890abcdefghijklmnopqrstuvwxyz", nil, nil, nil)
b64cap := base64Captcha.NewCaptcha(capd, base64Captcha.DefaultMemStore)
id, b64s, answer, err := b64cap.Generate()
if err != nil {
global.Logger.Error("Generate captcha failed: " + err.Error())
return Captcha{}
}
// 保存验证码到对应 IP
l.captchas[ip] = Captcha{
Id: id,
B64: b64s,
Code: answer,
ExpiresAt: time.Now().Add(5 * time.Minute),
}
return l.captchas[ip]
}
// VerifyCaptcha 验证指定 IP 的验证码
func (l *LoginLimiter) VerifyCaptcha(ip, code string) bool {
l.mu.RLock()
defer l.mu.RUnlock()
// 检查验证码是否存在且未过期
if captcha, exists := l.captchas[ip]; exists && time.Now().Before(captcha.ExpiresAt) {
return captcha.Code == code
}
return false
}
// RemoveCaptcha 移除指定 IP 的验证码
func (l *LoginLimiter) RemoveCaptcha(ip string) {
l.mu.Lock()
defer l.mu.Unlock()
delete(l.captchas, ip)
}
// CleanupExpired 清理过期的记录
func (l *LoginLimiter) CleanupExpired() {
l.mu.Lock()
defer l.mu.Unlock()
now := time.Now()
for ip, lastTime := range l.timestamp {
if now.Sub(lastTime) > l.expiry {
delete(l.failCount, ip)
delete(l.timestamp, ip)
delete(l.captchas, ip)
}
}
}
func (l *LoginLimiter) RemoveRecord(ip string) {
l.mu.Lock()
defer l.mu.Unlock()
delete(l.failCount, ip)
delete(l.timestamp, ip)
delete(l.captchas, ip)
}
var loginLimiter = NewLoginLimiter(3, 5*time.Minute)
// Login 登录
// @Tags 登录
// @Summary 登录
@@ -156,10 +32,16 @@ func (ct *Login) Login(c *gin.Context) {
response.Fail(c, 101, response.TranslateMsg(c, "PwdLoginDisabled"))
return
}
// 检查登录限制
loginLimiter := global.LoginLimiter
clientIp := c.ClientIP()
_, needCaptcha := loginLimiter.CheckSecurityStatus(clientIp)
f := &admin.Login{}
err := c.ShouldBindJSON(f)
clientIp := c.ClientIP()
if err != nil {
loginLimiter.RecordFailedAttempt(clientIp)
global.Logger.Warn(fmt.Sprintf("Login Fail: %s %s %s", "ParamsError", c.RemoteIP(), clientIp))
response.Fail(c, 101, response.TranslateMsg(c, "ParamsError")+err.Error())
return
@@ -167,13 +49,14 @@ func (ct *Login) Login(c *gin.Context) {
errList := global.Validator.ValidStruct(c, f)
if len(errList) > 0 {
loginLimiter.RecordFailedAttempt(clientIp)
global.Logger.Warn(fmt.Sprintf("Login Fail: %s %s %s", "ParamsError", c.RemoteIP(), clientIp))
response.Fail(c, 101, errList[0])
return
}
// 检查是否需要验证码
if loginLimiter.NeedsCaptcha(clientIp) {
if needCaptcha {
if f.Captcha == "" || !loginLimiter.VerifyCaptcha(clientIp, f.Captcha) {
response.Fail(c, 101, response.TranslateMsg(c, "CaptchaError"))
return
@@ -184,17 +67,22 @@ func (ct *Login) Login(c *gin.Context) {
if u.Id == 0 {
global.Logger.Warn(fmt.Sprintf("Login Fail: %s %s %s", "UsernameOrPasswordError", c.RemoteIP(), clientIp))
loginLimiter.RecordFailure(clientIp)
if loginLimiter.NeedsCaptcha(clientIp) {
loginLimiter.RemoveCaptcha(clientIp)
loginLimiter.RecordFailedAttempt(clientIp)
// 移除验证码，重新生成
loginLimiter.RemoveCaptcha(clientIp)
if _, needCaptcha = loginLimiter.CheckSecurityStatus(clientIp); needCaptcha {
response.Fail(c, 110, response.TranslateMsg(c, "UsernameOrPasswordError"))
} else {
response.Fail(c, 101, response.TranslateMsg(c, "UsernameOrPasswordError"))
}
response.Fail(c, 101, response.TranslateMsg(c, "UsernameOrPasswordError"))
return
}
if !service.AllService.UserService.CheckUserEnable(u) {
if loginLimiter.NeedsCaptcha(clientIp) {
if needCaptcha {
loginLimiter.RemoveCaptcha(clientIp)
response.Fail(c, 110, response.TranslateMsg(c, "UserDisabled"))
return
}
response.Fail(c, 101, response.TranslateMsg(c, "UserDisabled"))
return
@@ -209,23 +97,36 @@ func (ct *Login) Login(c *gin.Context) {
Platform: f.Platform,
})
// 成功后清除记录
loginLimiter.RemoveRecord(clientIp)
// 清理过期记录
go loginLimiter.CleanupExpired()
// 登录成功，清除登录限制
loginLimiter.RemoveAttempts(clientIp)
responseLoginSuccess(c, u, ut.Token)
}
func (ct *Login) Captcha(c *gin.Context) {
loginLimiter := global.LoginLimiter
clientIp := c.ClientIP()
if !loginLimiter.NeedsCaptcha(clientIp) {
banned, needCaptcha := loginLimiter.CheckSecurityStatus(clientIp)
if banned {
response.Fail(c, 101, response.TranslateMsg(c, "LoginBanned"))
return
}
if !needCaptcha {
response.Fail(c, 101, response.TranslateMsg(c, "NoCaptchaRequired"))
return
}
captcha := loginLimiter.GenerateCaptcha(clientIp)
err, captcha := loginLimiter.RequireCaptcha(clientIp)
if err != nil {
response.Fail(c, 101, response.TranslateMsg(c, "CaptchaError")+err.Error())
return
}
err, b64 := loginLimiter.DrawCaptcha(captcha.Content)
if err != nil {
response.Fail(c, 101, response.TranslateMsg(c, "CaptchaError")+err.Error())
return
}
response.Success(c, gin.H{
"captcha": captcha,
"captcha": gin.H{
"b64": b64,
},
})
}
@@ -257,12 +158,18 @@ func (ct *Login) Logout(c *gin.Context) {
// @Failure 500 {object} response.ErrorResponse
// @Router /admin/login-options [post]
func (ct *Login) LoginOptions(c *gin.Context) {
ip := c.ClientIP()
loginLimiter := global.LoginLimiter
clientIp := c.ClientIP()
banned, needCaptcha := loginLimiter.CheckSecurityStatus(clientIp)
if banned {
response.Fail(c, 101, response.TranslateMsg(c, "LoginBanned"))
return
}
ops := service.AllService.OauthService.GetOauthProviders()
response.Success(c, gin.H{
"ops": ops,
"register": global.Config.App.Register,
"need_captcha": loginLimiter.NeedsCaptcha(ip),
"need_captcha": needCaptcha,
})
}

8
http/controller/api/login.go
View File

@@ -31,10 +31,16 @@ func (l *Login) Login(c *gin.Context) {
response.Error(c, response.TranslateMsg(c, "PwdLoginDisabled"))
return
}
// 检查登录限制
loginLimiter := global.LoginLimiter
clientIp := c.ClientIP()
f := &api.LoginForm{}
err := c.ShouldBindJSON(f)
//fmt.Println(f)
if err != nil {
loginLimiter.RecordFailedAttempt(clientIp)
global.Logger.Warn(fmt.Sprintf("Login Fail: %s %s %s", "ParamsError", c.RemoteIP(), c.ClientIP()))
response.Error(c, response.TranslateMsg(c, "ParamsError")+err.Error())
return
@@ -42,6 +48,7 @@ func (l *Login) Login(c *gin.Context) {
errList := global.Validator.ValidStruct(c, f)
if len(errList) > 0 {
loginLimiter.RecordFailedAttempt(clientIp)
global.Logger.Warn(fmt.Sprintf("Login Fail: %s %s %s", "ParamsError", c.RemoteIP(), c.ClientIP()))
response.Error(c, errList[0])
return
@@ -50,6 +57,7 @@ func (l *Login) Login(c *gin.Context) {
u := service.AllService.UserService.InfoByUsernamePassword(f.Username, f.Password)
if u.Id == 0 {
loginLimiter.RecordFailedAttempt(clientIp)
global.Logger.Warn(fmt.Sprintf("Login Fail: %s %s %s", "UsernameOrPasswordError", c.RemoteIP(), c.ClientIP()))
response.Error(c, response.TranslateMsg(c, "UsernameOrPasswordError"))
return

2
http/http.go
View File

@@ -33,7 +33,7 @@ func ApiInit() {
g.NoRoute(func(c *gin.Context) {
c.String(http.StatusNotFound, "404 not found")
})
g.Use(middleware.Logger(), gin.Recovery())
g.Use(middleware.Logger(), middleware.Limiter(), gin.Recovery())
router.WebInit(g)
router.Init(g)
router.ApiInit(g)

22
http/middleware/limiter.go Normal file
View File

@@ -0,0 +1,22 @@
package middleware
import (
"github.com/gin-gonic/gin"
"github.com/lejianwen/rustdesk-api/v2/global"
"github.com/lejianwen/rustdesk-api/v2/http/response"
"net/http"
)
func Limiter() gin.HandlerFunc {
return func(c *gin.Context) {
loginLimiter := global.LoginLimiter
clientIp := c.ClientIP()
banned, _ := loginLimiter.CheckSecurityStatus(clientIp)
if banned {
response.Fail(c, http.StatusLocked, response.TranslateMsg(c, "Banned"))
c.Abort()
return
}
c.Next()
}
}