Integration Guide: Adding Rate Limiting to Your Go Web Application

go get github.com/Jouini-Mohamed-Chaker/go-ratelimit

package main

import (
    "fmt"
    "log"
    "net/http"
    "time"

    "github.com/Jouini-Mohamed-Chaker/go-ratelimit"
)

func main() {
    // Create rate limiter: 1000 requests per hour with 50 burst capacity
    limiter := ratelimit.NewWithConfig(ratelimit.Config{
        Limit:     1000,
        Window:    time.Hour,
        BurstSize: 50,
        SkipPaths: []string{"/health", "/metrics"}, // Health checks bypass rate limiting
    })

    // Your existing handlers
    http.HandleFunc("/api/users", handleUsers)
    http.HandleFunc("/api/posts", handlePosts)
    http.HandleFunc("/health", handleHealth)

    // Apply rate limiting to all routes
    server := &http.Server{
        Addr:    ":8080",
        Handler: limiter.Wrap(http.DefaultServeMux),
    }

    log.Println("Server starting on :8080 with rate limiting (1000 req/hour)")
    log.Fatal(server.ListenAndServe())
}

func handleUsers(w http.ResponseWriter, r *http.Request) {
    w.Header().Set("Content-Type", "application/json")
    fmt.Fprintf(w, `{"users": ["alice", "bob", "charlie"]}`)
}

func handlePosts(w http.ResponseWriter, r *http.Request) {
    w.Header().Set("Content-Type", "application/json")
    fmt.Fprintf(w, `{"posts": [{"id": 1, "title": "Hello World"}]}`)
}

func handleHealth(w http.ResponseWriter, r *http.Request) {
    w.WriteHeader(http.StatusOK)
    w.Write([]byte("OK"))
}

package main

import (
    "encoding/json"
    "net/http"
    "time"

    "github.com/gorilla/mux"
    "github.com/Jouini-Mohamed-Chaker/go-ratelimit"
)

func main() {
    // Create rate limiter with different limits for different endpoints
    apiLimiter := ratelimit.NewWithConfig(ratelimit.Config{
        Limit:      100,              // 100 requests per minute for API
        Window:     time.Minute,
        BurstSize:  10,
        TrustedIPs: []string{"127.0.0.1"}, // Local development
    })

    authLimiter := ratelimit.NewWithConfig(ratelimit.Config{
        Limit:  5,               // Strict limit for auth endpoints
        Window: time.Minute,
        BurstSize: 2,
    })

    r := mux.NewRouter()

    // API routes with standard rate limiting
    apiRouter := r.PathPrefix("/api").Subrouter()
    apiRouter.Use(func(next http.Handler) http.Handler {
        return apiLimiter.Wrap(next)
    })
    apiRouter.HandleFunc("/users", getUsers).Methods("GET")
    apiRouter.HandleFunc("/posts", getPosts).Methods("GET")

    // Auth routes with stricter rate limiting
    authRouter := r.PathPrefix("/auth").Subrouter()
    authRouter.Use(func(next http.Handler) http.Handler {
        return authLimiter.Wrap(next)
    })
    authRouter.HandleFunc("/login", handleLogin).Methods("POST")
    authRouter.HandleFunc("/register", handleRegister).Methods("POST")

    // Public routes without rate limiting
    r.HandleFunc("/health", handleHealth).Methods("GET")
    r.HandleFunc("/", handleHome).Methods("GET")

    http.ListenAndServe(":8080", r)
}

func getUsers(w http.ResponseWriter, r *http.Request) {
    users := []map[string]interface{}{
        {"id": 1, "name": "Alice"},
        {"id": 2, "name": "Bob"},
    }
    w.Header().Set("Content-Type", "application/json")
    json.NewEncoder(w).Encode(map[string]interface{}{"users": users})
}

func getPosts(w http.ResponseWriter, r *http.Request) {
    posts := []map[string]interface{}{
        {"id": 1, "title": "First Post", "author": "Alice"},
        {"id": 2, "title": "Second Post", "author": "Bob"},
    }
    w.Header().Set("Content-Type", "application/json")
    json.NewEncoder(w).Encode(map[string]interface{}{"posts": posts})
}

