From bb1b161ae2feac235ecef9e63da3f662f5cc9f44 Mon Sep 17 00:00:00 2001
From: kirari04 <lev@saminskij.ch>
Date: Tue, 11 Jun 2024 22:04:30 +0200
Subject: [PATCH] clean up implementation

---
 src/mod/dynamicproxy/Server.go    |  9 +++++----
 src/mod/dynamicproxy/ratelimit.go | 10 +++++++++-
 src/mod/dynamicproxy/typedef.go   |  4 ++--
 3 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/src/mod/dynamicproxy/Server.go b/src/mod/dynamicproxy/Server.go
index 8756964..f89fa26 100644
--- a/src/mod/dynamicproxy/Server.go
+++ b/src/mod/dynamicproxy/Server.go
@@ -73,11 +73,12 @@ func (h *ProxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		}
 
 		// Rate Limit Check
-		// if sep.RequireBasicAuth {
-		if err := handleRateLimit(w, r, sep); err != nil {
-			return
+		if sep.RequireRateLimit {
+			err := h.handleRateLimitRouting(w, r, sep)
+			if err != nil {
+				return
+			}
 		}
-		// }
 
 		//Validate basic auth
 		if sep.RequireBasicAuth {
diff --git a/src/mod/dynamicproxy/ratelimit.go b/src/mod/dynamicproxy/ratelimit.go
index cdcb96c..0ac4e9f 100644
--- a/src/mod/dynamicproxy/ratelimit.go
+++ b/src/mod/dynamicproxy/ratelimit.go
@@ -50,6 +50,14 @@ func (t *IpTable) Clear() {
 
 var ipTable = IpTable{}
 
+func (h *ProxyHandler) handleRateLimitRouting(w http.ResponseWriter, r *http.Request, pe *ProxyEndpoint) error {
+	err := handleRateLimit(w, r, pe)
+	if err != nil {
+		h.logRequest(r, false, 429, "ratelimit", pe.Domain)
+	}
+	return err
+}
+
 func handleRateLimit(w http.ResponseWriter, r *http.Request, pe *ProxyEndpoint) error {
 	ip, _, err := net.SplitHostPort(r.RemoteAddr)
 	if err != nil {
@@ -60,7 +68,7 @@ func handleRateLimit(w http.ResponseWriter, r *http.Request, pe *ProxyEndpoint)
 
 	ipTable.Increment(ip)
 
-	if ipTable.Exceeded(ip, 10) {
+	if ipTable.Exceeded(ip, int64(pe.RateLimit)) {
 		w.WriteHeader(429)
 		return errors.New("rate limit exceeded")
 	}
diff --git a/src/mod/dynamicproxy/typedef.go b/src/mod/dynamicproxy/typedef.go
index 660e663..d4f9a8a 100644
--- a/src/mod/dynamicproxy/typedef.go
+++ b/src/mod/dynamicproxy/typedef.go
@@ -125,8 +125,8 @@ type ProxyEndpoint struct {
 	BasicAuthExceptionRules []*BasicAuthExceptionRule //Path to exclude in a basic auth enabled proxy target
 
 	// Rate Limiting
-	EnableRateLimiting bool
-	RateLimiting       int // Rate limit in requests per second
+	RequireRateLimit bool
+	RateLimit        int64 // Rate limit in requests per second
 
 	//Access Control
 	AccessFilterUUID string //Access filter ID