Day 93 — Per-IP Rate Limiting with CIDR Awareness

Key Concepts

• X-Forwarded-For trust: Only trust XFF[len-1] (last hop, set by your trusted L4 proxy). Never trust the full chain from untrusted clients.
• IPv6 /48 grouping: IPv6 bots spread across a /48. Group by first 48 bits.
• Counter expiry: Un-evicted stale counters grow memory unbounded. Background goroutine cleans up.

Exercise — CIDR-Aware Rate Limiter

package ddos

import (
    "fmt"
    "net"
    "net/http"
    "strings"
    "sync"
    "time"
)

type counterEntry struct {
    bucket  *TokenBucket
    lastSee time.Time
}

type CIDRRateLimiter struct {
    mu      sync.RWMutex
    ipLimit  int     // per /32 or /128
    subLimit int     // per /24 or /48
    rate     float64 // tokens/second
    entries  map[string]*counterEntry
    done     chan struct{}
}

func NewCIDRRateLimiter(ipLimit, subLimit int, rate float64) *CIDRRateLimiter {
    rl := &CIDRRateLimiter{
        ipLimit:  ipLimit,
        subLimit: subLimit,
        rate:     rate,
        entries:  make(map[string]*counterEntry),
        done:     make(chan struct{}),
    }
    go rl.cleanup()
    return rl
}

// realIP extracts the real client IP, trusting only one XFF hop
func realIP(r *http.Request) string {
    xff := r.Header.Get("X-Forwarded-For")
    if xff == "" {
        host, _, _ := net.SplitHostPort(r.RemoteAddr)
        return host
    }
    // Trust only the rightmost entry added by our L4 proxy
    parts := strings.Split(xff, ",")
    ip := strings.TrimSpace(parts[len(parts)-1])
    if net.ParseIP(ip) == nil {
        host, _, _ := net.SplitHostPort(r.RemoteAddr)
        return host
    }
    return ip
}

// subnetKey returns the /24 (IPv4) or /48 (IPv6) key
func subnetKey(ipStr string) string {
    ip := net.ParseIP(ipStr)
    if ip == nil {
        return ipStr
    }
    if ip4 := ip.To4(); ip4 != nil {
        // IPv4: mask to /24
        return fmt.Sprintf("%d.%d.%d.0/24", ip4[0], ip4[1], ip4[2])
    }
    // IPv6: mask to /48 (first 6 bytes)
    return fmt.Sprintf("%02x%02x:%02x%02x:%02x%02x::/48",
        ip[0], ip[1], ip[2], ip[3], ip[4], ip[5])
}

func (rl *CIDRRateLimiter) getBucket(key string, limit int) *TokenBucket {
    rl.mu.Lock()
    defer rl.mu.Unlock()
    e, ok := rl.entries[key]
    if !ok {
        e = &counterEntry{bucket: NewTokenBucket(float64(limit), rl.rate)}
        rl.entries[key] = e
    }
    e.lastSee = time.Now()
    return e.bucket
}

func (rl *CIDRRateLimiter) Allow(r *http.Request) bool {
    ip := realIP(r)
    subnet := subnetKey(ip)

    // Both per-IP and per-subnet must allow
    if !rl.getBucket(ip, rl.ipLimit).Allow() {
        return false
    }
    if !rl.getBucket(subnet, rl.subLimit).Allow() {
        return false
    }
    return true
}

func (rl *CIDRRateLimiter) cleanup() {
    ticker := time.NewTicker(5 * time.Minute)
    defer ticker.Stop()
    for {
        select {
        case <-rl.done:
            return
        case <-ticker.C:
            cutoff := time.Now().Add(-10 * time.Minute)
            rl.mu.Lock()
            for k, e := range rl.entries {
                if e.lastSee.Before(cutoff) {
                    delete(rl.entries, k)
                }
            }
            rl.mu.Unlock()
        }
    }
}

func (rl *CIDRRateLimiter) Stop() { close(rl.done) }

func (rl *CIDRRateLimiter) Middleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        if !rl.Allow(r) {
            w.Header().Set("Retry-After", "1")
            http.Error(w, `{"error":"rate_limit_exceeded"}`, http.StatusTooManyRequests)
            return
        }
        next.ServeHTTP(w, r)
    })
}