6d16828b05
- Implements token bucket rate limiting with goroutine-safe RateLimiter - Adds support for IP-based, user-based, and combined rate limiting strategies - Creates HTTP middleware to apply rate limiting to API endpoints - Provides configurable rate, burst, and TTL parameters via config - Introduces RateLimitingConfig to app configuration - Includes comprehensive test coverage for all features - Integrates with existing HTTP handlers and channel managers
60 lines
1.7 KiB
Go
60 lines
1.7 KiB
Go
package ratelimit
|
|
|
|
import (
|
|
"net/http"
|
|
|
|
"github.com/sipeed/picoclaw/pkg/config"
|
|
)
|
|
|
|
// HTTPServer integrates rate limiting with an HTTP server
|
|
type HTTPServer struct {
|
|
config *config.Config
|
|
server *http.Server
|
|
limiter *RateLimiter
|
|
}
|
|
|
|
// NewHTTPServer creates a new HTTP server with integrated rate limiting
|
|
func NewHTTPServer(cfg *config.Config) *HTTPServer {
|
|
limiter := CreateDefaultRateLimiterFromConfig(cfg)
|
|
|
|
return &HTTPServer{
|
|
config: cfg,
|
|
limiter: limiter,
|
|
}
|
|
}
|
|
|
|
// ApplyRateLimit wraps an HTTP handler with rate limiting based on configuration
|
|
func (hs *HTTPServer) ApplyRateLimit(handler http.Handler) http.Handler {
|
|
if !hs.config.Tools.RateLimiting.Enabled {
|
|
return handler
|
|
}
|
|
|
|
middleware := NewRateLimitMiddleware(hs.limiter)
|
|
|
|
strategy := hs.config.Tools.RateLimiting.Strategy
|
|
if strategy == "" {
|
|
strategy = "ip" // default to IP-based limiting
|
|
}
|
|
|
|
return middleware.CreateMiddleware(strategy)(handler)
|
|
}
|
|
|
|
// ApplyRateLimitFunc wraps an HTTP handler function with rate limiting
|
|
func (hs *HTTPServer) ApplyRateLimitFunc(handlerFunc http.HandlerFunc) http.Handler {
|
|
return hs.ApplyRateLimit(handlerFunc)
|
|
}
|
|
|
|
// WrapRoute provides an easy way to wrap individual routes with rate limiting
|
|
func (hs *HTTPServer) WrapRoute(pattern string, handler http.Handler) (string, http.Handler) {
|
|
return pattern, hs.ApplyRateLimit(handler)
|
|
}
|
|
|
|
// WrapRouteFunc provides an easy way to wrap individual routes with rate limiting
|
|
func (hs *HTTPServer) WrapRouteFunc(pattern string, handlerFunc http.HandlerFunc) (string, http.Handler) {
|
|
return pattern, hs.ApplyRateLimit(handlerFunc)
|
|
}
|
|
|
|
// GetLimiter returns the rate limiter instance
|
|
func (hs *HTTPServer) GetLimiter() *RateLimiter {
|
|
return hs.limiter
|
|
}
|