Add rate limiting to the baker

alinz · Jan 28, 2024 · e2de7cb · e2de7cb
1 parent e2a61c5
commit e2de7cb
Show file tree

Hide file tree

Showing 15 changed files with 424 additions and 71 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -1,5 +1,5 @@
 # build stage
-FROM golang:1.20-alpine AS builder
+FROM golang:1.21-alpine AS builder
 ARG GIT_COMMIT
 ARG VERSION
 RUN apk --no-cache add build-base git mercurial gcc

diff --git a/README.md b/README.md
@@ -9,23 +9,24 @@
 
 # Introduction
 
-Baker.go is a dynamic http reverse proxy designed to be highly extensible.
+Baker.go is a dynamic HTTP reverse proxy designed to be highly extensible.
 
 # Features
 
 - [x] Include Docker driver to listen to Docker's events
-- [x] Has exposed a driver interface which can be easily hook to other orchestration engines
-- [x] Dynamic configuration, no need to restart reverse proxy in order to change the configuration
+- [x] Has exposed a driver interface that can be easily hooked to other orchestration engines
+- [x] Dynamic configuration, no need to restart reverse proxy to change the configuration
 - [x] Uses a custom trie data structure, to compute fast path pattern matching
-- [x] It can be use as library as it has implemented `http.Handler` interface
+- [x] It can be used as a library as it has implemented HTTP`.`Handler` interface
 - [x] Highly extendable as most of the components have exposed interfaces
-- [x] Middleware like feature to change the incoming and outgoing traffics
+- [x] Middleware-like feature to change the incoming and outgoing traffic
 - [x] load balancing by default
 - [x] Automatically updates and creates SSL certificates using `Let's Encrypt`
+- [x] Configurable Rate Limiter per Domain and Path
 
 # Usage
 
-First we need to run Baker inside docker. The following `docker-compose.yml`
+First, we need to run Baker inside docker. The following `docker-compose.yml`
 
 ```yml
 version: "3.5"
@@ -60,7 +61,7 @@ networks:
     driver: bridge
 ```
 
-Then for each service, a following `docker-compose` can be used. The only requirements is labels and networks. Make sure both baker and service has the same network interface
+Then for each service, the following `docker-compose` can be used. The only requirements are labels and networks. Make sure both baker and service have the same network interface
 
 ```yml
 version: "3.5"
@@ -84,7 +85,7 @@ networks:
       name: baker_net
 ```
 
