From e8127001e14a9995bbd105120ba2ddf5d0b7f29d Mon Sep 17 00:00:00 2001
From: Milot Mirdita <milot@mirdita.de>
Date: Fri, 30 Jun 2023 13:46:04 +0900
Subject: [PATCH] Add an option to overwrite which IP the rate-limiter chooses

---
 backend/config.go | 9 +++++----
 backend/server.go | 3 +++
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/backend/config.go b/backend/config.go
index b4847e8..7dda8ae 100644
--- 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
index d501e7e..45ab71a 100644
--- 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")
 		}