update ratelimit

README.md

+# RPC Proxy
+
+This is a simple RPC proxy server that forwards requests to an Ethereum RPC backend. It includes features like whitelisting, blacklisting, and rate limiting.
+
+## Environment Variables
+
+The following environment variables can be used to configure the proxy:
+
+### General Configuration
+
+- `ETH_RPC_BACKEND`: The URL of the real Ethereum RPC backend to forward requests to.
+- `MYSQL_DSN`: The Data Source Name (DSN) for the MySQL database connection (e.g., `user:password@tcp(127.0.0.1:3306)/dbname`).
+- `BLACKLIST_CONTRACT_ADDR`: The address of the blacklist contract on the Ethereum network.
+- `WHITELIST_FILE`: The path to the file containing whitelisted IP addresses or domains.
+- `BLACKLIST_CACHE_CLEANUP_INTERVAL`: The interval for cleaning up the blacklist cache (default: `5m`).
+
+### Rate Limiting Configuration
+
+The proxy implements a two-stage rate limiting mechanism: **Normal Mode** and **Penalty Mode**.
+
+#### Normal Mode
+Applies to all non-whitelisted IPs by default.
+
+- `RATE_LIMIT_NORMAL`: The number of requests allowed per 15 minutes for normal usage (default: `1000`).
+- `RATE_BURST_NORMAL`: The burst size for normal usage (default: `20`).
+
+#### Penalty Mode
+Applies to IPs that frequently exceed the rate limit.
+
+- `RATE_LIMIT_PENALTY`: The number of requests allowed per 15 minutes for penalized IPs (default: `100`).
+- `RATE_BURST_PENALTY`: The burst size for penalized IPs (default: `5`).
+- `PENALTY_DURATION`: The duration an IP stays in penalty mode (default: `20m`).
+- `PENALTY_TRIGGER_COUNT`: The number of consecutive rate limit violations required to trigger penalty mode (default: `5`).
+
+## Usage
+
+1. Set the required environment variables.
+2. Run the proxy server.
+3. The server listens on port `8545`.
+
+## Whitelist File Format
+
+The whitelist file should contain one entry per line. Lines starting with `#` or `//` are ignored.
+Wildcards `*` are supported.
+
+Example:
+```
+127.0.0.1
+example.com
+*.example.com
+```
+

main.go

