feat(webhook): add rate limiting to webhook endpoint (#1210)

Signed-off-by: Christopher Coco <[email protected]>
Signed-off-by: cjcocokrisp <[email protected]>

Author: cjcocokrisp

cmd/run.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math"
 	"os"
 	"strings"
 	"sync"
@@ -26,6 +27,8 @@ import (
 
 	"golang.org/x/sync/semaphore"
 
+	"go.uber.org/ratelimit"
+
 	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
@@ -216,6 +219,10 @@ func newRunCommand() *cobra.Command {
 				log.Infof("Starting webhook server on port %d", webhookCfg.Port)
 				webhookServer = webhook.NewWebhookServer(webhookCfg.Port, handler, cfg.KubeClient, argoClient)
 
+				if webhookCfg.RateLimitNumAllowedRequests > 0 {
+					webhookServer.RateLimiter = ratelimit.New(webhookCfg.RateLimitNumAllowedRequests, ratelimit.Per(time.Hour))
+				}
+
 				// Set updater config
 				webhookServer.UpdaterConfig = &argocd.UpdateConfiguration{
 					NewRegFN:               registry.NewClient,
@@ -337,6 +344,7 @@ func newRunCommand() *cobra.Command {
 	runCmd.Flags().StringVar(&webhookCfg.GHCRSecret, "ghcr-webhook-secret", env.GetStringVal("GHCR_WEBHOOK_SECRET", ""), "Secret for validating GitHub Container Registry webhooks")
 	runCmd.Flags().StringVar(&webhookCfg.QuaySecret, "quay-webhook-secret", env.GetStringVal("QUAY_WEBHOOK_SECRET", ""), "Secret for validating Quay webhooks")
 	runCmd.Flags().StringVar(&webhookCfg.HarborSecret, "harbor-webhook-secret", env.GetStringVal("HARBOR_WEBHOOK_SECRET", ""), "Secret for validating Harbor webhooks")
+	runCmd.Flags().IntVar(&webhookCfg.RateLimitNumAllowedRequests, "webhook-ratelimit-allowed", env.ParseNumFromEnv("WEBHOOK_RATELIMIT_ALLOWED", 0, 0, math.MaxInt), "The number of allowed requests in an hour for webhook rate limiting, setting to 0 disables ratelimiting")
 
 	return runCmd
 }

cmd/webhook.go

@@ -4,12 +4,14 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math"
 	"os"
 	"os/signal"
 	"strconv"
 	"strings"
 	"syscall"
 	"text/template"
+	"time"
 
 	"github.com/argoproj-labs/argocd-image-updater/pkg/argocd"
 	"github.com/argoproj-labs/argocd-image-updater/pkg/common"
@@ -21,16 +23,18 @@ import (
 
 	"github.com/argoproj/argo-cd/v2/util/askpass"
 	"github.com/spf13/cobra"
+	"go.uber.org/ratelimit"
 	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
 // WebhookConfig holds the options for the webhook server
 type WebhookConfig struct {
-	Port         int
-	DockerSecret string
-	GHCRSecret   string
-	QuaySecret   string
-	HarborSecret string
+	Port                        int
+	DockerSecret                string
+	GHCRSecret                  string
+	QuaySecret                  string
+	HarborSecret                string
+	RateLimitNumAllowedRequests int
 }
 
 // NewWebhookCommand creates a new webhook command
@@ -190,6 +194,7 @@ Supported registries:
 	webhookCmd.Flags().StringVar(&webhookCfg.GHCRSecret, "ghcr-webhook-secret", env.GetStringVal("GHCR_WEBHOOK_SECRET", ""), "Secret for validating GitHub Container Registry webhooks")
 	webhookCmd.Flags().StringVar(&webhookCfg.QuaySecret, "quay-webhook-secret", env.GetStringVal("QUAY_WEBHOOK_SECRET", ""), "Secret for validating Quay webhooks")
 	webhookCmd.Flags().StringVar(&webhookCfg.HarborSecret, "harbor-webhook-secret", env.GetStringVal("HARBOR_WEBHOOK_SECRET", ""), "Secret for validating Harbor webhooks")
+	webhookCmd.Flags().IntVar(&webhookCfg.RateLimitNumAllowedRequests, "webhook-ratelimit-allowed", env.ParseNumFromEnv("WEBHOOK_RATELIMIT_ALLOWED", 0, 0, math.MaxInt), "The number of allowed requests in an hour for webhook rate limiting, setting to 0 disables ratelimiting")
 
 	return webhookCmd
 }
@@ -238,6 +243,10 @@ func runWebhook(cfg *ImageUpdaterConfig, webhookCfg *WebhookConfig) error {
 	// Create webhook server
 	server := webhook.NewWebhookServer(webhookCfg.Port, handler, cfg.KubeClient, cfg.ArgoClient)
 
+	if webhookCfg.RateLimitNumAllowedRequests > 0 {
+		server.RateLimiter = ratelimit.New(webhookCfg.RateLimitNumAllowedRequests, ratelimit.Per(time.Hour))
+	}
+
 	// Set updater config
 	server.UpdaterConfig = &argocd.UpdateConfiguration{
 		NewRegFN:               registry.NewClient,
@@ -273,5 +282,6 @@ func runWebhook(cfg *ImageUpdaterConfig, webhookCfg *WebhookConfig) error {
 	if err := server.Stop(); err != nil {
 		log.Errorf("Error stopping webhook server: %v", err)
 	}
+
 	return nil
 }

docs/configuration/webhook.md

@@ -155,6 +155,21 @@ to apply them yourself.
 
 They are located in the `manifets/base/networking` directory.
 
+## Rate Limiting
+
+To prevent overloading from the `/webhook` endpoint which could cause Image 
+Updater to use too many resources rate limiting is implemented for the endpoint.
+
+The rate limiting allows for a certain amount of requests per hour. This setting
+is configurable and can be set with the configuration value below. If you go over
+the limit the request will wait until it is allowed. The rate limit value defaults 
+to 0 which means that it is disabled.
+```yaml
+data:
+  # How many requests can be made per second. The default is 0 meaning disabled.
+  webhook.ratelimit-allowed: <SOME_NUMBER>
+```
+
 ## Environment Variables
 
 The flags for both the `run` and `webhook` CLI commands can also be set via 
@@ -168,6 +183,7 @@ environment variables. Below is the list of which variables correspond to which
 |`GHCR_WEBHOOK_SECRET` |`--gchr-webhook-secret`|
 |`HARBOR_WEBHOOK_SECRET` |`--harbor-webhook-secret`|
 |`QUAY_WEBHOOK_SECRET` |`--quay-webhook-secret`|
+|`WEBHOOK_RATELIMIT_ALLOWED`|`--webhook-ratelimit-allowed`|
 
 ## Adding Support For Other Registries
 

docs/install/cmd/run.md

@@ -87,14 +87,20 @@ to workloads it found in need for upgrade.
 
 Secret for validating Docker Hub webhooks.
 
+Can also be set with the `DOCKER_WEBHOOK_SECRET` environment variable.
+
 **--enable-webhook *enabled***
 
 Enable webhook server for receiving registry events.
 
+Can also be set with the `ENABLE_WEBHOOK` environment variable.
+
 **--ghcr-webhook-secret *secret***
 
 Secret for validating GitHub container registry webhooks.
 
+Can also be set with the `GHCR_WEBHOOK_SECRET` environment variable.
+
 **--git-commit-email *email***
 
 E-Mail address to use for Git commits (default "[email protected]")
@@ -131,6 +137,8 @@ Can also be set using the *GIT_COMMIT_USER* environment variable.
 
 Secret for validating Harbor webhooks
 
+Can also be set with the `HARBOR_WEBHOOK_SECRET` environment variable.
+
 **--health-port *port***
 
 Specifies the local port to bind the health server to. The health server is
@@ -210,6 +218,8 @@ Argo CD Image Updater will exit after the first update cycle.
 
 Secret for validating Quay webhooks.
 
+Can also be set with the `QUAY_WEBHOOK_SECRET` environment variable.
+
 **--registries-conf-path *path***
 
 Load the registry configuration from file at *path*. Defaults to the path
@@ -225,4 +235,13 @@ whether to perform a cache warm-up on startup (default true)
 
 Port to listen on for webhook events (default 8082)
 
+Can also be set with the `WEBHOOK_PORT` environment variable.
+
+**--webhook-ratelimit-allowed *numRequests***
+
+The number of allowed requests in an hour for webhook rate limiting, setting to 0
+means that the rate limiting is disabled.
+
+Can also be set with the `WEBHOOK_RATELIMIT_ALLOWED` environment variable.
+
 [label selector syntax]: https://kubernetes.io/docs/concepts/overview/working-with-objects/labels/#label-selectors

docs/install/cmd/webhook.md

@@ -89,10 +89,14 @@ for images can only be specified from an environment variable.
 
 Secret for validating Docker Hub webhooks.
 
+Can also be set with the `DOCKER_WEBHOOK_SECRET` environment variable.
+
 **--ghcr-webhook-secret *secret***
 
 Secret for validating GitHub container registry secrets.
 
+Can also be set with the `GHCR_WEBHOOK_SECRET` environment variable.
+
 **--git-commit-email *email***
 
 E-Mail address to use for Git commits (default "[email protected]")
@@ -129,6 +133,8 @@ Can also be set using the *GIT_COMMIT_USER* environment variable.
 
 Secret for validating Harbor webhooks
 
+Can also be set with the `HARBOR_WEBHOOK_SECRET` environment variable.
+
 **-h, --help**
 
 help for run
@@ -175,6 +181,8 @@ application processing, specify a number of `1`.
 
 Secret for validating Quay webhooks
 
+Can also be set with the `QUAY_WEBHOOK_SECRET` environment variable.
+
 **--registries-conf-path *path***
 
 Load the registry configuration from file at *path*. Defaults to the path
@@ -186,4 +194,13 @@ default configuration should be used instead, specify the empty string, i.e.
 
 Port to listen on for webhook events (default 8080)
 
+Can also be set with the `WEBHOOK_PORT` environment variable.
+
+**--webhook-ratelimit-allowed *numRequests***
+
+The number of allowed requests in an hour for webhook rate limiting, setting to 0
+means that the rate limiting is disabled.
+
+Can also be set with the `WEBHOOK_RATELIMIT_ALLOWED` environment variable.
+
 [label selector syntax]: https://kubernetes.io/docs/concepts/overview/working-with-objects/labels/#label-selectors

manifests/base/deployment/argocd-image-updater-deployment.yaml

@@ -149,6 +149,12 @@ spec:
                 name: argocd-image-updater-secret
                 key: webhook.harbor-secret
                 optional: true
+        - name: WEBHOOK_RATELIMIT_ALLOWED
+          valueFrom:
+            configMapKeyRef:
+                name: argocd-image-updater-config
+                key: webhook.ratelimit-allowed
+                optional: true
         livenessProbe:
           httpGet:
             path: /healthz

manifests/install.yaml

@@ -257,6 +257,12 @@ spec:
               key: webhook.harbor-secret
               name: argocd-image-updater-secret
               optional: true
+        - name: WEBHOOK_RATELIMIT_ALLOWED
+          valueFrom:
+            configMapKeyRef:
+              key: webhook.ratelimit-allowed
+              name: argocd-image-updater-config
+              optional: true
         image: quay.io/argoprojlabs/argocd-image-updater:latest
         imagePullPolicy: Always
         livenessProbe:

pkg/webhook/server.go

@@ -13,6 +13,7 @@ import (
 	"github.com/argoproj-labs/argocd-image-updater/registry-scanner/pkg/log"
 
 	"github.com/argoproj/argo-cd/v2/pkg/apis/application/v1alpha1"
+	"go.uber.org/ratelimit"
 )
 
 // WebhookServer manages webhook endpoints and triggers update checks
@@ -33,16 +34,19 @@ type WebhookServer struct {
 	mutex sync.Mutex
 	// mutex for concurrent repo access
 	syncState *argocd.SyncIterationState
+	// rate limiter to limit requests in an interval
+	RateLimiter ratelimit.Limiter
 }
 
 // NewWebhookServer creates a new webhook server
 func NewWebhookServer(port int, handler *WebhookHandler, kubeClient *kube.ImageUpdaterKubernetesClient, argoClient argocd.ArgoCD) *WebhookServer {
 	return &WebhookServer{
-		Port:       port,
-		Handler:    handler,
-		KubeClient: kubeClient,
-		ArgoClient: argoClient,
-		syncState:  argocd.NewSyncIterationState(),
+		Port:        port,
+		Handler:     handler,
+		KubeClient:  kubeClient,
+		ArgoClient:  argoClient,
+		syncState:   argocd.NewSyncIterationState(),
+		RateLimiter: nil,
 	}
 }
 
@@ -96,6 +100,10 @@ func (s *WebhookServer) handleWebhook(w http.ResponseWriter, r *http.Request) {
 
 	// Process webhook asynchronously
 	go func() {
+		if s.RateLimiter != nil {
+			s.RateLimiter.Take()
+		}
+
 		err := s.processWebhookEvent(event)
 		if err != nil {
 			logCtx.Errorf("Failed to process webhook event: %v", err)

pkg/webhook/server_test.go

@@ -72,6 +72,15 @@ var (
 	}
 )
 
+type mockRateLimiter struct {
+	Called bool
+}
+
+func (m *mockRateLimiter) Take() time.Time {
+	m.Called = true
+	return time.Now()
+}
+
 // Helper function to create a mock server
 func createMockServer(t *testing.T, port int) *WebhookServer {
 	handler := NewWebhookHandler()
@@ -455,3 +464,39 @@ func TestParseImageList(t *testing.T) {
 		assert.Equal(t, "baz", imgs[0].KustomizeImage.ImageName)
 	})
 }
+
+// TestWebhookServerRateLimit tests to see if the webhook endpoint's rate limiting functionality works
+func TestWebhookServerRateLimit(t *testing.T) {
+	server := createMockServer(t, 8080)
+	mockArgoClient := server.ArgoClient.(*mocks.ArgoCD)
+	mockArgoClient.On("ListApplications", mock.Anything).Return([]v1alpha1.Application{}, nil).Maybe()
+
+	handler := NewDockerHubWebhook("")
+	assert.NotNil(t, handler, "Docker handler was nil")
+
+	server.Handler.RegisterHandler(handler)
+
+	mock := &mockRateLimiter{}
+	server.RateLimiter = mock
+
+	body := []byte(`{
+		"repository": {
+			"repo_name": "somepersononthisfakeregistry/myimagethatdoescoolstuff",
+			"name": "myimagethatdoescoolstuff",
+			"namespace": "randomplaceincluster"
+		},
+		"push_data": {
+			"tag": "v12.0.9"
+		}
+	}`)
+
+	req := httptest.NewRequest(http.MethodPost, "/webhook?type=docker", bytes.NewReader(body))
+	rec := httptest.NewRecorder()
+
+	server.handleWebhook(rec, req)
+
+	// Wait for thread to call it.
+	time.Sleep(time.Second)
+
+	assert.True(t, mock.Called, "Take was not called")
+}