Add an option to overwrite which IP the rate-limiter chooses

soedinglab · Jun 30, 2023 · e812700 · e812700
1 parent 9f03a9b
commit e812700
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/backend/config.go b/backend/config.go
@@ -188,10 +188,11 @@ type ConfigAuth struct {
 }
 
 type ConfigRateLimit struct {
-	Rate   float64 `json:"rate" validate:"required"`
-	Burst  int     `json:"burst" validate:"required"`
-	TTL    int     `json:"ttl" validate:"required"`
-	Reason string  `json:"reason"`
+	Rate           float64 `json:"rate" validate:"required"`
+	Burst          int     `json:"burst" validate:"required"`
+	TTL            int     `json:"ttl" validate:"required"`
+	IpLookupHeader string  `json:"ipheader"`
+	Reason         string  `json:"reason"`
 }
 
 type ConfigWorker struct {

diff --git a/backend/server.go b/backend/server.go
@@ -461,6 +461,9 @@ func server(jobsystem JobSystem, config ConfigRoot) {
 			SetBurst(config.Server.RateLimit.Burst).
 			SetMessageContentType("application/json; charset=utf-8").
 			SetMessage(string(b))
+		if config.Server.RateLimit.IpLookupHeader == "" {
+			lmt.SetIPLookups([]string{config.Server.RateLimit.IpLookupHeader})
+		}
 		if config.App == AppMMseqs2 || config.App == AppFoldSeek {
 			r.Handle("/ticket", tollbooth.LimitFuncHandler(lmt, ticketHandlerFunc)).Methods("POST")
 		}