feat: config and checker of waiting room

config_example/rules/Server.yml

@@ -14,4 +14,8 @@ VerifyBot:
   verify_baidu_bot: true
   verify_yandex_bot: true
   verify_sogou_bot: true
-  verify_apple_bot: true
+  verify_apple_bot: true
+WaitingRoom:
+  enabled: false
+  max_concurrent_user: 1000
+  session_timeout: 1800

internal/config/config.go

@@ -66,20 +66,22 @@ func LoadMainConfig(basePath string) (*MainConfig, error) {
 
 // RuleSet stores all rules
 type RuleSet struct {
-	IPAllowTrie   *dataType.TrieNode
+	IPAllowTrie     *dataType.TrieNode
-	IPBlockTrie   *dataType.TrieNode
+	IPBlockTrie     *dataType.TrieNode
-	URLAllowList  *dataType.URLRuleList
+	URLAllowList    *dataType.URLRuleList
-	URLBlockList  *dataType.URLRuleList
+	URLBlockList    *dataType.URLRuleList
-	CAPTCHARule   *dataType.CaptchaRule
+	CAPTCHARule     *dataType.CaptchaRule
-	VerifyBotRule *dataType.VerifyBotRule
+	VerifyBotRule   *dataType.VerifyBotRule
-	HTTPFloodRule *dataType.HTTPFloodRule
+	HTTPFloodRule   *dataType.HTTPFloodRule
+	WaitingRoomRule *dataType.WaitingRoomRule
 }
 
 // ruleSetWrapper
 type ruleSetWrapper struct {
-	CAPTCHARule   *dataType.CaptchaRule   `yaml:"CAPTCHA"`
+	CAPTCHARule     *dataType.CaptchaRule     `yaml:"CAPTCHA"`
-	VerifyBotRule *dataType.VerifyBotRule `yaml:"VerifyBot"`
+	VerifyBotRule   *dataType.VerifyBotRule   `yaml:"VerifyBot"`
-	HTTPFloodRule httpFloodRuleWrapper    `yaml:"HTTPFlood"`
+	HTTPFloodRule   httpFloodRuleWrapper      `yaml:"HTTPFlood"`
+	WaitingRoomRule *dataType.WaitingRoomRule `yaml:"WaitingRoom"`
 }
 
 type httpFloodRuleWrapper struct {
@@ -90,13 +92,14 @@ type httpFloodRuleWrapper struct {
 // LoadRules Load all rules from the specified path
 func LoadRules(rulePath string) (*RuleSet, error) {
 	rs := RuleSet{
-		IPAllowTrie:   &dataType.TrieNode{},
+		IPAllowTrie:     &dataType.TrieNode{},
-		IPBlockTrie:   &dataType.TrieNode{},
+		IPBlockTrie:     &dataType.TrieNode{},
-		URLAllowList:  &dataType.URLRuleList{},
+		URLAllowList:    &dataType.URLRuleList{},
-		URLBlockList:  &dataType.URLRuleList{},
+		URLBlockList:    &dataType.URLRuleList{},
-		CAPTCHARule:   &dataType.CaptchaRule{},
+		CAPTCHARule:     &dataType.CaptchaRule{},
-		VerifyBotRule: &dataType.VerifyBotRule{},
+		VerifyBotRule:   &dataType.VerifyBotRule{},
-		HTTPFloodRule: &dataType.HTTPFloodRule{},
+		HTTPFloodRule:   &dataType.HTTPFloodRule{},
+		WaitingRoomRule: &dataType.WaitingRoomRule{},
 	}
 
 	// Load IP Allow List
@@ -149,6 +152,9 @@ func loadServerRules(YAMLFile string, rs RuleSet) (*RuleSet, error) {
 
 	*rs.CAPTCHARule = *wrapper.CAPTCHARule
 	*rs.VerifyBotRule = *wrapper.VerifyBotRule
+	if wrapper.WaitingRoomRule != nil {
+		*rs.WaitingRoomRule = *wrapper.WaitingRoomRule
+	}
 
 	rs.HTTPFloodRule.HTTPFloodSpeedLimit = make(map[int64]int64)
 	rs.HTTPFloodRule.HTTPFloodSameURILimit = make(map[int64]int64)

internal/dataType/type.go

@@ -31,7 +31,14 @@ type HTTPFloodRule struct {
 	HTTPFloodSameURILimit map[int64]int64
 }
 
+type WaitingRoomRule struct {
+	Enabled           bool  `yaml:"enabled"`
+	MaxConcurrentUser int64 `yaml:"max_concurrent_user"`
+	SessionTimeout    int64 `yaml:"session_timeout"`
+}
+
 type SharedMemory struct {
 	HTTPFloodSpeedLimitCounter   *Counter
 	HTTPFloodSameURILimitCounter *Counter
+	WaitingRoom                  *WaitingRoom
 }

internal/server/checker.go

@@ -25,6 +25,7 @@ func CheckMain(w http.ResponseWriter, userRequestData dataType.UserRequest, rule
 	checkFuncs = append(checkFuncs, check.URLBlockList)
 	checkFuncs = append(checkFuncs, check.VerifyBot)
 	checkFuncs = append(checkFuncs, check.HTTPFlood)
+	checkFuncs = append(checkFuncs, check.WaitingRoom)
 	checkFuncs = append(checkFuncs, check.Captcha)
 
 	for _, checkFunc := range checkFuncs {
@@ -106,6 +107,40 @@ func CheckMain(w http.ResponseWriter, userRequestData dataType.UserRequest, rule
 			return
 		}
 
+	} else if bytes.Compare(decision.HTTPCode, []byte("WAITING_ROOM")) == 0 {
+		tpl, err := template.ParseFiles(cfg.ErrorPage + "/waiting_room.html")
+		if err != nil {
+			utils.LogError(userRequestData, fmt.Sprintf("Error parsing template: %v", err), "CheckMain")
+			http.Error(w, "500 - Internal Server Error", http.StatusInternalServerError)
+			return
+		}
+
+		sessionID := string(decision.ResponseData)
+		position, totalQueue := sharedMem.WaitingRoom.GetQueueInfo(sessionID, userRequestData, ruleSet.CAPTCHARule.SecretKey)
+
+		data := struct {
+			EdgeTag       string
+			ConnectIP     string
+			Date          string
+			QueuePosition int
+			TotalQueue    int
+		}{
+			EdgeTag:       cfg.NodeName,
+			ConnectIP:     userRequestData.RemoteIP,
+			Date:          time.Now().Format("2006-01-02 15:04:05"),
+			QueuePosition: position,
+			TotalQueue:    totalQueue,
+		}
+
+		w.Header().Set("Set-Cookie", "__torii_session_id="+sessionID+"; Path=/; Max-Age=86400; Priority=High; HttpOnly;")
+		w.Header().Set("Content-Type", "text/html; charset=utf-8")
+		w.WriteHeader(http.StatusServiceUnavailable)
+		if err = tpl.Execute(w, data); err != nil {
+			utils.LogError(userRequestData, fmt.Sprintf("Error executing template: %v", err), "CheckMain")
+			http.Error(w, "500 - Internal Server Error", http.StatusInternalServerError)
+			return
+		}
+
 	} else {
 		//should never happen
 		utils.LogError(userRequestData, fmt.Sprintf("Error access in wrong state: %v", decision), "CheckMain")

main.go

@@ -58,6 +58,7 @@ func main() {
 	sharedMem := &dataType.SharedMemory{
 		HTTPFloodSpeedLimitCounter:   dataType.NewCounter(max(runtime.NumCPU()*8, 16), utils.FindMaxRateTime(ruleSet.HTTPFloodRule.HTTPFloodSpeedLimit)),
 		HTTPFloodSameURILimitCounter: dataType.NewCounter(max(runtime.NumCPU()*8, 16), utils.FindMaxRateTime(ruleSet.HTTPFloodRule.HTTPFloodSameURILimit)),
+		WaitingRoom:                  dataType.NewWaitingRoom(ruleSet.WaitingRoomRule.MaxConcurrentUser, ruleSet.WaitingRoomRule.SessionTimeout),
 	}
 
 	//GC