func handleLogin(w http.ResponseWriter, r *http.Request) {
    // Simulate login logic
    w.Header().Set("Content-Type", "application/json")
    json.NewEncoder(w).Encode(map[string]string{
        "message": "Login successful",
        "token":   "fake-jwt-token",
    })
}

func handleRegister(w http.ResponseWriter, r *http.Request) {
    // Simulate registration logic
    w.Header().Set("Content-Type", "application/json")
    json.NewEncoder(w).Encode(map[string]string{
        "message": "Registration successful",
    })
}

func handleHealth(w http.ResponseWriter, r *http.Request) {
    w.WriteHeader(http.StatusOK)
    w.Write([]byte("OK"))
}

func handleHome(w http.ResponseWriter, r *http.Request) {
    w.Header().Set("Content-Type", "text/html")
    w.Write([]byte(`<h1>Welcome to My API</h1><p>Rate limiting is active!</p>`))
}

package main

import (
    "net/http"
    "time"

    "github.com/gin-gonic/gin"
    "github.com/Jouini-Mohamed-Chaker/go-ratelimit"
)

// Gin middleware adapter
func RateLimitMiddleware(limiter *ratelimit.Limiter) gin.HandlerFunc {
    return func(c *gin.Context) {
        // Create a wrapper to make Gin's context compatible
        wrapper := &ginResponseWriter{ResponseWriter: c.Writer, context: c}
        
        if limiter.Wrap(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
            c.Next()
        })).ServeHTTP(wrapper, c.Request); wrapper.written {
            c.Abort()
            return
        }
    }
}

type ginResponseWriter struct {
    gin.ResponseWriter
    context *gin.Context
    written bool
}

func (w *ginResponseWriter) Write(data []byte) (int, error) {
    w.written = true
    return w.ResponseWriter.Write(data)
}

func (w *ginResponseWriter) WriteHeader(statusCode int) {
    if statusCode != http.StatusOK {
        w.written = true
    }
    w.ResponseWriter.WriteHeader(statusCode)
}

func main() {
    r := gin.Default()

    // Create rate limiters for different route groups
    apiLimiter := ratelimit.NewWithConfig(ratelimit.Config{
        Limit:     200,
        Window:    time.Minute,
        BurstSize: 20,
    })

    authLimiter := ratelimit.NewWithConfig(ratelimit.Config{
        Limit:  10,
        Window: time.Minute,
        BurstSize: 3,
    })

    // API routes with rate limiting
    api := r.Group("/api")
    api.Use(RateLimitMiddleware(apiLimiter))
    {
        api.GET("/users", getUsers)
        api.GET("/posts", getPosts)
        api.POST("/posts", createPost)
    }

    // Auth routes with stricter limiting
    auth := r.Group("/auth")
    auth.Use(RateLimitMiddleware(authLimiter))
    {
        auth.POST("/login", login)
        auth.POST("/register", register)
        auth.POST("/forgot-password", forgotPassword)
    }

    // Public routes (no rate limiting)
    r.GET("/health", func(c *gin.Context) {
        c.JSON(200, gin.H{"status": "ok"})
    })

    r.GET("/", func(c *gin.Context) {
        c.HTML(200, "index.html", gin.H{
            "title": "My API with Rate Limiting",
        })
    })

    r.Run(":8080")
}

func getUsers(c *gin.Context) {
    c.JSON(200, gin.H{
        "users": []gin.H{
            {"id": 1, "name": "Alice"},
            {"id": 2, "name": "Bob"},
        },
    })
}

func getPosts(c *gin.Context) {
    c.JSON(200, gin.H{
        "posts": []gin.H{
            {"id": 1, "title": "First Post"},
            {"id": 2, "title": "Second Post"},
        },
    })
}

func createPost(c *gin.Context) {
    c.JSON(201, gin.H{
        "message": "Post created",
        "id":      123,
    })
}

func login(c *gin.Context) {
    c.JSON(200, gin.H{
        "message": "Login successful",
        "token":   "fake-jwt-token",
    })
}

func register(c *gin.Context) {
    c.JSON(201, gin.H{
        "message": "Registration successful",
    })
}

