feat(webhook): add rate limiting to the webhook endpoint

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

Author: cjcocokrisp

cmd/run.go

@@ -216,6 +216,11 @@ 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.RateLimitEnabled {
+					limiter := webhook.NewRateLimiter(webhookCfg.RateLimitNumAllowedRequests, webhookCfg.RateLimitWindow, webhookCfg.RateLimitCleanUpInterval)
+					webhookServer.RateLimiter = limiter
+				}
+
 				// Set updater config
 				webhookServer.UpdaterConfig = &argocd.UpdateConfiguration{
 					NewRegFN:               registry.NewClient,
@@ -271,6 +276,9 @@ func newRunCommand() *cobra.Command {
 						if err := webhookServer.Stop(); err != nil {
 							log.Errorf("Error stopping webhook server: %v", err)
 						}
+						if webhookCfg.RateLimitEnabled {
+							webhookServer.RateLimiter.StopCleanUp()
+						}
 					}
 					return nil
 				case err := <-whErrCh:
@@ -337,6 +345,10 @@ 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().BoolVar(&webhookCfg.RateLimitEnabled, "enable-webhook-ratelimit", false, "Enable rate limiting for the webhook endpoint")
+	runCmd.Flags().IntVar(&webhookCfg.RateLimitNumAllowedRequests, "webhook-ratelimit-num-allowed", 100, "The number of allowed requests in a window for webhook rate limiting")
+	runCmd.Flags().DurationVar(&webhookCfg.RateLimitWindow, "webhook-ratelimit-window", 2*time.Minute, "The duration for the window for the webhook rate limiting")
+	runCmd.Flags().DurationVar(&webhookCfg.RateLimitCleanUpInterval, "webhook-ratelimit-cleanup-interval", 1*time.Hour, "How often the rate limiter cleans up stale clients")
 
 	return runCmd
 }

cmd/webhook.go

@@ -10,6 +10,7 @@ import (
 	"strings"
 	"syscall"
 	"text/template"
+	"time"
 
 	"github.com/argoproj-labs/argocd-image-updater/pkg/argocd"
 	"github.com/argoproj-labs/argocd-image-updater/pkg/common"
@@ -26,11 +27,15 @@ import (
 
 // 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
+	RateLimitEnabled            bool
+	RateLimitNumAllowedRequests int
+	RateLimitWindow             time.Duration
+	RateLimitCleanUpInterval    time.Duration
 }
 
 // NewWebhookCommand creates a new webhook command
@@ -190,6 +195,10 @@ 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().BoolVar(&webhookCfg.RateLimitEnabled, "enable-webhook-ratelimit", false, "Enable rate limiting for the webhook endpoint")
+	webhookCmd.Flags().IntVar(&webhookCfg.RateLimitNumAllowedRequests, "webhook-ratelimit-num-allowed", 100, "The number of allowed requests in a window for webhook rate limiting")
+	webhookCmd.Flags().DurationVar(&webhookCfg.RateLimitWindow, "webhook-ratelimit-window", 2*time.Minute, "The duration for the window for the webhook rate limiting")
+	webhookCmd.Flags().DurationVar(&webhookCfg.RateLimitCleanUpInterval, "webhook-ratelimit-cleanup-interval", 1*time.Hour, "How often the rate limiter cleans up stale clients")
 
 	return webhookCmd
 }
@@ -238,6 +247,11 @@ func runWebhook(cfg *ImageUpdaterConfig, webhookCfg *WebhookConfig) error {
 	// Create webhook server
 	server := webhook.NewWebhookServer(webhookCfg.Port, handler, cfg.KubeClient, cfg.ArgoClient)
 
+	if webhookCfg.RateLimitEnabled {
+		limiter := webhook.NewRateLimiter(webhookCfg.RateLimitNumAllowedRequests, webhookCfg.RateLimitWindow, webhookCfg.RateLimitCleanUpInterval)
+		server.RateLimiter = limiter
+	}
+
 	// Set updater config
 	server.UpdaterConfig = &argocd.UpdateConfiguration{
 		NewRegFN:               registry.NewClient,
@@ -273,5 +287,10 @@ func runWebhook(cfg *ImageUpdaterConfig, webhookCfg *WebhookConfig) error {
 	if err := server.Stop(); err != nil {
 		log.Errorf("Error stopping webhook server: %v", err)
 	}
+
+	if webhookCfg.RateLimitEnabled {
+		server.RateLimiter.StopCleanUp()
+	}
+
 	return nil
 }

pkg/webhook/ratelimit.go

@@ -0,0 +1,91 @@
+package webhook
+
+import (
+	"sync"
+	"time"
+)
+
+// RateLimiter implements a sliding window rate limiting algorithm
+type RateLimiter struct {
+	mu       sync.Mutex
+	clients  map[string][]int64
+	lastSeen map[string]time.Time
+	window   time.Duration
+	allowed  int
+	done     chan bool
+}
+
+func NewRateLimiter(numRequets int, window time.Duration, cleanUpInterval time.Duration) *RateLimiter {
+	limiter := RateLimiter{
+		clients:  make(map[string][]int64),
+		lastSeen: make(map[string]time.Time),
+		window:   window,
+		allowed:  numRequets,
+		done:     make(chan bool),
+	}
+	go limiter.StartCleanUp(cleanUpInterval)
+	return &limiter
+}
+
+// Checks to see if a client has gone over the limit
+func (rl *RateLimiter) Allow(clientIP string) bool {
+	now := time.Now()
+	rl.mu.Lock()
+	defer rl.mu.Unlock()
+
+	if _, ok := rl.clients[clientIP]; !ok {
+		rl.clients[clientIP] = []int64{}
+	}
+
+	allow := false
+	windowStart := now.Unix() - int64(rl.window.Seconds())
+	filtered := []int64{}
+	for _, ts := range rl.clients[clientIP] {
+		if ts > windowStart {
+			filtered = append(filtered, ts)
+		}
+	}
+	rl.clients[clientIP] = filtered
+
+	if len(filtered) < rl.allowed {
+		rl.clients[clientIP] = append(filtered, now.Unix())
+		rl.lastSeen[clientIP] = now
+		allow = true
+	}
+
+	return allow
+}
+
+// Cleans up the clients map at an interval to prevent stale clients from taking up memory
+func (rl *RateLimiter) StartCleanUp(interval time.Duration) {
+	ticker := time.NewTicker(interval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-rl.done:
+			return
+		case <-ticker.C:
+			rl.CleanUp()
+		}
+	}
+}
+
+// Cleans up any clients that have not made a request in an amount of time over the window
+func (rl *RateLimiter) CleanUp() {
+	rl.mu.Lock()
+	defer rl.mu.Unlock()
+
+	for k, v := range rl.lastSeen {
+		now := time.Now()
+		if v.Add(rl.window).Before(now) {
+			delete(rl.clients, k)
+			delete(rl.lastSeen, k)
+		}
+	}
+}
+
+// Stop the clean up goroutine
+func (rl *RateLimiter) StopCleanUp() {
+	rl.done <- true
+}