diff --git a/internal/app/server.go b/internal/app/server.go index fbb7a31..7149f45 100644 --- a/internal/app/server.go +++ b/internal/app/server.go @@ -78,26 +78,30 @@ func RunServer(cfg *config.Config) error { } router := gin.Default() - // Set up rate limiting (1 request per second per IP) - rateLimitStore := ratelimit.InMemoryStore(&ratelimit.InMemoryOptions{ - Rate: 4 * time.Second, - Limit: 2, - }) - rateLimiter := ratelimit.RateLimiter(rateLimitStore, &ratelimit.Options{ - ErrorHandler: func(c *gin.Context, info ratelimit.Info) { - c.JSON(429, gin.H{ - "error": "rate_limit_exceeded", - "message": "Too many requests. Try again in " + time.Until(info.ResetTime).String(), - }) - }, - KeyFunc: func(c *gin.Context) string { - return c.ClientIP() - }, - }) - - // Register API routes with rate limiting apiGroup := router.Group("/api") - apiGroup.Use(rateLimiter) + + if cfg.RateLimit > 0 { + // Set up rate limiting (2x to handle burst) + rateLimitStore := ratelimit.InMemoryStore(&ratelimit.InMemoryOptions{ + Rate: 2 * time.Second, + Limit: 2 * cfg.RateLimit, + }) + rateLimiter := ratelimit.RateLimiter(rateLimitStore, &ratelimit.Options{ + ErrorHandler: func(c *gin.Context, info ratelimit.Info) { + c.JSON(429, gin.H{ + "error": "rate_limit_exceeded", + "message": "Too many requests. Try again in " + time.Until(info.ResetTime).String(), + }) + }, + KeyFunc: func(c *gin.Context) string { + return c.ClientIP() + }, + }) + + apiGroup.Use(rateLimiter) + } + + // Register API routes api.RegisterHandlers(apiGroup, handler) web.DeclareRoutes(cfg, router) diff --git a/internal/config/cli.go b/internal/config/cli.go index d19b90b..17f0ff6 100644 --- a/internal/config/cli.go +++ b/internal/config/cli.go @@ -38,6 +38,7 @@ func declareFlags(o *Config) { flag.DurationVar(&o.Analysis.HTTPTimeout, "http-timeout", o.Analysis.HTTPTimeout, "Timeout when performing HTTP query") flag.Var(&StringArray{&o.Analysis.RBLs}, "rbl", "Append a RBL (use this option multiple time to append multiple RBLs)") flag.DurationVar(&o.ReportRetention, "report-retention", o.ReportRetention, "How long to keep reports (e.g., 720h, 30d). 0 = keep forever") + flag.UintVar(&o.RateLimit, "rate-limit", o.RateLimit, "API rate limit (requests per second per IP)") flag.Var(&URL{&o.SurveyURL}, "survey-url", "URL for user feedback survey") // Others flags are declared in some other files likes sources, storages, ... when they need specials configurations diff --git a/internal/config/config.go b/internal/config/config.go index 668573a..6f35ec7 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -42,6 +42,7 @@ type Config struct { Email EmailConfig Analysis AnalysisConfig ReportRetention time.Duration // How long to keep reports. 0 = keep forever + RateLimit uint // API rate limit (requests per second per IP) SurveyURL url.URL // URL for user feedback survey } @@ -71,6 +72,7 @@ func DefaultConfig() *Config { DevProxy: "", Bind: ":8080", ReportRetention: 0, // Keep reports forever by default + RateLimit: 1, // is in fact 2 requests per 2 seconds per IP (default) Database: DatabaseConfig{ Type: "sqlite", DSN: "happydeliver.db",