func forgotPassword(c *gin.Context) {
    c.JSON(200, gin.H{
        "message": "Password reset email sent",
    })
}

package main

import (
    "net/http"
    "time"

    "github.com/labstack/echo/v4"
    "github.com/labstack/echo/v4/middleware"
    "github.com/Jouini-Mohamed-Chaker/go-ratelimit"
)

// Echo middleware adapter
func RateLimitMiddleware(limiter *ratelimit.Limiter) echo.MiddlewareFunc {
    return func(next echo.HandlerFunc) echo.HandlerFunc {
        return func(c echo.Context) error {
            req := c.Request()
            res := c.Response().Writer

            // Check if request should be rate limited
            handler := limiter.Wrap(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                // Continue to next handler
                if err := next(c); err != nil {
                    c.Error(err)
                }
            }))

            handler.ServeHTTP(res, req)
            return nil
        }
    }
}

func main() {
    e := echo.New()

    // Middleware
    e.Use(middleware.Logger())
    e.Use(middleware.Recover())

    // Rate limiters
    apiLimiter := ratelimit.NewWithConfig(ratelimit.Config{
        Limit:     500,
        Window:    time.Hour,
        BurstSize: 50,
        SkipPaths: []string{"/health", "/metrics"},
    })

    strictLimiter := ratelimit.NewWithConfig(ratelimit.Config{
        Limit:  15,
        Window: time.Minute,
        BurstSize: 5,
    })

    // API group with rate limiting
    api := e.Group("/api", RateLimitMiddleware(apiLimiter))
    api.GET("/users", getUsers)
    api.GET("/posts", getPosts)
    api.POST("/posts", createPost)

    // Auth group with strict rate limiting
    auth := e.Group("/auth", RateLimitMiddleware(strictLimiter))
    auth.POST("/login", login)
    auth.POST("/register", register)

    // Public routes
    e.GET("/health", health)
    e.GET("/", home)

    e.Logger.Fatal(e.Start(":8080"))
}

func getUsers(c echo.Context) error {
    return c.JSON(http.StatusOK, map[string]interface{}{
        "users": []map[string]interface{}{
            {"id": 1, "name": "Alice"},
            {"id": 2, "name": "Bob"},
        },
    })
}

func getPosts(c echo.Context) error {
    return c.JSON(http.StatusOK, map[string]interface{}{
        "posts": []map[string]interface{}{
            {"id": 1, "title": "Hello World"},
        },
    })
}

func createPost(c echo.Context) error {
    return c.JSON(http.StatusCreated, map[string]string{
        "message": "Post created successfully",
    })
}

func login(c echo.Context) error {
    return c.JSON(http.StatusOK, map[string]string{
        "message": "Login successful",
        "token":   "fake-jwt-token",
    })
}

func register(c echo.Context) error {
    return c.JSON(http.StatusCreated, map[string]string{
        "message": "Registration successful",
    })
}

func health(c echo.Context) error {
    return c.String(http.StatusOK, "OK")
}

func home(c echo.Context) error {
    return c.HTML(http.StatusOK, "<h1>Welcome to My API</h1><p>Rate limiting active!</p>")
}

package main

import (
    "log"
    "net/http"
    "os"
    "strconv"
    "strings"
    "time"

    "github.com/Jouini-Mohamed-Chaker/go-ratelimit"
)

type AppConfig struct {
    Port            string
    RateLimit       int
    RateWindow      time.Duration
    BurstSize       int
    TrustedIPs      []string
    SkipPaths       []string
    EnableRateLimit bool
}

func loadConfig() *AppConfig {
    config := &AppConfig{
        Port:            getEnv("PORT", "8080"),
        RateLimit:       getEnvInt("RATE_LIMIT", 1000),
        RateWindow:      getEnvDuration("RATE_WINDOW", "1h"),
        BurstSize:       getEnvInt("BURST_SIZE", 100),
        TrustedIPs:      getEnvSlice("TRUSTED_IPS", []string{"127.0.0.1"}),
        SkipPaths:       getEnvSlice("SKIP_PATHS", []string{"/health", "/metrics"}),
        EnableRateLimit: getEnvBool("ENABLE_RATE_LIMIT", true),
    }
    return config
}

