controller: support rate-limiting incoming requests

helps constraining resource usage and against flood attacks.

Signed-off-by: Ramkumar Chinchani <rchincha@cisco.com>

README.md

@@ -40,6 +40,7 @@ https://zotregistry.io
   * Layer deduplication using hard links when content is identical
 * Serve [multiple storage paths (and backends)](./examples/config-multiple.json) using a single zot server
 * Pull and synchronize from other dist-spec conformant registries [sync](#sync)
+* Supports ratelimiting including per HTTP method
 * Swagger based documentation
 * Single binary for _all_ the above features
 * Released under Apache 2.0 License

examples/config-ratelimit.json

@@ -0,0 +1,24 @@
+{
+    "version": "0.1.0-dev",
+    "storage": {
+        "rootDirectory": "/tmp/zot"
+    },
+    "http": {
+        "address": "127.0.0.1",
+        "port": "8080",
+        "ReadOnly": false,
+        "Ratelimit": {
+            "Rate": 10,
+            "Methods": [
+                {
+                    "Method": "GET",
+                    "Rate": 5
+                }
+            ]
+        }
+    },
+    "log": {
+        "level": "debug",
+        "output": "/tmp/zot.log"
+    }
+}

go.mod

@@ -16,6 +16,7 @@ require (
 	github.com/containerd/containerd v1.5.9 // indirect
 	github.com/containers/common v0.44.3
 	github.com/containers/image/v5 v5.16.1
+	github.com/didip/tollbooth/v6 v6.1.1
 	github.com/docker/distribution v2.7.1+incompatible
 	github.com/dustin/go-humanize v1.0.0
 	github.com/fsnotify/fsnotify v1.5.1
@@ -167,6 +168,7 @@ require (
 	github.com/go-openapi/swag v0.19.15 // indirect
 	github.com/go-openapi/validate v0.20.3 // indirect
 	github.com/go-piv/piv-go v1.9.0 // indirect
+	github.com/go-pkgz/expirable-cache v0.0.3 // indirect
 	github.com/go-playground/locales v0.13.0 // indirect
 	github.com/go-playground/universal-translator v0.17.0 // indirect
 	github.com/go-playground/validator v9.31.0+incompatible // indirect

go.sum

@@ -708,6 +708,8 @@ github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cu
 github.com/dgryski/go-sip13 v0.0.0-20181026042036-e10d5fee7954/go.mod h1:vAd38F8PWV+bWy6jNmig1y/TA+kYO4g3RSRF0IAv0no=
 github.com/dgryski/trifles v0.0.0-20200323201526-dd97f9abfb48 h1:fRzb/w+pyskVMQ+UbP35JkH8yB7MYb4q/qhBarqZE6g=
 github.com/dgryski/trifles v0.0.0-20200323201526-dd97f9abfb48/go.mod h1:if7Fbed8SFyPtHLHbg49SI7NAdJiC5WIA09pe59rfAA=
+github.com/didip/tollbooth/v6 v6.1.1 h1:Nt7PvWLa9Y94OrykXsFNBinVRQIu8xdy4avpl99Dc1M=
+github.com/didip/tollbooth/v6 v6.1.1/go.mod h1:xjcse6CTHCLuOkzsWrEgdy9WPJFv+p/x6v+MyfP+O9s=
 github.com/dimchansky/utfbom v1.1.0/go.mod h1:rO41eb7gLfo8SF1jd9F8HplJm1Fewwi4mQvIirEdv+8=
 github.com/dimchansky/utfbom v1.1.1 h1:vV6w1AhK4VMnhBno/TPVCoK9U/LP0PkLCS9tbxHdi/U=
 github.com/dimchansky/utfbom v1.1.1/go.mod h1:SxdoEBH5qIqFocHMyGOXVAybYJdr71b1Q/j0mACtrfE=
@@ -1001,6 +1003,8 @@ github.com/go-openapi/validate v0.20.3 h1:GZPPhhKSZrE8HjB4eEkoYAZmoWA4+tCemSgINH
 github.com/go-openapi/validate v0.20.3/go.mod h1:goDdqVGiigM3jChcrYJxD2joalke3ZXeftD16byIjA4=
 github.com/go-piv/piv-go v1.9.0 h1:P6j2gjfP7zO7T3nCk/jwCgsvFRwB8shEqAJ4q85jgXc=
 github.com/go-piv/piv-go v1.9.0/go.mod h1:NZ2zmjVkfFaL/CF8cVQ/pXdXtuj110zEKGdJM6fJZZM=
+github.com/go-pkgz/expirable-cache v0.0.3 h1:rTh6qNPp78z0bQE6HDhXBHUwqnV9i09Vm6dksJLXQDc=
+github.com/go-pkgz/expirable-cache v0.0.3/go.mod h1:+IauqN00R2FqNRLCLA+X5YljQJrwB179PfiAoMPlTlQ=
 github.com/go-playground/assert/v2 v2.0.1/go.mod h1:VDjEfimB/XKnb+ZQfWdccd7VUvScMdVu0Titje2rxJ4=
 github.com/go-playground/locales v0.12.1/go.mod h1:IUMDtCfWo/w/mtMfIE/IG2K+Ey3ygWanZIBtBW0W2TM=
 github.com/go-playground/locales v0.13.0 h1:HyWk6mgj5qFqCT5fjGBuRArbVDfE4hi8+e8ceBS/t7Q=

pkg/api/config/config.go

@@ -47,6 +47,16 @@ type BearerConfig struct {
 	Cert    string
 }
 
+type MethodRatelimitConfig struct {
+	Method string
+	Rate   int
+}
+
+type RatelimitConfig struct {
+	Rate    *int                    // requests per second
+	Methods []MethodRatelimitConfig `mapstructure:",omitempty"`
+}
+
 type HTTPConfig struct {
 	Address          string
 	Port             string
@@ -56,6 +66,7 @@ type HTTPConfig struct {
 	Realm            string
 	AllowReadAccess  bool             `mapstructure:",omitempty"`
 	ReadOnly         bool             `mapstructure:",omitempty"`
+	Ratelimit        *RatelimitConfig `mapstructure:",omitempty"`
 }
 
 type LDAPConfig struct {

pkg/api/controller.go

@@ -8,7 +8,10 @@ import (
 	"io/ioutil"
 	"net"
 	"net/http"
+	"runtime"
+	"strings"
 	goSync "sync"
+	"syscall"
 	"time"
 
 	"github.com/docker/distribution/registry/storage/driver/factory"
@@ -68,6 +71,27 @@ func DefaultHeaders() mux.MiddlewareFunc {
 	}
 }
 
+func DumpRuntimeParams(log log.Logger) {
+	var rLimit syscall.Rlimit
+
+	evt := log.Info().Int("cpus", runtime.NumCPU())
+
+	err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, &rLimit)
+	if err == nil {
+		evt = evt.Uint64("max. open files", rLimit.Cur)
+	}
+
+	if content, err := ioutil.ReadFile("/proc/sys/net/core/somaxconn"); err == nil {
+		evt = evt.Str("listen backlog", strings.TrimSuffix(string(content), "\n"))
+	}
+
+	if content, err := ioutil.ReadFile("/proc/sys/user/max_inotify_watches"); err == nil {
+		evt = evt.Str("max. inotify watches", strings.TrimSuffix(string(content), "\n"))
+	}
+
+	evt.Msg("runtime params")
+}
+
 func (c *Controller) Run() error {
 	// validate configuration
 	if err := c.Config.Validate(c.Log); err != nil {
@@ -79,8 +103,25 @@ func (c *Controller) Run() error {
 	// print the current configuration, but strip secrets
 	c.Log.Info().Interface("params", c.Config.Sanitize()).Msg("configuration settings")
 
+	// print the current runtime environment
+	DumpRuntimeParams(c.Log)
+
+	// setup HTTP API router
 	engine := mux.NewRouter()
-	engine.Use(DefaultHeaders(),
+
+	// rate-limit HTTP requests if enabled
+	if c.Config.HTTP.Ratelimit != nil {
+		if c.Config.HTTP.Ratelimit.Rate != nil {
+			engine.Use(RateLimiter(c, *c.Config.HTTP.Ratelimit.Rate))
+		}
+
+		for _, mrlim := range c.Config.HTTP.Ratelimit.Methods {
+			engine.Use(MethodRateLimiter(c, mrlim.Method, mrlim.Rate))
+		}
+	}
+
+	engine.Use(
+		DefaultHeaders(),
 		SessionLogger(c),
 		handlers.RecoveryHandler(handlers.RecoveryLogger(c.Log),
 			handlers.PrintRecoveryStack(false)))

pkg/api/controller_test.go

@@ -388,6 +388,127 @@ func TestHtpasswdFiveCreds(t *testing.T) {
 	})
 }
 
+func TestRatelimit(t *testing.T) {
+	Convey("Make a new controller", t, func() {
+		port := GetFreePort()
+		baseURL := GetBaseURL(port)
+		conf := config.New()
+		conf.HTTP.Port = port
+
+		rate := 1
+		conf.HTTP.Ratelimit = &config.RatelimitConfig{
+			Rate: &rate,
+		}
+		ctlr := api.NewController(conf)
+		dir, err := ioutil.TempDir("", "oci-repo-test")
+		if err != nil {
+			panic(err)
+		}
+		defer os.RemoveAll(dir)
+		ctlr.Config.Storage.RootDirectory = dir
+
+		go startServer(ctlr)
+		defer stopServer(ctlr)
+		WaitTillServerReady(baseURL)
+
+		Convey("Ratelimit", func() {
+			client := resty.New()
+			// first request should succeed
+			resp, err := client.R().Get(baseURL + "/v2/")
+			So(err, ShouldBeNil)
+			So(resp, ShouldNotBeNil)
+			So(resp.StatusCode(), ShouldEqual, http.StatusOK)
+			// second request back-to-back should fail
+			resp, err = client.R().Get(baseURL + "/v2/")
+			So(err, ShouldBeNil)
+			So(resp, ShouldNotBeNil)
+			So(resp.StatusCode(), ShouldEqual, http.StatusTooManyRequests)
+		})
+	})
+
+	Convey("Make a new controller", t, func() {
+		port := GetFreePort()
+		baseURL := GetBaseURL(port)
+		conf := config.New()
+		conf.HTTP.Port = port
+
+		conf.HTTP.Ratelimit = &config.RatelimitConfig{
+			Methods: []config.MethodRatelimitConfig{
+				{
+					Method: http.MethodGet,
+					Rate:   1,
+				},
+			},
+		}
+		ctlr := api.NewController(conf)
+		dir, err := ioutil.TempDir("", "oci-repo-test")
+		if err != nil {
+			panic(err)
+		}
+		defer os.RemoveAll(dir)
+		ctlr.Config.Storage.RootDirectory = dir
+
+		go startServer(ctlr)
+		defer stopServer(ctlr)
+		WaitTillServerReady(baseURL)
+		Convey("Method Ratelimit", func() {
+			client := resty.New()
+			// first request should succeed
+			resp, err := client.R().Get(baseURL + "/v2/")
+			So(err, ShouldBeNil)
+			So(resp, ShouldNotBeNil)
+			So(resp.StatusCode(), ShouldEqual, http.StatusOK)
+			// second request back-to-back should fail
+			resp, err = client.R().Get(baseURL + "/v2/")
+			So(err, ShouldBeNil)
+			So(resp, ShouldNotBeNil)
+			So(resp.StatusCode(), ShouldEqual, http.StatusTooManyRequests)
+		})
+	})
+
+	Convey("Make a new controller", t, func() {
+		port := GetFreePort()
+		baseURL := GetBaseURL(port)
+		conf := config.New()
+		conf.HTTP.Port = port
+
+		rate := 1
+		conf.HTTP.Ratelimit = &config.RatelimitConfig{
+			Rate: &rate, // this dominates
+			Methods: []config.MethodRatelimitConfig{
+				{
+					Method: http.MethodGet,
+					Rate:   100,
+				},
+			},
+		}
+		ctlr := api.NewController(conf)
+		dir, err := ioutil.TempDir("", "oci-repo-test")
+		if err != nil {
+			panic(err)
+		}
+		defer os.RemoveAll(dir)
+		ctlr.Config.Storage.RootDirectory = dir
+
+		go startServer(ctlr)
+		defer stopServer(ctlr)
+		WaitTillServerReady(baseURL)
+		Convey("Global and Method Ratelimit", func() {
+			client := resty.New()
+			// first request should succeed
+			resp, err := client.R().Get(baseURL + "/v2/")
+			So(err, ShouldBeNil)
+			So(resp, ShouldNotBeNil)
+			So(resp.StatusCode(), ShouldEqual, http.StatusOK)
+			// second request back-to-back should fail
+			resp, err = client.R().Get(baseURL + "/v2/")
+			So(err, ShouldBeNil)
+			So(resp, ShouldNotBeNil)
+			So(resp.StatusCode(), ShouldEqual, http.StatusTooManyRequests)
+		})
+	})
+}
+
 func TestBasicAuth(t *testing.T) {
 	Convey("Make a new controller", t, func() {
 		port := GetFreePort()

pkg/api/session.go

@@ -7,6 +7,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/didip/tollbooth/v6"
 	"github.com/gorilla/mux"
 	"zotregistry.io/zot/pkg/extensions/monitoring"
 	"zotregistry.io/zot/pkg/log"
@@ -25,7 +26,7 @@ func (w *statusWriter) WriteHeader(status int) {
 
 func (w *statusWriter) Write(b []byte) (int, error) {
 	if w.status == 0 {
-		w.status = 200
+		w.status = http.StatusOK
 	}
 
 	n, err := w.ResponseWriter.Write(b)
@@ -34,6 +35,35 @@ func (w *statusWriter) Write(b []byte) (int, error) {
 	return n, err
 }
 
+// RateLimiter limits handling of incoming requests.
+func RateLimiter(ctlr *Controller, rate int) mux.MiddlewareFunc {
+	ctlr.Log.Info().Int("rate", rate).Msg("ratelimiter enabled")
+
+	limiter := tollbooth.NewLimiter(float64(rate), nil)
+	limiter.SetMessage(http.StatusText(http.StatusTooManyRequests)).
+		SetStatusCode(http.StatusTooManyRequests).
+		SetOnLimitReached(nil)
+
+	return func(next http.Handler) http.Handler {
+		return tollbooth.LimitHandler(limiter, next)
+	}
+}
+
+// MethodRateLimiter limits handling of incoming requests.
+func MethodRateLimiter(ctlr *Controller, method string, rate int) mux.MiddlewareFunc {
+	ctlr.Log.Info().Str("method", method).Int("rate", rate).Msg("per-method ratelimiter enabled")
+
+	limiter := tollbooth.NewLimiter(float64(rate), nil)
+	limiter.SetMethods([]string{method}).
+		SetMessage(http.StatusText(http.StatusTooManyRequests)).
+		SetStatusCode(http.StatusTooManyRequests).
+		SetOnLimitReached(nil)
+
+	return func(next http.Handler) http.Handler {
+		return tollbooth.LimitHandler(limiter, next)
+	}
+}
+
 // SessionLogger logs session details.
 func SessionLogger(ctlr *Controller) mux.MiddlewareFunc {
 	logger := ctlr.Log.With().Str("module", "http").Logger()