From 1a977b2bce74defc40a9254832a4a93e2169ce2b Mon Sep 17 00:00:00 2001
From: Roi Feng <37480123+Rayzggz@users.noreply.github.com>
Date: Thu, 24 Apr 2025 21:02:19 -0400
Subject: [PATCH] feat: Block HTTP FLOOD

---
 internal/check/HTTPFlood.go |  6 ++----
 internal/server/checker.go  | 24 ++++++++++++++++++++++++
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/internal/check/HTTPFlood.go b/internal/check/HTTPFlood.go
index d088ff2..45f1d90 100644
--- a/internal/check/HTTPFlood.go
+++ b/internal/check/HTTPFlood.go
@@ -17,8 +17,7 @@ func HTTPFlood(reqData dataType.UserRequest, ruleSet *config.RuleSet, decision *
 	for window, limit := range ruleSet.HTTPFloodRule.HTTPFloodSpeedLimit {
 		if sharedMem.HTTPFloodSpeedLimitCounter.Query(ipKey, window) > limit {
 			log.Printf("HTTPFlood rate limit exceeded: IP %s, window %d, limit %d", ipKey, window, limit)
-			//decision.SetResponse(action.Done, []byte("403"), nil)
-			decision.Set(action.Continue)
+			decision.SetCode(action.Done, []byte("429"))
 			return
 		}
 	}
@@ -26,8 +25,7 @@ func HTTPFlood(reqData dataType.UserRequest, ruleSet *config.RuleSet, decision *
 	for window, limit := range ruleSet.HTTPFloodRule.HTTPFloodSameURILimit {
 		if sharedMem.HTTPFloodSameURILimitCounter.Query(uriKey, window) > limit {
 			log.Printf("HTTPFlood URI rate limit exceeded: IP %s, URI %s, window %d, limit %d", ipKey, reqData.Uri, window, limit)
-			//decision.SetResponse(action.Done, []byte("403"), nil)
-			decision.Set(action.Continue)
+			decision.SetCode(action.Done, []byte("429"))
 			return
 		}
 	}
diff --git a/internal/server/checker.go b/internal/server/checker.go
index ffb835e..7039f2c 100644
--- a/internal/server/checker.go
+++ b/internal/server/checker.go
@@ -81,6 +81,30 @@ func CheckMain(w http.ResponseWriter, userRequestData dataType.UserRequest, rule
 			return
 		}
 
+	} else if bytes.Compare(decision.HTTPCode, []byte("429")) == 0 {
+		tpl, err := template.ParseFiles(cfg.ErrorPage + "/429.html")
+		if err != nil {
+			log.Printf("Error template: %v", err)
+			http.Error(w, "500 - Internal Server Error", http.StatusInternalServerError)
+			return
+		}
+		data := struct {
+			EdgeTag   string
+			ConnectIP string
+			Date      string
+		}{
+			EdgeTag:   cfg.NodeName,
+			ConnectIP: userRequestData.RemoteIP,
+			Date:      time.Now().Format("2006-01-02 15:04:05"),
+		}
+		w.WriteHeader(http.StatusTooManyRequests)
+		w.Header().Set("Content-Type", "text/html; charset=utf-8")
+		if err = tpl.Execute(w, data); err != nil {
+			log.Printf("Error template: %v", err)
+			http.Error(w, "500 - Internal Server Error", http.StatusInternalServerError)
+			return
+		}
+
 	} else {
 		//should never happen
 		log.Printf("Error access in wrong state: %v", decision)