func getEnv(key, defaultValue string) string {
    if value := os.Getenv(key); value != "" {
        return value
    }
    return defaultValue
}

func getEnvInt(key string, defaultValue int) int {
    if value := os.Getenv(key); value != "" {
        if intValue, err := strconv.Atoi(value); err == nil {
            return intValue
        }
    }
    return defaultValue
}

func getEnvDuration(key string, defaultValue string) time.Duration {
    value := getEnv(key, defaultValue)
    if duration, err := time.ParseDuration(value); err == nil {
        return duration
    }
    // Fallback parsing
    if duration, err := time.ParseDuration(defaultValue); err == nil {
        return duration
    }
    return time.Hour
}

func getEnvBool(key string, defaultValue bool) bool {
    if value := os.Getenv(key); value != "" {
        if boolValue, err := strconv.ParseBool(value); err == nil {
            return boolValue
        }
    }
    return defaultValue
}

func getEnvSlice(key string, defaultValue []string) []string {
    if value := os.Getenv(key); value != "" {
        return strings.Split(value, ",")
    }
    return defaultValue
}

func main() {
    config := loadConfig()

    var handler http.Handler = http.DefaultServeMux

    // Setup your routes
    http.HandleFunc("/api/users", handleAPI)
    http.HandleFunc("/api/posts", handleAPI)
    http.HandleFunc("/health", handleHealth)

    // Apply rate limiting if enabled
    if config.EnableRateLimit {
        limiter := ratelimit.NewWithConfig(ratelimit.Config{
            Limit:      config.RateLimit,
            Window:     config.RateWindow,
            BurstSize:  config.BurstSize,
            SkipPaths:  config.SkipPaths,
            TrustedIPs: config.TrustedIPs,
        })

        // Custom rate limit handler for better API responses
        limiter.OnLimitReached = func(w http.ResponseWriter, r *http.Request) {
            w.Header().Set("Content-Type", "application/json")
            w.Header().Set("Retry-After", "60")
            w.WriteHeader(http.StatusTooManyRequests)
            w.Write([]byte(`{
                "error": "rate_limit_exceeded",
                "message": "Too many requests. Please slow down.",
                "retry_after": 60
            }`))
        }

        handler = limiter.Wrap(handler)
        log.Printf("Rate limiting enabled: %d requests per %v (burst: %d)", 
            config.RateLimit, config.RateWindow, config.BurstSize)
    } else {
        log.Println("Rate limiting disabled")
    }

    log.Printf("Server starting on port %s", config.Port)
    log.Fatal(http.ListenAndServe(":"+config.Port, handler))
}

func handleAPI(w http.ResponseWriter, r *http.Request) {
    w.Header().Set("Content-Type", "application/json")
    w.Write([]byte(`{"message": "API response", "path": "` + r.URL.Path + `"}`))
}

func handleHealth(w http.ResponseWriter, r *http.Request) {
    w.WriteHeader(http.StatusOK)
    w.Write([]byte("OK"))
}

PORT=8080
RATE_LIMIT=2000
RATE_WINDOW=1h
BURST_SIZE=200
TRUSTED_IPS=127.0.0.1,10.0.0.0/8,192.168.0.0/16
SKIP_PATHS=/health,/metrics,/favicon.ico
ENABLE_RATE_LIMIT=true

FROM golang:1.21-alpine AS builder

WORKDIR /app
COPY go.mod go.sum ./
RUN go mod download

COPY . .
RUN CGO_ENABLED=0 GOOS=linux go build -o main .

FROM alpine:latest
RUN apk --no-cache add ca-certificates
WORKDIR /root/

COPY --from=builder /app/main .

EXPOSE 8080
CMD ["./main"]

version: '3.8'
services:
  api:
    build: .
    ports:
      - "8080:8080"
    env_file:
      - .env
    restart: unless-stopped
    healthcheck:
      test: ["CMD", "wget", "--quiet", "--tries=1", "--spider", "http://localhost:8080/health"]
      interval: 30s
      timeout: 10s
      retries: 3

package main

