Rate limiting

Author: mahesh-hegde

Dockerfile

@@ -3,7 +3,7 @@
 ARG distroless_tag=nonroot
 
 FROM golang:1.25.4-trixie AS builder
-ARG build_tags="json1,fts5"
+ARG build_tags="json1,fts5,native_sqlite"
 RUN apt-get update && apt-get install -y build-essential && apt-get clean
 WORKDIR /app
 COPY go.mod go.sum ./
@@ -27,6 +27,6 @@ COPY --chown=nonroot:nonroot --from=builder /app/bin/dhee /app/bin/dhee
 COPY --chown=nonroot:nonroot --from=builder /app/data/dhee.db /app/data/dhee.db
 COPY --chown=nonroot:nonroot ./data/config.json /app/data/config.json
 
-CMD ["./bin/dhee", "server", "--data-dir", "./data", "--store", "sqlite", "--address", "0.0.0.0", "--cert-dir", "/app/certs", "--acme", "--rate-limit", "8"]
+CMD ["./bin/dhee", "server", "--data-dir", "./data", "--store", "sqlite", "--address", "0.0.0.0", "--cert-dir", "/app/certs", "--acme", "--rate-limit", "8", "--global-rate-limit", "64"]
 
 EXPOSE 8080

app/config/config.go

@@ -57,6 +57,7 @@ type ServerRuntimeConfig struct {
 	CertDir            string
 	AcmeEnabled        bool
 	RateLimit          int // 0 for inifinite
+	GlobalRateLimit    int // 0 for infinite
 	BehindLoadBalancer bool
 	GzipLevel          int // 0 to disable
 }

app/server/controllers.go

@@ -19,20 +19,43 @@ const MAX_CONCURRENT_REGEX_SEARCHES = 20
 
 // DheeController handles all HTTP requests.
 type DheeController struct {
-	ds           *dictionary.DictionaryService
-	es           *excerpts.ExcerptService
-	conf         *config.DheeConfig
-	regexLimiter chan struct{}
+	ds            *dictionary.DictionaryService
+	es            *excerpts.ExcerptService
+	conf          *config.DheeConfig
+	sconf         *config.ServerRuntimeConfig
+	regexLimiter  chan struct{}
+	globalLimiter chan struct{}
 }
 
 // NewDheeController creates a new controller instance and initializes the regex limiter.
