feat: APIratelimit headers and doc (#9206)

* feat: apiratelimit headers

Signed-off-by: Tianchu Zhao <evantczhao@gmail.com>

* docs: argoserver apiratelimit

Signed-off-by: Tianchu Zhao <evantczhao@gmail.com>

* fix: lint

Signed-off-by: Tianchu Zhao <evantczhao@gmail.com>

docs/argo-server.md

@@ -164,3 +164,14 @@ Argo Server does not perform authentication directly. It delegates this to eithe
 ### IP Address Logging
 
 Argo Server does not log the IP addresses of API requests. We recommend you put the Argo Server behind a load balancer, and that load balancer is configured to log the IP addresses of requests that return authentication or authorization errors.
+
+### Rate Limiting
+
+> v3.4 and after
+
+Argo Server by default rate limits to 1000 per IP per minute, you can configure it through `--api-rate-limit`. You can access additional information through the following headers.
+
+* `X-Rate-Limit-Limit` - the rate limit ceiling that is applicable for the current request.
+* `X-Rate-Limit-Remaining` - the number of requests left for the current rate-limit window.
+* `X-Rate-Limit-Reset` - the time at which the rate limit resets, specified in UTC time.
+* `Retry-After` - indicate when a client should retry requests (when the rate limit expires), in UTC time.

server/apiserver/argoserver.go

@@ -64,6 +64,7 @@ import (
 	"github.com/argoproj/argo-workflows/v3/workflow/hydrator"
 
 	limiter "github.com/sethvargo/go-limiter"
+	"github.com/sethvargo/go-limiter/httplimit"
 	"github.com/sethvargo/go-limiter/memorystore"
 )
 
@@ -318,10 +319,15 @@ func (as *argoServer) newGRPCServer(instanceIDService instanceid.Service, offloa
 func (as *argoServer) newHTTPServer(ctx context.Context, port int, artifactServer *artifacts.ArtifactServer) *http.Server {
 	endpoint := fmt.Sprintf("localhost:%d", port)
 
+	ratelimit_middleware, err := httplimit.NewMiddleware(as.apiRateLimiter, httplimit.IPKeyFunc())
+	if err != nil {
+		log.Fatal(err)
+	}
+
 	mux := http.NewServeMux()
 	httpServer := http.Server{
 		Addr:      endpoint,
-		Handler:   as.httpLimit(accesslog.Interceptor(mux)),
+		Handler:   ratelimit_middleware.Handle(accesslog.Interceptor(mux)),
 		TLSConfig: as.tlsConfig,
 	}
 	dialOpts := []grpc.DialOption{
@@ -413,26 +419,3 @@ func (as *argoServer) checkServeErr(name string, err error) {
 		log.Infof("graceful shutdown %s", name)
 	}
 }
-
-func (as *argoServer) httpLimit(next http.Handler) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		// Get the IP address for the current user.
-		ip, _, err := net.SplitHostPort(r.RemoteAddr)
-		if err != nil {
-			http.Error(w, "Internal Server Error", http.StatusInternalServerError)
-			return
-		}
-		ctx := r.Context()
-		_, _, _, ok, err := as.apiRateLimiter.Take(ctx, ip)
-		if err != nil {
-			http.Error(w, "Internal Server Error", http.StatusInternalServerError)
-			return
-		}
-		if !ok {
-			http.Error(w, http.StatusText(429), http.StatusTooManyRequests)
-			return
-		}
-
-		next.ServeHTTP(w, r)
-	})
-}

test/e2e/argo_server_test.go

@@ -1979,6 +1979,19 @@ func (s *ArgoServerSuite) TestSensorService() {
 	})
 }
 
+func (s *ArgoServerSuite) TestRateLimitHeader() {
+	s.Run("GetRateLimit", func() {
+		resp := s.e().GET("/api/v1/version").
+			Expect().
+			Status(200)
+
+		resp.Header("X-RateLimit-Limit").NotEmpty()
+		resp.Header("X-RateLimit-Remaining").NotEmpty()
+		resp.Header("X-RateLimit-Reset").NotEmpty()
+		resp.Header("Retry-After").Empty()
+	})
+}
+
 func TestArgoServerSuite(t *testing.T) {
 	suite.Run(t, new(ArgoServerSuite))
 }