diff --git a/src/mod/dynamicproxy/ratelimit.go b/src/mod/dynamicproxy/ratelimit.go index e9884d7..e173fb7 100644 --- a/src/mod/dynamicproxy/ratelimit.go +++ b/src/mod/dynamicproxy/ratelimit.go @@ -120,7 +120,7 @@ func handleRateLimit(w http.ResponseWriter, r *http.Request, pe *ProxyEndpoint) return errors.New("rate limit exceeded") } - log.Println("Rate limit check", ip, ipTable.GetCount(ip)) + // log.Println("Rate limit check", ip, ipTable.GetCount(ip)) return nil } diff --git a/src/mod/dynamicproxy/ratelimit_benchmark_test.go b/src/mod/dynamicproxy/ratelimit_benchmark_test.go new file mode 100644 index 0000000..ed5b846 --- /dev/null +++ b/src/mod/dynamicproxy/ratelimit_benchmark_test.go @@ -0,0 +1,92 @@ +package dynamicproxy + +import ( + "net/http" + "net/http/httptest" + "sync" + "testing" +) + +func BenchmarkRateLimitSyncMapInt64(b *testing.B) { + ipTableSyncMapInt64 = IpTableSyncMapInt64{} // Reset the table + + 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) { + ipTableSyncMapAtomicInt64 = IpTableSyncMapAtomicInt64{} // Reset the table + + 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 BenchmarkRateLimitSyncMapInt64Concurrent(b *testing.B) { + ipTableSyncMapInt64 = IpTableSyncMapInt64{} // Reset the table + + 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) { + ipTableSyncMapAtomicInt64 = IpTableSyncMapAtomicInt64{} // Reset the table + + 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 new file mode 100644 index 0000000..48c749b --- /dev/null +++ b/src/mod/dynamicproxy/ratelimitb.go @@ -0,0 +1,79 @@ +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 new file mode 100644 index 0000000..88d67ea --- /dev/null +++ b/src/mod/dynamicproxy/ratelimitc.go @@ -0,0 +1,78 @@ +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) + } +} diff --git a/src/mod/dynamicproxy/typedef.go b/src/mod/dynamicproxy/typedef.go index 3386aa6..660e663 100644 --- a/src/mod/dynamicproxy/typedef.go +++ b/src/mod/dynamicproxy/typedef.go @@ -124,6 +124,10 @@ type ProxyEndpoint struct { BasicAuthCredentials []*BasicAuthCredentials //Basic auth credentials BasicAuthExceptionRules []*BasicAuthExceptionRule //Path to exclude in a basic auth enabled proxy target + // Rate Limiting + EnableRateLimiting bool + RateLimiting int // Rate limit in requests per second + //Access Control AccessFilterUUID string //Access filter ID