feat(server): Support supplying a list of headers when keying IPs for…

… rate limiting (#12199)

Signed-off-by: Yuan Tang <terrytangyuan@gmail.com>

.spelling

@@ -64,6 +64,7 @@ InsideBoard
 Invocators
 IAM-based
 Istio
+IPs
 Jemison
 JetBrains
 KNative
@@ -235,4 +236,4 @@ CronWorkflow
 CronWorkflows
 maxFailures
 maxSuccess
-gitops
+gitops

docs/environment-variables.md

@@ -122,6 +122,7 @@ data:
 | `DISABLE_VALUE_LIST_RETRIEVAL_KEY_PATTERN` | `string` | `""`    | Disable the retrieval of the list of label values for keys based on this regular expression.                            |
 | `FIRST_TIME_USER_MODAL`                    | `bool`   | `true`  | Show this modal.                                                                                                        |
 | `FEEDBACK_MODAL`                           | `bool`   | `true`  | Show this modal.                                                                                                        |
+| `IP_KEY_FUNC_HEADERS`                      | `string` | `""`    | List of comma separated request headers containing IPs to use for rate limiting. For example, "X-Forwarded-For,X-Real-IP". By default, uses the request's remote IP address.          |
 | `NEW_VERSION_MODAL`                        | `bool`   | `true`  | Show this modal.                                                                                                        |
 | `POD_NAMES`                                | `string` | `v2`    | Whether to have pod names contain the template name (v2) or be the node id (v1) - should be set the same for Controller |
 | `GRPC_MESSAGE_SIZE`                        | `string` | `104857600` | Use different GRPC Max message size for Server (supporting huge workflows).                                         |

server/apiserver/argoserver.go

@@ -7,6 +7,7 @@ import (
 	"net"
 	"net/http"
 	"os"
+	"strings"
 	"time"
 
 	"github.com/gorilla/handlers"
@@ -325,16 +326,21 @@ func (as *argoServer) newGRPCServer(instanceIDService instanceid.Service, offloa
 // using grpc-gateway as a proxy to the gRPC server.
 func (as *argoServer) newHTTPServer(ctx context.Context, port int, artifactServer *artifacts.ArtifactServer) *http.Server {
 	endpoint := fmt.Sprintf("localhost:%d", port)
+	ipKeyFunc := httplimit.IPKeyFunc()
+	if ipKeyFuncHeadersStr := env.GetString("IP_KEY_FUNC_HEADERS", ""); ipKeyFuncHeadersStr != "" {
+		ipKeyFuncHeaders := strings.Split(ipKeyFuncHeadersStr, ",")
+		ipKeyFunc = httplimit.IPKeyFunc(ipKeyFuncHeaders...)
+	}
 
-	ratelimit_middleware, err := httplimit.NewMiddleware(as.apiRateLimiter, httplimit.IPKeyFunc())
+	rateLimitMiddleware, err := httplimit.NewMiddleware(as.apiRateLimiter, ipKeyFunc)
 	if err != nil {
 		log.Fatal(err)
 	}
 
 	mux := http.NewServeMux()
 	httpServer := http.Server{
 		Addr:      endpoint,
-		Handler:   ratelimit_middleware.Handle(accesslog.Interceptor(mux)),
+		Handler:   rateLimitMiddleware.Handle(accesslog.Interceptor(mux)),
 		TLSConfig: as.tlsConfig,
 	}
 	dialOpts := []grpc.DialOption{