修改 LoginStatus 用 context 传递

bjdgyc · Nov 12, 2024 · 2b757b6 · 2b757b6
1 parent 9ef2954
commit 2b757b6
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 9 deletions.
diff --git a/server/handler/antiBruteForce.go b/server/handler/antiBruteForce.go
@@ -1,6 +1,7 @@
 package handler
 
 import (
+	"context"
 	"encoding/xml"
 	"io"
 	"net"
@@ -16,15 +17,24 @@ var lockManager = admin.GetLockManager()
 
 const loginStatusKey = "login_status"
 
+type HttpContext struct {
+	LoginStatus bool // 登录状态
+}
+
 // 防爆破中间件
 func antiBruteForce(next http.Handler) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	return http.HandlerFunc(func(w http.ResponseWriter, old_r *http.Request) {
 		// 防爆破功能全局开关
 		if !base.Cfg.AntiBruteForce {
-			next.ServeHTTP(w, r)
+			next.ServeHTTP(w, old_r)
 			return
 		}
 
+		// 非并发安全
+		hc := &HttpContext{}
+		ctx := context.WithValue(context.Background(), loginStatusKey, hc)
+		r := old_r.WithContext(ctx)
+
 		body, err := io.ReadAll(r.Body)
 		if err != nil {
 			http.Error(w, "Failed to read request body", http.StatusBadRequest)
@@ -96,15 +106,17 @@ func antiBruteForce(next http.Handler) http.Handler {
 		next.ServeHTTP(w, r)
 
 		// 检查登录状态
-		Status, _ := lockManager.LoginStatus.Load(loginStatusKey)
-		loginStatus, _ := Status.(bool)
+		// Status, _ := lockManager.LoginStatus.Load(loginStatusKey)
+		// loginStatus, _ := Status.(bool)
+
+		loginStatus := hc.LoginStatus
 
 		// 更新用户登录状态
 		lockManager.UpdateGlobalIPLock(ip, now, loginStatus)
 		lockManager.UpdateGlobalUserLock(username, now, loginStatus)
 		lockManager.UpdateUserIPLock(username, ip, now, loginStatus)
 
 		// 清除登录状态
-		lockManager.LoginStatus.Delete(loginStatusKey)
+		// lockManager.LoginStatus.Delete(loginStatusKey)
 	})
 }
diff --git a/server/handler/link_auth.go b/server/handler/link_auth.go
@@ -94,7 +94,10 @@ func LinkAuth(w http.ResponseWriter, r *http.Request) {
 	// TODO 用户密码校验
 	err = dbdata.CheckUser(cr.Auth.Username, cr.Auth.Password, cr.GroupSelect)
 	if err != nil {
-		lockManager.LoginStatus.Store(loginStatusKey, false) // 记录登录失败状态
+		// lockManager.LoginStatus.Store(loginStatusKey, false) // 记录登录失败状态
+		hc := r.Context().Value(loginStatusKey).(*HttpContext)
+		hc.LoginStatus = false
+
 		base.Warn(err, r.RemoteAddr)
 		ua.Info = err.Error()
 		ua.Status = dbdata.UserAuthFail
@@ -119,7 +122,10 @@ func LinkAuth(w http.ResponseWriter, r *http.Request) {
 	}
 	// 用户otp验证
 	if base.Cfg.AuthAloneOtp && !v.DisableOtp {
-		lockManager.LoginStatus.Store(loginStatusKey, true) // 重置OTP验证计数
+		// lockManager.LoginStatus.Store(loginStatusKey, true) // 重置OTP验证计数
+		hc := r.Context().Value(loginStatusKey).(*HttpContext)
+		hc.LoginStatus = true
+
 		sessionID, err := GenerateSessionID()
 		if err != nil {
 			base.Error("Failed to generate session ID: ", err)

diff --git a/server/handler/link_auth_otp.go b/server/handler/link_auth_otp.go
@@ -109,7 +109,9 @@ func DeleteCookie(w http.ResponseWriter, name string) {
 	http.SetCookie(w, cookie)
 }
 func CreateSession(w http.ResponseWriter, r *http.Request, authSession *AuthSession) {
-	lockManager.LoginStatus.Store(loginStatusKey, true) // 更新登录成功状态
+	// lockManager.LoginStatus.Store(loginStatusKey, true) // 更新登录成功状态
+	hc := r.Context().Value(loginStatusKey).(*HttpContext)
+	hc.LoginStatus = true
 
 	cr := authSession.ClientRequest
 	ua := authSession.UserActLog
@@ -201,7 +203,9 @@ func LinkAuth_otp(w http.ResponseWriter, r *http.Request) {
 		// 	http.Error(w, "TooManyError, please login again", http.StatusBadRequest)
 		// 	return
 		// }
-		lockManager.LoginStatus.Store(loginStatusKey, false) // 记录登录失败状态
+		// lockManager.LoginStatus.Store(loginStatusKey, false) // 记录登录失败状态
+		hc := r.Context().Value(loginStatusKey).(*HttpContext)
+		hc.LoginStatus = false
 
 		base.Warn("OTP 动态码错误", username, r.RemoteAddr)
 		ua.Info = "OTP 动态码错误"