-func NewDheeController(dictStore dictionary.DictStore, excerptStore excerpts.ExcerptStore, conf *config.DheeConfig, transliterator *transliteration.Transliterator) *DheeController {
-	return &DheeController{
+func NewDheeController(dictStore dictionary.DictStore, excerptStore excerpts.ExcerptStore, conf *config.DheeConfig, sconf *config.ServerRuntimeConfig, transliterator *transliteration.Transliterator) *DheeController {
+	controller := &DheeController{
 		ds:           dictionary.NewDictionaryService(dictStore, conf, transliterator),
 		es:           excerpts.NewExcerptService(dictStore, excerptStore, conf, transliterator),
 		conf:         conf,
+		sconf:        sconf,
 		regexLimiter: make(chan struct{}, MAX_CONCURRENT_REGEX_SEARCHES), // limit to 20 concurrent regex searches
 	}
+	if sconf.GlobalRateLimit > 0 {
+		controller.globalLimiter = make(chan struct{}, sconf.GlobalRateLimit)
+	}
+	for i := 0; i < sconf.GlobalRateLimit; i++ {
+		controller.globalLimiter <- struct{}{}
+	}
+	return controller
+}
+
+func (c *DheeController) GlobalRateLimitMiddleware(next echo.HandlerFunc) echo.HandlerFunc {
+	return func(ctx echo.Context) error {
+		// Acquire a token from the global limiter, but if context is done, return 429
+		select {
+		case <-c.globalLimiter:
+			defer func() { c.globalLimiter <- struct{}{} }()
+			return next(ctx)
+		case <-ctx.Request().Context().Done():
+			return echo.NewHTTPError(http.StatusTooManyRequests, "Global concurrent request limit reached, please try later")
+		}
+	}
 }
 
 func (c *DheeController) GetHome(ctx echo.Context) error {

app/server/echo_server.go

@@ -100,6 +100,10 @@ func StartServer(controller *DheeController, dheeConf *config.DheeConfig, server
 		e.Use(middleware.RateLimiterWithConfig(config))
 	}
 
+	if serverConf.GlobalRateLimit > 0 {
+		e.Use(controller.GlobalRateLimitMiddleware)
+	}
+
 	if serverConf.GzipLevel != 0 {
 		e.Use(middleware.GzipWithConfig(middleware.GzipConfig{Level: serverConf.GzipLevel, MinLength: 512}))
 	}

cmd/dhee/main.go

@@ -116,7 +116,7 @@ func runServer() {
 	var address, dataDir, store string
 	var port int
 	var cpuProfile, memProfile string
-	var serverConfig config.ServerRuntimeConfig
+	var serverConf config.ServerRuntimeConfig
 
 	jsonHandler := slog.NewJSONHandler(os.Stdout, nil)
 	slog.SetDefault(slog.New(jsonHandler))
@@ -127,13 +127,14 @@ func runServer() {
 	flags.StringVar(&cpuProfile, "cpu-profile", "", "write cpu profile to file")
 	flags.StringVar(&memProfile, "mem-profile", "", "write memory profile to file")
 
-	flags.StringVarP(&serverConfig.Addr, "address", "a", "localhost", "Server address to bind")
-	flags.IntVarP(&serverConfig.Port, "port", "p", 8080, "Server port to bind")
-	flags.StringVar(&serverConfig.CertDir, "cert-dir", "", "directory to read/write TLS certs for ACME")
-	flags.BoolVar(&serverConfig.AcmeEnabled, "acme", false, "use ACME to renew TLS certificates")
-	flags.BoolVar(&serverConfig.BehindLoadBalancer, "behind-load-balancer", false, "Certain behaviors when behind a load balancer (e.g., trusting X-Forwarded-For header)")
-	flags.IntVar(&serverConfig.GzipLevel, "gzip-level", 1, "Gzip compression level (1-9), or 0 to disable gzip")
-	flags.IntVar(&serverConfig.RateLimit, "rate-limit", 0, "Number of requests per second for rate limiting")
+	flags.StringVarP(&serverConf.Addr, "address", "a", "localhost", "Server address to bind")
+	flags.IntVarP(&serverConf.Port, "port", "p", 8080, "Server port to bind")
+	flags.StringVar(&serverConf.CertDir, "cert-dir", "", "directory to read/write TLS certs for ACME")
+	flags.BoolVar(&serverConf.AcmeEnabled, "acme", false, "use ACME to renew TLS certificates")
+	flags.BoolVar(&serverConf.BehindLoadBalancer, "behind-load-balancer", false, "Certain behaviors when behind a load balancer (e.g., trusting X-Forwarded-For header)")
+	flags.IntVar(&serverConf.GzipLevel, "gzip-level", 1, "Gzip compression level (1-9), or 0 to disable gzip")
+	flags.IntVar(&serverConf.RateLimit, "rate-limit", 0, "Number of requests per second for rate limiting")
+	flags.IntVar(&serverConf.GlobalRateLimit, "global-rate-limit", 0, "Global request rate limit per second")
 
 	flags.Parse(os.Args[2:])
 
@@ -213,8 +214,8 @@ func runServer() {
 		os.Exit(1)
 	}
 
-	controller := server.NewDheeController(dictStore, excerptStore, conf, transliterator)
-	server.StartServer(controller, conf, serverConfig)
+	controller := server.NewDheeController(dictStore, excerptStore, conf, &serverConf, transliterator)
+	server.StartServer(controller, conf, serverConf)
 }
 
 func runIndex() {