feat: waiting room check api

2025-06-16 12:11:22 +08:00 · 2025-06-09 18:48:48 -04:00
parent fcb08478d2
commit aa1e760a79
2 changed files with 162 additions and 2 deletions
--- a/internal/server/checker.go
+++ b/internal/server/checker.go
@ -10,6 +10,7 @@ import (
 	"server_torii/internal/config"
 	"server_torii/internal/dataType"
 	"server_torii/internal/utils"
+	"strings"
 	"time"
 )

@ -115,8 +116,16 @@ func CheckMain(w http.ResponseWriter, userRequestData dataType.UserRequest, rule
 			return
 		}

-		sessionID := string(decision.ResponseData)
-		position, totalQueue := sharedMem.WaitingRoom.GetQueueInfo(sessionID, userRequestData, ruleSet.CAPTCHARule.SecretKey)
+		responseData := string(decision.ResponseData)
+		parts := strings.Split(responseData, "|")
+		if len(parts) < 2 {
+			utils.LogError(userRequestData, fmt.Sprintf("Invalid responseData format: %s", responseData), "CheckMain")
+			http.Error(w, "500 - Internal Server Error", http.StatusInternalServerError)
+			return
+		}
+		sessionID := parts[0]
+		userKey := parts[1]
+		position, totalQueue := sharedMem.WaitingRoom.GetQueueInfo(sessionID, userKey)

 		data := struct {
 			EdgeTag       string
--- a/internal/server/torii.go
+++ b/internal/server/torii.go
@ -2,6 +2,7 @@ package server

 import (
 	"bytes"
+	"encoding/json"
 	"html/template"
 	"net/http"
 	"server_torii/internal/action"
@ -20,6 +21,12 @@ func CheckTorii(w http.ResponseWriter, r *http.Request, reqData dataType.UserReq
 		check.CheckCaptcha(r, reqData, ruleSet, decision)
 	} else if reqData.Uri == cfg.WebPath+"/health_check" {
 		decision.SetResponse(action.Done, []byte("200"), []byte("ok"))
+	} else if reqData.Uri == cfg.WebPath+"/waiting_room/status" {
+		handleWaitingRoomStatus(w, r, reqData, ruleSet, sharedMem)
+		return
+	} else if reqData.Uri == cfg.WebPath+"/waiting_room/join" {
+		handleWaitingRoomJoin(w, r, reqData, ruleSet, sharedMem)
+		return
 	}
 	if bytes.Compare(decision.HTTPCode, []byte("200")) == 0 {
 		if bytes.Compare(decision.ResponseData, []byte("ok")) == 0 {
@ -88,3 +95,147 @@ func CheckTorii(w http.ResponseWriter, r *http.Request, reqData dataType.UserReq
 		}
 	}
 }
+
+func handleWaitingRoomStatus(w http.ResponseWriter, r *http.Request, reqData dataType.UserRequest, ruleSet *config.RuleSet, sharedMem *dataType.SharedMemory) {
+	if !ruleSet.WaitingRoomRule.Enabled {
+		w.WriteHeader(http.StatusNotFound)
+		err := json.NewEncoder(w).Encode(map[string]interface{}{
+			"error": "Waiting room not enabled",
+		})
+		if err != nil {
+			return
+		}
+		return
+	}
+
+	sessionID := reqData.ToriiSessionID
+	userKey := check.GenerateUserKey(reqData)
+
+	var position, totalQueue int
+	var canEnter bool
+
+	if sessionID != "" {
+		validSessionID := check.VerifyWaitingRoomSessionID(sessionID, reqData, ruleSet.CAPTCHARule.SecretKey, ruleSet.WaitingRoomRule.SessionTimeout)
+		if validSessionID {
+			canEnter, _ = sharedMem.WaitingRoom.CanEnterSite(sessionID, userKey)
+			position, totalQueue = sharedMem.WaitingRoom.GetQueueInfo(sessionID, userKey)
+		}
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(http.StatusOK)
+	err := json.NewEncoder(w).Encode(map[string]interface{}{
+		"position":   position,
+		"totalQueue": totalQueue,
+		"canEnter":   canEnter,
+		"sessionID":  sessionID,
+		"inQueue":    sessionID != "" && position > 0,
+	})
+	if err != nil {
+		return
+	}
+}
+
+func handleWaitingRoomJoin(w http.ResponseWriter, r *http.Request, reqData dataType.UserRequest, ruleSet *config.RuleSet, sharedMem *dataType.SharedMemory) {
+	if !ruleSet.WaitingRoomRule.Enabled {
+		w.WriteHeader(http.StatusNotFound)
+		err := json.NewEncoder(w).Encode(map[string]interface{}{
+			"error": "Waiting room not enabled",
+		})
+		if err != nil {
+			return
+		}
+		return
+	}
+
+	if r.Method != http.MethodPost {
+		w.WriteHeader(http.StatusMethodNotAllowed)
+		err := json.NewEncoder(w).Encode(map[string]interface{}{
+			"error": "Method not allowed",
+		})
+		if err != nil {
+			return
+		}
+		return
+	}
+
+	userKey := check.GenerateUserKey(reqData)
+	sessionID := reqData.ToriiSessionID
+
+	if sessionID != "" {
+		validSessionID := check.VerifyWaitingRoomSessionID(sessionID, reqData, ruleSet.CAPTCHARule.SecretKey, ruleSet.WaitingRoomRule.SessionTimeout)
+		if validSessionID {
+			canEnter, _ := sharedMem.WaitingRoom.CanEnterSite(sessionID, userKey)
+			if canEnter {
+				sharedMem.WaitingRoom.AddToActiveSession(sessionID, userKey)
+				w.Header().Set("Content-Type", "application/json")
+				w.WriteHeader(http.StatusOK)
+				err := json.NewEncoder(w).Encode(map[string]interface{}{
+					"success":  true,
+					"canEnter": true,
+					"message":  "可以进入网站",
+				})
+				if err != nil {
+					return
+				}
+				return
+			}
+			position, totalQueue := sharedMem.WaitingRoom.GetQueueInfo(sessionID, userKey)
+			w.Header().Set("Content-Type", "application/json")
+			w.WriteHeader(http.StatusOK)
+			err := json.NewEncoder(w).Encode(map[string]interface{}{
+				"success":    true,
+				"canEnter":   false,
+				"position":   position,
+				"totalQueue": totalQueue,
+				"message":    "已在队列中，请等待",
+			})
+			if err != nil {
+				return
+			}
+			return
+		}
+	}
+
+	newSessionID := genWaitingRoomSessionID(reqData, ruleSet.CAPTCHARule.SecretKey)
+
+	canEnter, _ := sharedMem.WaitingRoom.CanEnterSite("", userKey)
+	if canEnter {
+		sharedMem.WaitingRoom.AddToActiveSession(newSessionID, userKey)
+		w.Header().Set("Set-Cookie", "__torii_session_id="+newSessionID+"; Path=/; Max-Age=86400; Priority=High; HttpOnly;")
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusOK)
+		err := json.NewEncoder(w).Encode(map[string]interface{}{
+			"success":   true,
+			"canEnter":  true,
+			"sessionID": newSessionID,
+			"message":   "成功加入，可以进入网站",
+		})
+		if err != nil {
+			return
+		}
+		return
+	}
+
+	sharedMem.WaitingRoom.AddToQueue(newSessionID, userKey)
+	position, totalQueue := sharedMem.WaitingRoom.GetQueueInfo(newSessionID, userKey)
+
+	w.Header().Set("Set-Cookie", "__torii_session_id="+newSessionID+"; Path=/; Max-Age=86400; Priority=High; HttpOnly;")
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(http.StatusOK)
+	err := json.NewEncoder(w).Encode(map[string]interface{}{
+		"success":    true,
+		"canEnter":   false,
+		"sessionID":  newSessionID,
+		"position":   position,
+		"totalQueue": totalQueue,
+		"message":    "已加入排队",
+	})
+	if err != nil {
+		return
+	}
+}
+
+func genWaitingRoomSessionID(reqData dataType.UserRequest, secretKey string) string {
+	return check.GenWaitingRoomSessionID(reqData, secretKey)
+}