-The service, should expose a REST endpoint which returns a configuration, the configuration endpoint act as a health check and providing realtime configuration:
+The service should expose a REST endpoint that returns a configuration, the configuration endpoint acts as a health check and provides real-time configuration:
 
 ```json
 [
@@ -104,8 +105,8 @@ The service, should expose a REST endpoint which returns a configuration, the co
     "ready": true,
     "rules": [
       {
-        "name": "ReplacePath",
-        "config": {
+        "type": "ReplacePath",
+        "args": {
           "search": "/sample1",
           "replace": "",
           "times": 1
@@ -122,14 +123,13 @@ At the moment, there are 2 middlewares provided by default
 
 ### ReplacePath
 
-Remove a specific path from incoming request. Service will be receiving the modified path.
-
-in order to use this middleware, simply add the following rule to rules section of the configuration
+Remove a specific path from an incoming request. Service will be receiving the modified path.
+to use this middleware, simply add the following rule to the rules section of the configuration
 
 ```json
 {
-  "name": "ReplacePath",
-  "config": {
+  "type": "ReplacePath",
+  "args": {
     "search": "/sample1",
     "replace": "",
     "times": 1
@@ -139,16 +139,32 @@ in order to use this middleware, simply add the following rule to rules section
 
 ### AppendPath
 
-Add a path at the beginning and end of path
-
-in order to use this middleware, simply add the following rule to rules section of the configuration
+Add a path at the beginning and end of the path
+to use this middleware, simply add the following rule to the rules section of the configuration
 
 ```json
 {
-  "name": "AppendPath",
-  "config": {
+  "type": "AppendPath",
+  "args": {
     "begin": "/begin",
     "end": "/end"
   }
 }
 ```
+
+### RateLimiter
+
+Add a rate limiter for a specific domain and path
+to use this middleware, simply add the following rule to the riles sections of the configuration
+
+```json
+{
+  "type": "RateLimiter",
+  "args": {
+    "request_limit": 100,
+    "window_duration": "60s"
+  }
+}
+```
+
+the above configuration means, in each 1 min, 100 request should be routed per individual IP address
diff --git a/baker.go b/baker.go
@@ -7,6 +7,7 @@ import (
 	"net/http/httputil"
 	"net/netip"
 	"net/url"
+	"strings"
 	"time"
 
 	"github.com/alinz/baker.go/pkg/collection"
@@ -27,6 +28,15 @@ type Endpoint struct {
 	Ready  bool   `json:"ready"`
 }
 
+func (e *Endpoint) getHashKey() string {
+	var sb strings.Builder
+
+	sb.WriteString(e.Domain)
+	sb.WriteString(e.Path)
+
+	return sb.String()
+}
+
 type Container struct {
 	ID   string         `json:"id"`
 	Addr netip.AddrPort `json:"addr"`
@@ -118,7 +128,7 @@ func (s *Service) Add(container *Container, endpoint *Endpoint) {
 	})
 }
 
-func (s *Service) Remove(container *Container) {
+func (s *Service) Remove(container *Container) int {
 	value, ok := s.containers.Get(container.ID)
 	if ok {
 		log.Info().
@@ -128,7 +138,7 @@ func (s *Service) Remove(container *Container) {
 			Msg("an exisiting container is removed")
 	}
 
-	s.containers.Remove(container.ID)
+	return s.containers.Remove(container.ID)
 }
 
 func (s *Service) Select() (*Container, *Endpoint, bool) {
@@ -146,12 +156,14 @@ func NewService() *Service {
 }
 
 type Server struct {
-	domains    *Domains
-	rules      map[string]rule.BuilderFunc
-	containers *collection.Set[string, *Container]
-	done       chan struct{}
-	http       httpclient.GetterFunc
-	refMap     *collection.Map[*value]
+	domains            *Domains
+	rules              map[string]rule.BuilderFunc
+	pingDuration       time.Duration
+	containers         *collection.Set[string, *Container]
+	done               chan struct{}
+	http               httpclient.GetterFunc
+	refMap             *collection.Map[*value]
+	middlewareCacheMap *collection.Map[rule.Middleware]
 }
 
 var _ http.Handler = &Server{}
@@ -161,7 +173,7 @@ func (s *Server) pinger() {
 		select {
 		case <-s.done:
 			return
-		case <-time.After(10 * time.Second):
+		case <-time.After(s.pingDuration):
 			s.containers.Iterate(func(id string, container *Container) bool {
 				configPath := fmt.Sprintf("http://%s%s", container.Addr, container.Path)
 				body, err := s.http(configPath)
@@ -242,7 +254,7 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		},
 	}
 
-	rules, err := s.getMiddlewares(endpoint.Rules)
+	rules, err := s.getMiddlewares(endpoint)
 	if err != nil {
 		w.WriteHeader(http.StatusInternalServerError)
 		json.NewEncoder(w).Encode(struct {
@@ -262,22 +274,38 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	s.apply(proxy, rules...).ServeHTTP(w, r)
 }
 
-func (s *Server) getMiddlewares(rules []Rule) ([]rule.Middleware, error) {
-	if len(rules) == 0 {
+func (s *Server) getMiddlewares(endpoint *Endpoint) ([]rule.Middleware, error) {
+	if len(endpoint.Rules) == 0 {
 		return rule.Empty, nil
 	}
 
 	middlewares := make([]rule.Middleware, 0)
 
-	for _, rule := range rules {
-		builder, ok := s.rules[rule.Type]
+	for _, r := range endpoint.Rules {
+		builder, ok := s.rules[r.Type]
 		if !ok {
-			return nil, fmt.Errorf("failed to find rule builder for %s", rule.Type)
+			return nil, fmt.Errorf("failed to find rule builder for %s", r.Type)
 		}
 
-		middleware, err := builder(rule.Args)
+		middleware, err := builder(r.Args)
 		if err != nil {
-			return nil, fmt.Errorf("failed to parse args for rule %s: %w", rule.Type, err)
+			return nil, fmt.Errorf("failed to parse args for rule %s: %w", r.Type, err)
+		}
+
+		if middleware.IsCachable() {
+			middleware = s.middlewareCacheMap.GetAndUpdate(endpoint.getHashKey(), func(old rule.Middleware, found bool) rule.Middleware {
+				// NOTE: the reason we are doing this is because we want to update the middleware
+				// and we don;t want to recreate some internal state of the middleware over and over
+				// The responsibility of initializing the internal state of middleware is on the
+				// UpdateMiddleware method.
+				var current rule.Middleware
+				if found {
+					current = old
+				} else {
+					current = middleware
+				}
+				return current.UpdateMiddelware(middleware)
+			})
 		}
 
 		middlewares = append(middlewares, middleware)
@@ -294,18 +322,46 @@ func (s *Server) apply(next http.Handler, rules ...rule.Middleware) http.Handler
 	return next
 }
 
-func New(containers <-chan *Container, rules ...rule.RegisterFunc) *Server {
-	s := &Server{
-		domains:    NewDomains(),
-		rules:      make(map[string]rule.BuilderFunc),
-		containers: collection.NewSet[string, *Container](),
-		done:       make(chan struct{}, 1),
-		http:       httpclient.New(),
-		refMap:     collection.NewMap[*value](),
+type bakerOption struct {
+	rules        map[string]rule.BuilderFunc
+	pingDuration time.Duration
+}
+
+type bakerOptionFunc func(*bakerOption)
+
+func WithPingDuration(d time.Duration) bakerOptionFunc {
+	return func(o *bakerOption) {
+		o.pingDuration = d
+	}
+}
+
+func WithRules(rules ...rule.RegisterFunc) bakerOptionFunc {
+	return func(o *bakerOption) {
+		for _, rule := range rules {
+			rule(o.rules)
+		}
+	}
+}
+
+func New(containers <-chan *Container, optFuncs ...bakerOptionFunc) *Server {
+	opt := &bakerOption{
+		rules:        make(map[string]rule.BuilderFunc),
+		pingDuration: 10 * time.Second,
 	}
 
-	for _, rule := range rules {
-		rule(s.rules)
+	for _, optFunc := range optFuncs {
+		optFunc(opt)
+	}
+
+	s := &Server{
+		domains:            NewDomains(),
+		rules:              opt.rules,
+		pingDuration:       opt.pingDuration,
+		containers:         collection.NewSet[string, *Container](),
+		done:               make(chan struct{}, 1),
+		http:               httpclient.New(),
+		refMap:             collection.NewMap[*value](),
+		middlewareCacheMap: collection.NewMap[rule.Middleware](),
 	}
 
 	go s.pinger()
@@ -333,10 +389,16 @@ func New(containers <-chan *Container, rules ...rule.RegisterFunc) *Server {
 
 					s.refMap.Delete(container.ID)
 
-					s.domains.
+					remaining := s.domains.
 						Paths(value.endpoint.Domain, false).
 						Service(value.endpoint.Path, false).
 						Remove(value.container)
+
+					// NOTE: if there is no more containers for this endpoint
+					// we can remove the middleware from the cache
+					if remaining == 0 {
+						s.middlewareCacheMap.Delete(value.endpoint.getHashKey())
+					}
 				}
 			}
 		}

diff --git a/baker_test.go b/baker_test.go
@@ -1,9 +1,15 @@
 package baker_test
 
 import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
 	"testing"
+	"time"
 
 	"github.com/alinz/baker.go"
+	"github.com/alinz/baker.go/confutil"
+	"github.com/alinz/baker.go/rule"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -41,3 +47,37 @@ func TestDomains(t *testing.T) {
 		assert.Equal(t, endpoint1, endpoint)
 	}
 }
+
+func TestBaker(t *testing.T) {
+	containers := MockDriver(t, confutil.NewEndpoints().New("example.com", "/*", true))
+
+	baker := baker.New(
+		containers,
+		baker.WithPingDuration(2*time.Second),
+		baker.WithRules(
+			rule.RegisterAppendPath(),
+			rule.RegisterReplacePath(),
+			rule.RegisterRateLimiter(),
+		),
+	)
+
+	s := httptest.NewServer(baker)
+	t.Cleanup(s.Close)
+
+	time.Sleep(3 * time.Second)
+
+	httpClient := http.Client{}
+	req, err := http.NewRequest("GET", fmt.Sprintf("%s/manifest.json", s.URL), nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	req.Host = "example.com"
+
+	resp, err := httpClient.Do(req)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	assert.Equal(t, http.StatusOK, resp.StatusCode)
+}
diff --git a/cmd/baker/main.go b/cmd/baker/main.go
@@ -80,8 +80,12 @@ https://github.com/alinz/baker.go
 
 	baker := baker.New(
 		containers,
-		rule.RegisterAppendPath(),
-		rule.RegisterReplacePath(),
+		baker.WithPingDuration(10*time.Second),
+		baker.WithRules(
+			rule.RegisterAppendPath(),
+			rule.RegisterReplacePath(),
+			rule.RegisterRateLimiter(),
+		),
 	)
 
 	if acmeEnable {