@@ -82,13 +82,22 @@ var (
 	// Rate limiting
 	visitors   = make(map[string]*visitor)
 	visitorsMu sync.Mutex
-	rateLimit  rate.Limit
-	rateBurst  int
+
+	// Rate limit settings
+	normalLimit     rate.Limit
+	normalBurst     int
+	penaltyLimit    rate.Limit
+	penaltyBurst    int
+	penaltyDuration time.Duration
+	penaltyTrigger  int
 )
 
 type visitor struct {
 	limiter       *rate.Limiter
 	lastSeen      time.Time
+	punishedUntil time.Time
+	blockedCount  int
+	lastBlockTime time.Time
 }
 
 func main() {
@@ -152,23 +161,52 @@ func main() {
 	startBlacklistCacheJanitor(cleanupInterval)
 
 	// Initialize rate limiter settings
-	// Default: 500 requests per 15 minutes
-	// rate.Every calculates the Limit (events/sec) for a given interval between events.
-	// Interval = 15 minutes / 500 requests
-	limit := rate.Every(15 * time.Minute / 500)
-	if s := os.Getenv("RATE_LIMIT"); s != "" {
+	// Normal Mode: High capacity for normal usage
+	nLimit := 1000.0
+	if s := os.Getenv("RATE_LIMIT_NORMAL"); s != "" {
+		if val, err := strconv.ParseFloat(s, 64); err == nil {
+			nLimit = val
+		}
+	}
+	normalLimit = rate.Every(15 * time.Minute / time.Duration(nLimit))
+
+	normalBurst = 20 // default burst
+	if s := os.Getenv("RATE_BURST_NORMAL"); s != "" {
+		if val, err := strconv.Atoi(s); err == nil {
+			normalBurst = val
+		}
+	}
+
+	// Penalty Mode: Strict limit for abusive IPs
+	pLimit := 100.0
+	if s := os.Getenv("RATE_LIMIT_PENALTY"); s != "" {
 		if val, err := strconv.ParseFloat(s, 64); err == nil {
-			limit = rate.Limit(val)
+			pLimit = val
 		}
 	}
-	rateLimit = limit
+	penaltyLimit = rate.Every(15 * time.Minute / time.Duration(pLimit))
 
-	rateBurst = 50 // default burst
-	if s := os.Getenv("RATE_BURST"); s != "" {
+	penaltyBurst = 5 // default burst
+	if s := os.Getenv("RATE_BURST_PENALTY"); s != "" {
 		if val, err := strconv.Atoi(s); err == nil {
-			rateBurst = val
+			penaltyBurst = val
+		}
+	}
+
+	penaltyDuration = 20 * time.Minute
+	if s := os.Getenv("PENALTY_DURATION"); s != "" {
+		if d, err := time.ParseDuration(s); err == nil {
+			penaltyDuration = d
+		}
 	}
+
+	penaltyTrigger = 5
+	if s := os.Getenv("PENALTY_TRIGGER_COUNT"); s != "" {
+		if val, err := strconv.Atoi(s); err == nil {
+			penaltyTrigger = val
 		}
+	}
+
 	startRateLimitCleaner()
 
 	http.HandleFunc("/", proxyHandler)
@@ -208,8 +246,8 @@ func proxyHandler(w http.ResponseWriter, r *http.Request) {
 	}
 
 	// Rate limiting for non-whitelisted requests
-	limiter := getVisitor(ip)
-	if !limiter.Allow() {
+	// ip is already extracted above
+	if !checkRateLimit(ip) {
 		http.Error(w, "Too Many Requests", http.StatusTooManyRequests)
 		log.Printf("rate limit exceeded for IP: %s", ip)
 		return
@@ -605,18 +643,52 @@ func startWhitelistWatcher(path string) {
 	}()
 }
 
-func getVisitor(ip string) *rate.Limiter {
+func checkRateLimit(ip string) bool {
 	visitorsMu.Lock()
 	defer visitorsMu.Unlock()
 
 	v, exists := visitors[ip]
 	if !exists {
-		limiter := rate.NewLimiter(rateLimit, rateBurst)
-		v = &visitor{limiter, time.Now()}
+		limiter := rate.NewLimiter(normalLimit, normalBurst)
+		v = &visitor{limiter: limiter, lastSeen: time.Now()}
 		visitors[ip] = v
 	}
 	v.lastSeen = time.Now()
-	return v.limiter
+
+	// Check if penalty expired
+	if !v.punishedUntil.IsZero() {
+		if time.Now().After(v.punishedUntil) {
+			// Penalty expired, restore normal settings
+			v.limiter.SetLimit(normalLimit)
+			v.limiter.SetBurst(normalBurst)
+			v.punishedUntil = time.Time{}
+			v.blockedCount = 0
+			log.Printf("IP %s penalty expired, restoring normal limits", ip)
+		}
+	}
+
+	allowed := v.limiter.Allow()
+
+	if !allowed {
+		v.blockedCount++
+		v.lastBlockTime = time.Now()
+
+		// Only trigger penalty if not already punished
+		if v.punishedUntil.IsZero() && v.blockedCount >= penaltyTrigger {
+			v.punishedUntil = time.Now().Add(penaltyDuration)
+			v.limiter.SetLimit(penaltyLimit)
+			v.limiter.SetBurst(penaltyBurst)
+			log.Printf("IP %s entered penalty mode (high traffic), limited to %v rps until %v", ip, penaltyLimit, v.punishedUntil)
+		}
+	} else {
+		// If allowed and not punished, maybe decay blockedCount?
+		// If it's been a while since last block, reset counter to avoid long-term accumulation
+		if v.punishedUntil.IsZero() && v.blockedCount > 0 && time.Since(v.lastBlockTime) > 1*time.Minute {
+			v.blockedCount = 0
+		}
+	}
+
+	return allowed
 }
 
 func startRateLimitCleaner() {
@@ -625,6 +697,10 @@ func startRateLimitCleaner() {
 			time.Sleep(5 * time.Minute)
 			visitorsMu.Lock()
 			for ip, v := range visitors {
+				// Keep punished visitors until penalty expires
+				if !v.punishedUntil.IsZero() && time.Now().Before(v.punishedUntil) {
+					continue
+				}
 				if time.Since(v.lastSeen) > 15*time.Minute {
 					delete(visitors, ip)
 				}