import (
    "crypto/sha256"
    "fmt"
    "net/http"
    "sync"
    "time"

    "github.com/Jouini-Mohamed-Chaker/go-ratelimit"
)

type UserLimiterManager struct {
    limiters map[string]*ratelimit.Limiter
    mu       sync.RWMutex
    config   ratelimit.Config
}

func NewUserLimiterManager(config ratelimit.Config) *UserLimiterManager {
    return &UserLimiterManager{
        limiters: make(map[string]*ratelimit.Limiter),
        config:   config,
    }
}

func (m *UserLimiterManager) GetLimiter(userID string) *ratelimit.Limiter {
    m.mu.RLock()
    limiter, exists := m.limiters[userID]
    m.mu.RUnlock()

    if exists {
        return limiter
    }

    // Create new limiter for user
    m.mu.Lock()
    defer m.mu.Unlock()

    // Double-check after acquiring write lock
    if limiter, exists := m.limiters[userID]; exists {
        return limiter
    }

    limiter = ratelimit.NewWithConfig(m.config)
    m.limiters[userID] = limiter
    return limiter
}

// Cleanup old limiters periodically (optional)
func (m *UserLimiterManager) Cleanup() {
    m.mu.Lock()
    defer m.mu.Unlock()

    for userID, limiter := range m.limiters {
        // Remove limiters that haven't been used recently
        // This is a simple implementation - you might want more sophisticated cleanup
        if limiter.RemainingTokens() == float64(m.config.Limit) {
            delete(m.limiters, userID)
        }
    }
}

func getUserIDFromRequest(r *http.Request) string {
    // Try to get user ID from various sources
    if userID := r.Header.Get("X-User-ID"); userID != "" {
        return userID
    }
    
    if userID := r.URL.Query().Get("user_id"); userID != "" {
        return userID
    }

    // Fall back to IP-based identification
    clientIP := getClientIP(r)
    hash := sha256.Sum256([]byte(clientIP))
    return fmt.Sprintf("ip-%x", hash[:8]) // Use first 8 bytes of hash
}

func getClientIP(r *http.Request) string {
    if ip := r.Header.Get("X-Forwarded-For"); ip != "" {
        return ip
    }
    if ip := r.Header.Get("X-Real-IP"); ip != "" {
        return ip
    }
    return r.RemoteAddr
}

func main() {
    // Create per-user rate limiter manager
    userManager := NewUserLimiterManager(ratelimit.Config{
        Limit:     100,           // 100 requests per user
        Window:    time.Hour,     // per hour
        BurstSize: 10,
    })

    // Global rate limiter for overall protection
    globalLimiter := ratelimit.NewWithConfig(ratelimit.Config{
        Limit:     10000,         // 10k total requests
        Window:    time.Hour,     // per hour
        BurstSize: 1000,
        SkipPaths: []string{"/health"},
    })

    // Start cleanup goroutine
    go func() {
        ticker := time.NewTicker(time.Hour)
        defer ticker.Stop()
        for range ticker.C {
            userManager.Cleanup()
        }
    }()

    handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        // First check global rate limit
        if !globalLimiter.Allow() {
            http.Error(w, "Global rate limit exceeded", http.StatusTooManyRequests)
            return
        }

        // Then check per-user rate limit
        userID := getUserIDFromRequest(r)
        userLimiter := userManager.GetLimiter(userID)
        
        if !userLimiter.Allow() {
            w.Header().Set("X-User-ID", userID)
            http.Error(w, "User rate limit exceeded", http.StatusTooManyRequests)
            return
        }

        // Add headers for debugging
        w.Header().Set("X-User-ID", userID)
        w.Header().Set("X-Remaining-Global", fmt.Sprintf("%.0f", globalLimiter.RemainingTokens()))
        w.Header().Set("X-Remaining-User", fmt.Sprintf("%.0f", userLimiter.RemainingTokens()))

        // Your API logic here
        w.Header().Set("Content-Type", "application/json")
        fmt.Fprintf(w, `{"message": "Success", "user_id": "%s"}`, userID)
    })

    http.Handle("/api/", handler)
    http.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
        w.Write([]byte("OK"))
    })

    log.Println("Server starting with per-user rate limiting...")
    http.ListenAndServe(":8080", nil)
}

