Add rate limiter to receiver and event http servers

Signed-off-by: Philip Laine <[email protected]>

controllers/suite_test.go

@@ -24,9 +24,13 @@ import (
 	"net/http/httptest"
 	"path/filepath"
 	"testing"
+	"time"
 
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
+	"github.com/sethvargo/go-limiter/memorystore"
+	prommetrics "github.com/slok/go-http-metrics/metrics/prometheus"
+	"github.com/slok/go-http-metrics/middleware"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/client-go/kubernetes/scheme"
@@ -51,6 +55,7 @@ import (
 var cfg *rest.Config
 var k8sClient client.Client
 var testEnv *envtest.Environment
+var eventMdlw middleware.Middleware
 
 func TestAPIs(t *testing.T) {
 	RegisterFailHandler(Fail)
@@ -70,6 +75,12 @@ var _ = BeforeSuite(func(done Done) {
 		CRDDirectoryPaths: []string{filepath.Join("..", "config", "crd", "bases")},
 	}
 
+	eventMdlw = middleware.New(middleware.Config{
+		Recorder: prommetrics.NewRecorder(prommetrics.Config{
+			Prefix: "event",
+		}),
+	})
+
 	var err error
 	cfg, err = testEnv.Start()
 	Expect(err).ToNot(HaveOccurred())
@@ -127,10 +138,14 @@ var _ = Describe("Event handlers", func() {
 		Expect(k8sClient.Create(ctx, &provider)).To(Succeed())
 
 		By("Creating and starting event server")
+		store, err := memorystore.New(&memorystore.Config{
+			Interval: 5 * time.Minute,
+		})
+		Expect(err).ShouldNot(HaveOccurred())
 		// TODO let OS assign port number
 		eventServer := server.NewEventServer("127.0.0.1:56789", logf.Log, k8sClient)
 		stopCh = make(chan struct{})
-		go eventServer.ListenAndServe(stopCh)
+		go eventServer.ListenAndServe(stopCh, eventMdlw, store)
 	})
 
 	AfterEach(func() {

go.mod

@@ -11,11 +11,15 @@ require (
 	github.com/getsentry/sentry-go v0.10.0
 	github.com/go-logr/logr v0.3.0
 	github.com/google/go-github/v32 v32.1.0
+	github.com/gorilla/mux v1.8.0
 	github.com/hashicorp/go-retryablehttp v0.6.8
 	github.com/ktrysmt/go-bitbucket v0.6.5
 	github.com/microsoft/azure-devops-go-api/azuredevops v1.0.0-b5
 	github.com/onsi/ginkgo v1.14.1
 	github.com/onsi/gomega v1.10.2
+	github.com/prometheus/client_golang v1.7.1
+	github.com/sethvargo/go-limiter v0.6.0
+	github.com/slok/go-http-metrics v0.9.0
 	github.com/spf13/pflag v1.0.5
 	github.com/stretchr/testify v1.6.1
 	github.com/whilp/git-urls v1.0.0

internal/server/event_server.go

@@ -17,13 +17,25 @@ limitations under the License.
 package server
 
 import (
+	"bytes"
 	"context"
+	"crypto/sha256"
+	"encoding/json"
+	"fmt"
+	"io/ioutil"
 	"net/http"
 	"os"
+	"strings"
 	"time"
 
 	"github.com/go-logr/logr"
+	"github.com/sethvargo/go-limiter"
+	"github.com/sethvargo/go-limiter/httplimit"
+	"github.com/slok/go-http-metrics/middleware"
+	middlewarestd "github.com/slok/go-http-metrics/middleware/std"
 	"sigs.k8s.io/controller-runtime/pkg/client"
+
+	"github.com/fluxcd/pkg/runtime/events"
 )
 
 // EventServer handles event POST requests
@@ -43,10 +55,14 @@ func NewEventServer(port string, logger logr.Logger, kubeClient client.Client) *
 }
 
 // ListenAndServe starts the HTTP server on the specified port
-func (s *EventServer) ListenAndServe(stopCh <-chan struct{}) {
+func (s *EventServer) ListenAndServe(stopCh <-chan struct{}, mdlw middleware.Middleware, store limiter.Store) {
+	limitMiddleware, err := httplimit.NewMiddleware(store, eventKeyFunc)
+	if err != nil {
+		s.logger.Error(err, "Event server crashed")
+		os.Exit(1)
+	}
 	mux := http.NewServeMux()
-	mux.HandleFunc("/", s.handleEvent())
-
+	mux.Handle("/", middlewarestd.Handler("", mdlw, limitMiddleware.Handle(http.HandlerFunc(s.handleEvent()))))
 	srv := &http.Server{
 		Addr:    s.port,
 		Handler: mux,
@@ -70,3 +86,27 @@ func (s *EventServer) ListenAndServe(stopCh <-chan struct{}) {
 		s.logger.Info("Event server stopped")
 	}
 }
+
+func eventKeyFunc(r *http.Request) (string, error) {
+	body, err := ioutil.ReadAll(r.Body)
+	if err != nil {
+		return "", err
+	}
+
+	event := &events.Event{}
+	err = json.Unmarshal(body, event)
+	if err != nil {
+		return "", err
+	}
+
+	r.Body = ioutil.NopCloser(bytes.NewBuffer(body))
+
+	comps := []string{"event", event.InvolvedObject.Name, event.InvolvedObject.Namespace, event.InvolvedObject.Kind, event.Message}
+	revString, ok := event.Metadata["revision"]
+	if ok {
+		comps = append(comps, revString)
+	}
+	val := strings.Join(comps, "/")
+	digest := sha256.Sum256([]byte(val))
+	return fmt.Sprintf("%x", digest), nil
+}

internal/server/event_server_test.go

@@ -0,0 +1,147 @@
+/*
+Copyright 2021 The Flux authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package server
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/onsi/gomega"
+	"github.com/sethvargo/go-limiter/httplimit"
+	"github.com/sethvargo/go-limiter/memorystore"
+	corev1 "k8s.io/api/core/v1"
+
+	"github.com/fluxcd/pkg/runtime/events"
+)
+
+func TestEventKeyFunc(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+
+	// Setup middleware
+	store, err := memorystore.New(&memorystore.Config{
+		Interval: 10 * time.Minute,
+	})
+	g.Expect(err).ShouldNot(gomega.HaveOccurred())
+	middleware, err := httplimit.NewMiddleware(store, eventKeyFunc)
+	g.Expect(err).ShouldNot(gomega.HaveOccurred())
+	handler := middleware.Handle(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	// Make request
+	tests := []struct {
+		involvedObject corev1.ObjectReference
+		severity       string
+		message        string
+		rateLimit      bool
+	}{
+		{
+			involvedObject: corev1.ObjectReference{
+				APIVersion: "kustomize.toolkit.fluxcd.io/v1beta1",
+				Kind:       "Kustomization",
+				Name:       "1",
+				Namespace:  "1",
+			},
+			severity:  events.EventSeverityInfo,
+			message:   "Health check passed",
+			rateLimit: false,
+		},
+		{
+			involvedObject: corev1.ObjectReference{
+				APIVersion: "kustomize.toolkit.fluxcd.io/v1beta1",
+				Kind:       "Kustomization",
+				Name:       "1",
+				Namespace:  "1",
+			},
+			severity:  events.EventSeverityInfo,
+			message:   "Health check passed",
+			rateLimit: true,
+		},
+		{
+			involvedObject: corev1.ObjectReference{
+				APIVersion: "kustomize.toolkit.fluxcd.io/v1beta1",
+				Kind:       "Kustomization",
+				Name:       "1",
+				Namespace:  "1",
+			},
+			severity:  events.EventSeverityError,
+			message:   "Health check timed out for [Deployment 'foo/bar']",
+			rateLimit: false,
+		},
+		{
+			involvedObject: corev1.ObjectReference{
+				APIVersion: "kustomize.toolkit.fluxcd.io/v1beta1",
+				Kind:       "Kustomization",
+				Name:       "2",
+				Namespace:  "2",
+			},
+			severity:  events.EventSeverityInfo,
+			message:   "Health check passed",
+			rateLimit: false,
+		},
+		{
+			involvedObject: corev1.ObjectReference{
+				APIVersion: "kustomize.toolkit.fluxcd.io/v1beta1",
+				Kind:       "Kustomization",
+				Name:       "3",
+				Namespace:  "3",
+			},
+			severity:  events.EventSeverityInfo,
+			message:   "Health check passed",
+			rateLimit: false,
+		},
+		{
+			involvedObject: corev1.ObjectReference{
+				APIVersion: "kustomize.toolkit.fluxcd.io/v1beta1",
+				Kind:       "Kustomization",
+				Name:       "2",
+				Namespace:  "2",
+			},
+			severity:  events.EventSeverityInfo,
+			message:   "Health check passed",
+			rateLimit: true,
+		},
+	}
+	for i, tt := range tests {
+		t.Run(fmt.Sprintf("%v", i), func(t *testing.T) {
+			event := events.Event{
+				InvolvedObject: tt.involvedObject,
+				Severity:       tt.severity,
+				Message:        tt.message,
+			}
+			eventData, err := json.Marshal(event)
+			g.Expect(err).ShouldNot(gomega.HaveOccurred())
+
+			req := httptest.NewRequest("POST", "/", bytes.NewBuffer(eventData))
+			g.Expect(err).ShouldNot(gomega.HaveOccurred())
+			res := httptest.NewRecorder()
+			handler.ServeHTTP(res, req)
+
+			if tt.rateLimit {
+				g.Expect(res.Code).Should(gomega.Equal(429))
+				g.Expect(res.Header().Get("X-Ratelimit-Remaining")).Should(gomega.Equal("0"))
+			} else {
+				g.Expect(res.Code).Should(gomega.Equal(200))
+			}
+		})
+	}
+}

internal/server/receiver_server.go

@@ -18,11 +18,19 @@ package server
 
 import (
 	"context"
+	"crypto/sha256"
+	"fmt"
 	"net/http"
+	"net/url"
 	"os"
+	"strings"
 	"time"
 
 	"github.com/go-logr/logr"
+	"github.com/sethvargo/go-limiter"
+	"github.com/sethvargo/go-limiter/httplimit"
+	"github.com/slok/go-http-metrics/middleware"
+	middlewarestd "github.com/slok/go-http-metrics/middleware/std"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 )
 
@@ -43,11 +51,14 @@ func NewReceiverServer(port string, logger logr.Logger, kubeClient client.Client
 }
 
 // ListenAndServe starts the HTTP server on the specified port
-func (s *ReceiverServer) ListenAndServe(stopCh <-chan struct{}) {
+func (s *ReceiverServer) ListenAndServe(stopCh <-chan struct{}, mdlw middleware.Middleware, store limiter.Store) {
+	limitMiddleware, err := httplimit.NewMiddleware(store, receiverKeyFunc)
+	if err != nil {
+		s.logger.Error(err, "Receiver server crashed")
+		os.Exit(1)
+	}
 	mux := http.DefaultServeMux
-
-	mux.HandleFunc("/hook/", s.handlePayload())
-
+	mux.Handle("/hook/", middlewarestd.Handler("", mdlw, limitMiddleware.Handle(http.HandlerFunc(s.handlePayload()))))
 	srv := &http.Server{
 		Addr:    s.port,
 		Handler: mux,
@@ -71,3 +82,10 @@ func (s *ReceiverServer) ListenAndServe(stopCh <-chan struct{}) {
 		s.logger.Info("Receiver server stopped")
 	}
 }
+
+func receiverKeyFunc(r *http.Request) (string, error) {
+	id := url.PathEscape(strings.TrimLeft(r.RequestURI, "/hook/"))
+	val := strings.Join([]string{"receiver", id}, "/")
+	digest := sha256.Sum256([]byte(val))
+	return fmt.Sprintf("%x", digest), nil
+}