diff --git a/config_example/rules/Server.yml b/config_example/rules/Server.yml index dae22ee..b8696f4 100644 --- a/config_example/rules/Server.yml +++ b/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 \ No newline at end of file + verify_apple_bot: true +WaitingRoom: + enabled: false + max_concurrent_user: 1000 + session_timeout: 1800 \ No newline at end of file diff --git a/internal/config/config.go b/internal/config/config.go index 3b70c0f..0e584b5 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -66,20 +66,22 @@ func LoadMainConfig(basePath string) (*MainConfig, error) { // RuleSet stores all rules type RuleSet struct { - IPAllowTrie *dataType.TrieNode - IPBlockTrie *dataType.TrieNode - URLAllowList *dataType.URLRuleList - URLBlockList *dataType.URLRuleList - CAPTCHARule *dataType.CaptchaRule - VerifyBotRule *dataType.VerifyBotRule - HTTPFloodRule *dataType.HTTPFloodRule + IPAllowTrie *dataType.TrieNode + IPBlockTrie *dataType.TrieNode + URLAllowList *dataType.URLRuleList + URLBlockList *dataType.URLRuleList + CAPTCHARule *dataType.CaptchaRule + VerifyBotRule *dataType.VerifyBotRule + HTTPFloodRule *dataType.HTTPFloodRule + WaitingRoomRule *dataType.WaitingRoomRule } // ruleSetWrapper type ruleSetWrapper struct { - CAPTCHARule *dataType.CaptchaRule `yaml:"CAPTCHA"` - VerifyBotRule *dataType.VerifyBotRule `yaml:"VerifyBot"` - HTTPFloodRule httpFloodRuleWrapper `yaml:"HTTPFlood"` + CAPTCHARule *dataType.CaptchaRule `yaml:"CAPTCHA"` + VerifyBotRule *dataType.VerifyBotRule `yaml:"VerifyBot"` + 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{}, - IPBlockTrie: &dataType.TrieNode{}, - URLAllowList: &dataType.URLRuleList{}, - URLBlockList: &dataType.URLRuleList{}, - CAPTCHARule: &dataType.CaptchaRule{}, - VerifyBotRule: &dataType.VerifyBotRule{}, - HTTPFloodRule: &dataType.HTTPFloodRule{}, + IPAllowTrie: &dataType.TrieNode{}, + IPBlockTrie: &dataType.TrieNode{}, + URLAllowList: &dataType.URLRuleList{}, + URLBlockList: &dataType.URLRuleList{}, + CAPTCHARule: &dataType.CaptchaRule{}, + VerifyBotRule: &dataType.VerifyBotRule{}, + 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) diff --git a/internal/dataType/type.go b/internal/dataType/type.go index 5c77443..a159166 100644 --- a/internal/dataType/type.go +++ b/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 } diff --git a/internal/server/checker.go b/internal/server/checker.go index 679c39b..6136f20 100644 --- a/internal/server/checker.go +++ b/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") diff --git a/main.go b/main.go index a9a4302..82052a8 100644 --- a/main.go +++ b/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