Removing Benchmark & Updated implementation

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")