From 95453431517ca3b4b52404a70d14016e9c4e266e Mon Sep 17 00:00:00 2001 From: kirari04 Date: Tue, 11 Jun 2024 16:56:59 +0200 Subject: [PATCH] Removing Benchmark & Updated implementation --- src/mod/dynamicproxy/ratelimit.go | 97 +++---------- .../dynamicproxy/ratelimit_benchmark_test.go | 128 ------------------ src/mod/dynamicproxy/ratelimitb.go | 79 ----------- src/mod/dynamicproxy/ratelimitc.go | 78 ----------- 4 files changed, 21 insertions(+), 361 deletions(-) delete mode 100644 src/mod/dynamicproxy/ratelimit_benchmark_test.go delete mode 100644 src/mod/dynamicproxy/ratelimitb.go delete mode 100644 src/mod/dynamicproxy/ratelimitc.go diff --git a/src/mod/dynamicproxy/ratelimit.go b/src/mod/dynamicproxy/ratelimit.go index e173fb7..cdcb96c 100644 --- a/src/mod/dynamicproxy/ratelimit.go +++ b/src/mod/dynamicproxy/ratelimit.go @@ -2,109 +2,55 @@ package dynamicproxy import ( "errors" + "log" "net" "net/http" "sync" + "sync/atomic" "time" - - "log" ) -/* - ratelimit.go - - This file handles the ratelimit on proxy endpoints - if RateLimit is set to true -*/ - -// idk what this was for -// 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, "host", pe.Domain) -// } -// return err -// } - +// IpTable is a rate limiter implementation using sync.Map with atomic int64 type IpTable struct { - sync.RWMutex - table map[string]*IpTableValue + table sync.Map } -// Get the ip from the table -func (t *IpTable) Get(ip string) (*IpTableValue, bool) { - t.RLock() - defer t.RUnlock() - v, ok := t.table[ip] - return v, ok -} - -// Clear the ip from the table -func (t *IpTable) Clear() { - t.Lock() - defer t.Unlock() - t.table = make(map[string]*IpTableValue) -} - -// Increment the count of requests for a given ip -// init ip in ipTable if not exists +// Increment the count of requests for a given IP func (t *IpTable) Increment(ip string) { - t.Lock() - defer t.Unlock() - v, ok := t.table[ip] - if !ok { - v = &IpTableValue{Count: 0, LastHit: time.Now()} - } - v.Count++ - t.table[ip] = v + v, _ := t.table.LoadOrStore(ip, new(int64)) + atomic.AddInt64(v.(*int64), 1) } -// Check if the ip is in the table and if it is, check if the count is less than the limit +// Check if the IP is in the table and if it is, check if the count is less than the limit func (t *IpTable) Exceeded(ip string, limit int64) bool { - t.RLock() - defer t.RUnlock() - v, ok := t.table[ip] + v, ok := t.table.Load(ip) if !ok { return false } - if v.Count < limit { - return false - } - return true + count := atomic.LoadInt64(v.(*int64)) + return count >= limit } -// Get the count of requests for a given ip -// returns 0 if ip is not in the table +// Get the count of requests for a given IP func (t *IpTable) GetCount(ip string) int64 { - t.RLock() - defer t.RUnlock() - v, ok := t.table[ip] + v, ok := t.table.Load(ip) if !ok { return 0 } - return v.Count + return atomic.LoadInt64(v.(*int64)) } -type IpTableValue struct { - Count int64 - LastHit time.Time +// Clear the IP table +func (t *IpTable) Clear() { + t.table.Range(func(key, value interface{}) bool { + t.table.Delete(key) + return true + }) } -var ipTable IpTable = IpTable{table: make(map[string]*IpTableValue)} +var ipTable = IpTable{} -// Handle rate limit logic -// do not write to http.ResponseWriter if err return is not nil (already handled by this function) func handleRateLimit(w http.ResponseWriter, r *http.Request, pe *ProxyEndpoint) error { - // if len(pe.BasicAuthExceptionRules) > 0 { - // //Check if the current path matches the exception rules - // for _, exceptionRule := range pe.BasicAuthExceptionRules { - // if strings.HasPrefix(r.RequestURI, exceptionRule.PathPrefix) { - // //This path is excluded from basic auth - // return nil - // } - // } - // } - ip, _, err := net.SplitHostPort(r.RemoteAddr) if err != nil { w.WriteHeader(500) @@ -114,7 +60,6 @@ func handleRateLimit(w http.ResponseWriter, r *http.Request, pe *ProxyEndpoint) ipTable.Increment(ip) - // if ipTable.Exceeded(ip, pe.RateLimit) { if ipTable.Exceeded(ip, 10) { w.WriteHeader(429) return errors.New("rate limit exceeded") diff --git a/src/mod/dynamicproxy/ratelimit_benchmark_test.go b/src/mod/dynamicproxy/ratelimit_benchmark_test.go deleted file mode 100644 index f19a40b..0000000 --- a/src/mod/dynamicproxy/ratelimit_benchmark_test.go +++ /dev/null @@ -1,128 +0,0 @@ -package dynamicproxy - -import ( - "net/http" - "net/http/httptest" - "sync" - "testing" -) - -func BenchmarkRateLimitLockMapInt64(b *testing.B) { - go InitRateLimit() - handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - handleRateLimit(w, r, &ProxyEndpoint{RateLimiting: 10}) - }) - - request := httptest.NewRequest("GET", "/", nil) - request.RemoteAddr = "192.168.1.1:1234" - - b.ResetTimer() - - for i := 0; i < b.N; i++ { - recorder := httptest.NewRecorder() - handler.ServeHTTP(recorder, request) - } -} - -func BenchmarkRateLimitSyncMapInt64(b *testing.B) { - go InitRateLimitSyncMapInt64() - handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - handleRateLimitSyncMapInt64(w, r, &ProxyEndpoint{RateLimiting: 10}) - }) - - request := httptest.NewRequest("GET", "/", nil) - request.RemoteAddr = "192.168.1.1:1234" - - b.ResetTimer() - - for i := 0; i < b.N; i++ { - recorder := httptest.NewRecorder() - handler.ServeHTTP(recorder, request) - } -} - -func BenchmarkRateLimitSyncMapAtomicInt64(b *testing.B) { - go InitRateLimitSyncMapAtomicInt64() - handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - handleRateLimitSyncMapAtomicInt64(w, r, &ProxyEndpoint{RateLimiting: 10}) - }) - - request := httptest.NewRequest("GET", "/", nil) - request.RemoteAddr = "192.168.1.1:1234" - - b.ResetTimer() - - for i := 0; i < b.N; i++ { - recorder := httptest.NewRecorder() - handler.ServeHTTP(recorder, request) - } -} - -func BenchmarkRateLimitLockMapInt64Concurrent(b *testing.B) { - go InitRateLimit() - handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - handleRateLimit(w, r, &ProxyEndpoint{RateLimiting: 10}) - }) - - request := httptest.NewRequest("GET", "/", nil) - request.RemoteAddr = "192.168.1.1:1234" - - b.ResetTimer() - - var wg sync.WaitGroup - for i := 0; i < b.N; i++ { - wg.Add(1) - go func() { - defer wg.Done() - recorder := httptest.NewRecorder() - handler.ServeHTTP(recorder, request) - }() - } - wg.Wait() -} - -func BenchmarkRateLimitSyncMapInt64Concurrent(b *testing.B) { - go InitRateLimitSyncMapInt64() - handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - handleRateLimitSyncMapInt64(w, r, &ProxyEndpoint{RateLimiting: 10}) - }) - - request := httptest.NewRequest("GET", "/", nil) - request.RemoteAddr = "192.168.1.1:1234" - - b.ResetTimer() - - var wg sync.WaitGroup - for i := 0; i < b.N; i++ { - wg.Add(1) - go func() { - defer wg.Done() - recorder := httptest.NewRecorder() - handler.ServeHTTP(recorder, request) - }() - } - wg.Wait() -} - -func BenchmarkRateLimitSyncMapAtomicInt64Concurrent(b *testing.B) { - go InitRateLimitSyncMapAtomicInt64() - handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - handleRateLimitSyncMapAtomicInt64(w, r, &ProxyEndpoint{RateLimiting: 10}) - }) - - request := httptest.NewRequest("GET", "/", nil) - request.RemoteAddr = "192.168.1.1:1234" - - b.ResetTimer() - - var wg sync.WaitGroup - for i := 0; i < b.N; i++ { - wg.Add(1) - go func() { - defer wg.Done() - recorder := httptest.NewRecorder() - handler.ServeHTTP(recorder, request) - }() - } - wg.Wait() -} diff --git a/src/mod/dynamicproxy/ratelimitb.go b/src/mod/dynamicproxy/ratelimitb.go deleted file mode 100644 index 48c749b..0000000 --- a/src/mod/dynamicproxy/ratelimitb.go +++ /dev/null @@ -1,79 +0,0 @@ -package dynamicproxy - -import ( - "errors" - "log" - "net" - "net/http" - "sync" - "time" -) - -// IpTableSyncMapInt64 is a rate limiter implementation using sync.Map with int64 -type IpTableSyncMapInt64 struct { - table sync.Map -} - -// Increment the count of requests for a given IP -func (t *IpTableSyncMapInt64) Increment(ip string) { - v, _ := t.table.LoadOrStore(ip, new(int64)) - count := v.(*int64) - *count++ -} - -// Check if the IP is in the table and if it is, check if the count is less than the limit -func (t *IpTableSyncMapInt64) Exceeded(ip string, limit int64) bool { - v, ok := t.table.Load(ip) - if !ok { - return false - } - count := v.(*int64) - return *count >= limit -} - -// Get the count of requests for a given IP -func (t *IpTableSyncMapInt64) GetCount(ip string) int64 { - v, ok := t.table.Load(ip) - if !ok { - return 0 - } - count := v.(*int64) - return *count -} - -// Clear the IP table -func (t *IpTableSyncMapInt64) Clear() { - t.table.Range(func(key, value interface{}) bool { - t.table.Delete(key) - return true - }) -} - -var ipTableSyncMapInt64 = IpTableSyncMapInt64{} - -func handleRateLimitSyncMapInt64(w http.ResponseWriter, r *http.Request, pe *ProxyEndpoint) error { - ip, _, err := net.SplitHostPort(r.RemoteAddr) - if err != nil { - w.WriteHeader(500) - log.Println("Error resolving remote address", r.RemoteAddr, err) - return errors.New("internal server error") - } - - ipTableSyncMapInt64.Increment(ip) - - if ipTableSyncMapInt64.Exceeded(ip, 10) { - w.WriteHeader(429) - return errors.New("rate limit exceeded") - } - - // log.Println("Rate limit check", ip, ipTableSyncMapInt64.GetCount(ip)) - - return nil -} - -func InitRateLimitSyncMapInt64() { - for { - ipTableSyncMapInt64.Clear() - time.Sleep(time.Second) - } -} diff --git a/src/mod/dynamicproxy/ratelimitc.go b/src/mod/dynamicproxy/ratelimitc.go deleted file mode 100644 index 88d67ea..0000000 --- a/src/mod/dynamicproxy/ratelimitc.go +++ /dev/null @@ -1,78 +0,0 @@ -package dynamicproxy - -import ( - "errors" - "log" - "net" - "net/http" - "sync" - "sync/atomic" - "time" -) - -// IpTableSyncMapAtomicInt64 is a rate limiter implementation using sync.Map with atomic int64 -type IpTableSyncMapAtomicInt64 struct { - table sync.Map -} - -// Increment the count of requests for a given IP -func (t *IpTableSyncMapAtomicInt64) Increment(ip string) { - v, _ := t.table.LoadOrStore(ip, new(int64)) - atomic.AddInt64(v.(*int64), 1) -} - -// Check if the IP is in the table and if it is, check if the count is less than the limit -func (t *IpTableSyncMapAtomicInt64) Exceeded(ip string, limit int64) bool { - v, ok := t.table.Load(ip) - if !ok { - return false - } - count := atomic.LoadInt64(v.(*int64)) - return count >= limit -} - -// Get the count of requests for a given IP -func (t *IpTableSyncMapAtomicInt64) GetCount(ip string) int64 { - v, ok := t.table.Load(ip) - if !ok { - return 0 - } - return atomic.LoadInt64(v.(*int64)) -} - -// Clear the IP table -func (t *IpTableSyncMapAtomicInt64) Clear() { - t.table.Range(func(key, value interface{}) bool { - t.table.Delete(key) - return true - }) -} - -var ipTableSyncMapAtomicInt64 = IpTableSyncMapAtomicInt64{} - -func handleRateLimitSyncMapAtomicInt64(w http.ResponseWriter, r *http.Request, pe *ProxyEndpoint) error { - ip, _, err := net.SplitHostPort(r.RemoteAddr) - if err != nil { - w.WriteHeader(500) - log.Println("Error resolving remote address", r.RemoteAddr, err) - return errors.New("internal server error") - } - - ipTableSyncMapAtomicInt64.Increment(ip) - - if ipTableSyncMapAtomicInt64.Exceeded(ip, 10) { - w.WriteHeader(429) - return errors.New("rate limit exceeded") - } - - // log.Println("Rate limit check", ip, ipTableSyncMapAtomicInt64.GetCount(ip)) - - return nil -} - -func InitRateLimitSyncMapAtomicInt64() { - for { - ipTableSyncMapAtomicInt64.Clear() - time.Sleep(time.Second) - } -}