Merge pull request #347 from wsczx/dev

解决防爆并行问题

server/admin/lockmanager.go

@@ -2,7 +2,6 @@ package admin
 
 import (
 	"encoding/json"
-	"fmt"
 	"io"
 	"net"
 	"net/http"
@@ -31,8 +30,8 @@ type IPWhitelists struct {
 }
 
 type LockManager struct {
-	mu            sync.Mutex
+	mu sync.Mutex
-	LoginStatus   sync.Map                         // 登录状态
+	// LoginStatus   sync.Map                         // 登录状态
 	ipLocks       map[string]*LockState            // 全局IP锁定状态
 	userLocks     map[string]*LockState            // 全局用户锁定状态
 	ipUserLocks   map[string]map[string]*LockState // 单用户IP锁定状态
@@ -46,7 +45,7 @@ var once sync.Once
 func GetLockManager() *LockManager {
 	once.Do(func() {
 		lockmanager = &LockManager{
-			LoginStatus:  sync.Map{},
+			// LoginStatus:  sync.Map{},
 			ipLocks:      make(map[string]*LockState),
 			userLocks:    make(map[string]*LockState),
 			ipUserLocks:  make(map[string]map[string]*LockState),
@@ -89,7 +88,7 @@ func UnlockUser(w http.ResponseWriter, r *http.Request) {
 	}
 
 	if lockinfo.State == nil {
-		RespError(w, RespInternalErr, fmt.Errorf("未找到锁定用户！"))
+		RespError(w, RespInternalErr, "未找到锁定用户！")
 		return
 	}
 	lm := GetLockManager()
@@ -111,7 +110,7 @@ func UnlockUser(w http.ResponseWriter, r *http.Request) {
 	}
 
 	if state == nil || !state.Locked {
-		RespError(w, RespInternalErr, fmt.Errorf("锁定状态未找到或已解锁"))
+		RespError(w, RespInternalErr, "锁定状态未找到或已解锁")
 		return
 	}
 
@@ -238,14 +237,14 @@ func (lm *LockManager) CleanupExpiredLocks() {
 	defer lm.mu.Unlock()
 
 	for ip, state := range lm.ipLocks {
-		if !lm.CheckLockState(state, now, base.Cfg.GlobalIPLockTime) ||
+		if !lm.CheckLockState(state, now, base.Cfg.GlobalIPBanResetTime) ||
 			now.Sub(state.LastAttempt) > time.Duration(base.Cfg.GlobalLockStateExpirationTime)*time.Second {
 			delete(lm.ipLocks, ip)
 		}
 	}
 
 	for user, state := range lm.userLocks {
-		if !lm.CheckLockState(state, now, base.Cfg.GlobalUserLockTime) ||
+		if !lm.CheckLockState(state, now, base.Cfg.GlobalUserBanResetTime) ||
 			now.Sub(state.LastAttempt) > time.Duration(base.Cfg.GlobalLockStateExpirationTime)*time.Second {
 			delete(lm.userLocks, user)
 		}
@@ -253,7 +252,7 @@ func (lm *LockManager) CleanupExpiredLocks() {
 
 	for user, ipMap := range lm.ipUserLocks {
 		for ip, state := range ipMap {
-			if !lm.CheckLockState(state, now, base.Cfg.LockTime) ||
+			if !lm.CheckLockState(state, now, base.Cfg.BanResetTime) ||
 				now.Sub(state.LastAttempt) > time.Duration(base.Cfg.GlobalLockStateExpirationTime)*time.Second {
 				delete(ipMap, ip)
 				if len(ipMap) == 0 {
@@ -410,3 +409,57 @@ func (lm *LockManager) Unlock(state *LockState) {
 	state.LockTime = time.Time{}
 	state.Locked = false
 }
+
+// 检查锁定状态
+func (lm *LockManager) CheckLocked(username, ipaddr string) bool {
+	if !base.Cfg.AntiBruteForce {
+		return true
+	}
+
+	ip, _, err := net.SplitHostPort(ipaddr) // 提取纯 IP 地址，去掉端口号
+	if err != nil {
+		base.Error("检查锁定状态失败,提取IP地址错误:", ipaddr)
+		return true
+	}
+	now := time.Now()
+
+	// 检查IP是否在白名单中
+	if lm.IsWhitelisted(ip) {
+		return true
+	}
+
+	// 检查全局 IP 锁定
+	if base.Cfg.MaxGlobalIPBanCount > 0 && lm.CheckGlobalIPLock(ip, now) {
+		base.Warn("IP", ip, "is globally locked. Try again later.")
+		return false
+	}
+
+	// 检查全局用户锁定
+	if base.Cfg.MaxGlobalUserBanCount > 0 && lm.CheckGlobalUserLock(username, now) {
+		base.Warn("User", username, "is globally locked. Try again later.")
+		return false
+	}
+
+	// 检查单个用户的 IP 锁定
+	if base.Cfg.MaxBanCount > 0 && lm.CheckUserIPLock(username, ip, now) {
+		base.Warn("IP", ip, "is locked for user", username, "Try again later.")
+		return false
+	}
+
+	return true
+}
+
+// 更新用户登录状态
+func (lm *LockManager) UpdateLoginStatus(username, ipaddr string, loginStatus bool) {
+	ip, _, err := net.SplitHostPort(ipaddr) // 提取纯 IP 地址，去掉端口号
+	if err != nil {
+		base.Error("更新登录状态失败,提取IP地址错误:", ipaddr)
+		return
+	}
+	now := time.Now()
+
+	// 更新用户登录状态
+	lm.UpdateGlobalIPLock(ip, now, loginStatus)
+	lm.UpdateGlobalUserLock(username, now, loginStatus)
+	lm.UpdateUserIPLock(username, ip, now, loginStatus)
+}

server/handler/antiBruteForce.go

@@ -9,11 +9,10 @@ import (
 	"strings"
 	"time"
 
-	"github.com/bjdgyc/anylink/admin"
 	"github.com/bjdgyc/anylink/base"
 )
 
-var lockManager = admin.GetLockManager()
+// var lockManager = admin.GetLockManager()
 
 const loginStatusKey = "login_status"
 

server/handler/link_auth.go

@@ -77,6 +77,13 @@ func LinkAuth(w http.ResponseWriter, r *http.Request) {
 		w.WriteHeader(http.StatusBadRequest)
 		return
 	}
+
+	// 锁定状态判断
+	if !lockManager.CheckLocked(cr.Auth.Username, r.RemoteAddr) {
+		w.WriteHeader(http.StatusTooManyRequests)
+		return
+	}
+
 	// 用户活动日志
 	ua := &dbdata.UserActLog{
 		Username:        cr.Auth.Username,
@@ -95,8 +102,9 @@ func LinkAuth(w http.ResponseWriter, r *http.Request) {
 	err = dbdata.CheckUser(cr.Auth.Username, cr.Auth.Password, cr.GroupSelect)
 	if err != nil {
 		// lockManager.LoginStatus.Store(loginStatusKey, false) // 记录登录失败状态
-		hc := r.Context().Value(loginStatusKey).(*HttpContext)
+		// hc := r.Context().Value(loginStatusKey).(*HttpContext)
-		hc.LoginStatus = false
+		// hc.LoginStatus = false
+		lockManager.UpdateLoginStatus(cr.Auth.Username, r.RemoteAddr, false) // 记录登录失败状态
 
 		base.Warn(err, r.RemoteAddr)
 		ua.Info = err.Error()
@@ -123,8 +131,9 @@ func LinkAuth(w http.ResponseWriter, r *http.Request) {
 	// 用户otp验证
 	if base.Cfg.AuthAloneOtp && !v.DisableOtp {
 		// lockManager.LoginStatus.Store(loginStatusKey, true) // 重置OTP验证计数
-		hc := r.Context().Value(loginStatusKey).(*HttpContext)
+		// hc := r.Context().Value(loginStatusKey).(*HttpContext)
-		hc.LoginStatus = true
+		// hc.LoginStatus = true
+		lockManager.UpdateLoginStatus(cr.Auth.Username, r.RemoteAddr, true) // 重置OTP验证计数
 
 		sessionID, err := GenerateSessionID()
 		if err != nil {

server/handler/link_auth_otp.go

@@ -9,6 +9,7 @@ import (
 	"net/http"
 	"sync"
 
+	"github.com/bjdgyc/anylink/admin"
 	"github.com/bjdgyc/anylink/base"
 	"github.com/bjdgyc/anylink/dbdata"
 	"github.com/bjdgyc/anylink/pkg/utils"
@@ -16,6 +17,7 @@ import (
 )
 
 var SessStore = NewSessionStore()
+var lockManager = admin.GetLockManager()
 
 // const maxOtpErrCount = 3
 
@@ -110,12 +112,13 @@ func DeleteCookie(w http.ResponseWriter, name string) {
 }
 func CreateSession(w http.ResponseWriter, r *http.Request, authSession *AuthSession) {
 	// lockManager.LoginStatus.Store(loginStatusKey, true) // 更新登录成功状态
-	hc := r.Context().Value(loginStatusKey).(*HttpContext)
+	// hc := r.Context().Value(loginStatusKey).(*HttpContext)
-	hc.LoginStatus = true
+	// hc.LoginStatus = true
-
 	cr := authSession.ClientRequest
 	ua := authSession.UserActLog
 
+	lockManager.UpdateLoginStatus(cr.Auth.Username, r.RemoteAddr, true) // 更新登录成功状态
+
 	sess := sessdata.NewSession("")
 	sess.Username = cr.Auth.Username
 	sess.Group = cr.GroupSelect
@@ -196,6 +199,13 @@ func LinkAuth_otp(w http.ResponseWriter, r *http.Request) {
 	otpSecret := sessionData.ClientRequest.Auth.OtpSecret
 	otp := cr.Auth.SecondaryPassword
 
+	// 锁定状态判断
+	if !lockManager.CheckLocked(username, r.RemoteAddr) {
+		w.WriteHeader(http.StatusTooManyRequests)
+		SessStore.DeleteAuthSession(sessionID)
+		return
+	}
+
 	// 动态码错误
 	if !dbdata.CheckOtp(username, otp, otpSecret) {
 		// if sessionData.AddOtpErrCount(1) > maxOtpErrCount {
@@ -204,8 +214,9 @@ func LinkAuth_otp(w http.ResponseWriter, r *http.Request) {
 		// 	return
 		// }
 		// lockManager.LoginStatus.Store(loginStatusKey, false) // 记录登录失败状态
-		hc := r.Context().Value(loginStatusKey).(*HttpContext)
+		// hc := r.Context().Value(loginStatusKey).(*HttpContext)
-		hc.LoginStatus = false
+		// hc.LoginStatus = false
+		lockManager.UpdateLoginStatus(username, r.RemoteAddr, false) // 记录登录失败状态
 
 		base.Warn("OTP 动态码错误", username, r.RemoteAddr)
 		ua.Info = "OTP 动态码错误"

server/handler/server.go

@@ -111,10 +111,12 @@ func initRoute() http.Handler {
 	})
 
 	r.HandleFunc("/", LinkHome).Methods(http.MethodGet)
-	r.Handle("/", antiBruteForce(http.HandlerFunc(LinkAuth))).Methods(http.MethodPost)
+	r.HandleFunc("/", LinkAuth).Methods(http.MethodPost)
+	// r.Handle("/", antiBruteForce(http.HandlerFunc(LinkAuth))).Methods(http.MethodPost)
 	r.HandleFunc("/CSCOSSLC/tunnel", LinkTunnel).Methods(http.MethodConnect)
 	r.HandleFunc("/otp_qr", LinkOtpQr).Methods(http.MethodGet)
-	r.Handle("/otp-verification", antiBruteForce(http.HandlerFunc(LinkAuth_otp))).Methods(http.MethodPost)
+	r.HandleFunc("/otp-verification", LinkAuth_otp).Methods(http.MethodPost)
+	// r.Handle("/otp-verification", antiBruteForce(http.HandlerFunc(LinkAuth_otp))).Methods(http.MethodPost)
 	r.HandleFunc(fmt.Sprintf("/profile_%s.xml", base.Cfg.ProfileName), func(w http.ResponseWriter, r *http.Request) {
 		b, _ := os.ReadFile(base.Cfg.Profile)
 		w.Write(b)