package main

import (
    "expvar"
    "fmt"
    "net/http"
    "time"

    "github.com/Jouini-Mohamed-Chaker/go-ratelimit"
)

var (
    rateLimitStats = expvar.NewMap("ratelimit")
)

func setupMetrics(limiter *ratelimit.Limiter) {
    // Update metrics every 10 seconds
    go func() {
        ticker := time.NewTicker(10 * time.Second)
        defer ticker.Stop()

        for range ticker.C {
            total, allowed, blocked := limiter.GetStats()
            remaining := limiter.RemainingTokens()

            rateLimitStats.Set("total_requests", &expvar.Int{})
            rateLimitStats.Get("total_requests").(*expvar.Int).Set(total)

            rateLimitStats.Set("allowed_requests", &expvar.Int{})
            rateLimitStats.Get("allowed_requests").(*expvar.Int).Set(allowed)

            rateLimitStats.Set("blocked_requests", &expvar.Int{})
            rateLimitStats.Get("blocked_requests").(*expvar.Int).Set(blocked)

            rateLimitStats.Set("remaining_tokens", &expvar.Float{})
            rateLimitStats.Get("remaining_tokens").(*expvar.Float).Set(remaining)

            if total > 0 {
                blockRate := float64(blocked) / float64(total) * 100
                rateLimitStats.Set("block_rate_percent", &expvar.Float{})
                rateLimitStats.Get("block_rate_percent").(*expvar.Float).Set(blockRate)
            }
        }
    }()
}

func main() {
    limiter := ratelimit.NewWithConfig(ratelimit.Config{
        Limit:     100,
        Window:    time.Minute,
        BurstSize: 10,
        SkipPaths: []string{"/health", "/metrics", "/debug/vars"},
    })

    // Setup metrics collection
    setupMetrics(limiter)

    // Add custom callbacks for real-time monitoring
    limiter.OnLimitReached = func(w http.ResponseWriter, r *http.Request) {
        // Log rate limit events
        fmt.Printf("[RATE_LIMIT] %s %s from %s\n", 
            r.Method, r.URL.Path, r.RemoteAddr)
        
        w.Header().Set("Content-Type", "application/json")
        w.WriteHeader(http.StatusTooManyRequests)
        w.Write([]byte(`{
            "error": "rate_limit_exceeded",
            "message": "Please slow down your requests",
            "retry_after": 60
        }`))
    }

    limiter.OnAllow = func(r *http.Request, remaining float64) {
        // Log when running low on tokens
        if remaining < 10 {
            fmt.Printf("[RATE_LIMIT] Low tokens: %.0f remaining for %s\n", 
                remaining, r.RemoteAddr)
        }
    }

    // Your API routes
    http.HandleFunc("/api/data", func(w http.ResponseWriter, r *http.Request) {
        w.Header().Set("Content-Type", "application/json")
        fmt.Fprintf(w, `{"data": "Hello World", "timestamp": "%s"}`, 
            time.Now().Format(time.RFC3339))
    })

    // Metrics endpoint (expvar provides /debug/vars automatically)
    http.HandleFunc("/metrics", func(w http.ResponseWriter, r *http.Request) {
        total, allowed, blocked := limiter.GetStats()
        remaining := limiter.RemainingTokens()

        w.Header().Set("Content-Type", "application/json")
        fmt.Fprintf(w, `{
            "rate_limit": {
                "total_requests": %d,
                "allowed_requests": %d,
                "blocked_requests": %d,
                "remaining_tokens": %.2f,
                "block_rate_percent": %.2f
            }
        }`, total, allowed, blocked, remaining, 
        func() float64 {
            if total > 0 {
                return float64(blocked) / float64(total) * 100
            }
            return 0
        }())
    })

    http.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
        w.Write([]byte("OK"))
    })

    // Apply rate limiting
    server := &http.Server{
        Addr:    ":8080",
        Handler: limiter.Wrap(http.DefaultServeMux),
    }

    fmt.Println("Server starting on :8080")
    fmt.Println("Metrics available at /metrics and /debug/vars")
    server.ListenAndServe()
}