Rate limiting (#1385)

* Initial rate limiting

* Move rate limiting to client API

* Update rate limits to hopefully be self-cleaning

* Use X-Forwarded-For, add comments

* Reduce rate limit threshold

* Tweak interval

* Configurable backoff

* Review comments, set cleanup interval to 30 seconds

* Allow generate-config to produce sane CI config

* Fix Complement dockerfile

6 files changed, 200 insertions, 2 deletions

diff --git a/build/scripts/Complement.Dockerfile b/build/scripts/Complement.Dockerfile
index 32c5234b..de51f16d 100644
--- a/build/scripts/Complement.Dockerfile
+++ b/build/scripts/Complement.Dockerfile
@@ -12,8 +12,7 @@ COPY . .
 RUN go build ./cmd/dendrite-monolith-server
 RUN go build ./cmd/generate-keys
 RUN go build ./cmd/generate-config
-RUN ./generate-config > dendrite.yaml
-RUN sed -i "s/disable_tls_validation: false/disable_tls_validation: true/g" dendrite.yaml
+RUN ./generate-config --ci > dendrite.yaml
 RUN ./generate-keys --private-key matrix_key.pem --tls-cert server.crt --tls-key server.key
 
 ENV SERVER_NAME=localhost
diff --git a/clientapi/routing/rate_limiting.go b/clientapi/routing/rate_limiting.go
new file mode 100644
index 00000000..16e3c056
--- /dev/null
+++ b/clientapi/routing/rate_limiting.go
@@ -0,0 +1,99 @@
+package routing
+
+import (
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/matrix-org/dendrite/clientapi/jsonerror"
+	"github.com/matrix-org/dendrite/internal/config"
+	"github.com/matrix-org/util"
+)
+
+type rateLimits struct {
+	limits           map[string]chan struct{}
+	limitsMutex      sync.RWMutex
+	enabled          bool
+	requestThreshold int64
+	cooloffDuration  time.Duration
+}
+
+func newRateLimits(cfg *config.RateLimiting) *rateLimits {
+	l := &rateLimits{
+		limits:           make(map[string]chan struct{}),
+		enabled:          cfg.Enabled,
+		requestThreshold: cfg.Threshold,
+		cooloffDuration:  time.Duration(cfg.CooloffMS) * time.Millisecond,
+	}
+	if l.enabled {
+		go l.clean()
+	}
+	return l
+}
+
+func (l *rateLimits) clean() {
+	for {
+		// On a 30 second interval, we'll take an exclusive write
+		// lock of the entire map and see if any of the channels are
+		// empty. If they are then we will close and delete them,
+		// freeing up memory.
+		time.Sleep(time.Second * 30)
+		l.limitsMutex.Lock()
+		for k, c := range l.limits {
+			if len(c) == 0 {
+				close(c)
+				delete(l.limits, k)
+			}
+		}
+		l.limitsMutex.Unlock()
+	}
+}
+
+func (l *rateLimits) rateLimit(req *http.Request) *util.JSONResponse {
+	// If rate limiting is disabled then do nothing.
+	if !l.enabled {
+		return nil
+	}
+
+	// Lock the map long enough to check for rate limiting. We hold it
+	// for longer here than we really need to but it makes sure that we
+	// also don't conflict with the cleaner goroutine which might clean
+	// up a channel after we have retrieved it otherwise.
+	l.limitsMutex.RLock()
+	defer l.limitsMutex.RUnlock()
+
+	// First of all, work out if X-Forwarded-For was sent to us. If not
+	// then we'll just use the IP address of the caller.
+	caller := req.RemoteAddr
+	if forwardedFor := req.Header.Get("X-Forwarded-For"); forwardedFor != "" {
+		caller = forwardedFor
+	}
+
+	// Look up the caller's channel, if they have one. If they don't then
+	// let's create one.
+	rateLimit, ok := l.limits[caller]
+	if !ok {
+		l.limits[caller] = make(chan struct{}, l.requestThreshold)
+		rateLimit = l.limits[caller]
+	}
+
+	// Check if the user has got free resource slots for this request.
+	// If they don't then we'll return an error.
+	select {
+	case rateLimit <- struct{}{}:
+	default:
+		// We hit the rate limit. Tell the client to back off.
+		return &util.JSONResponse{
+			Code: http.StatusTooManyRequests,
+			JSON: jsonerror.LimitExceeded("You are sending too many requests too quickly!", l.cooloffDuration.Milliseconds()),
+		}
+	}
+
+	// After the time interval, drain a resource from the rate limiting
+	// channel. This will free up space in the channel for new requests.
+	go func() {
+		<-time.After(l.cooloffDuration)
+		<-rateLimit
+	}()
+	return nil
+}
diff --git a/clientapi/routing/routing.go b/clientapi/routing/routing.go
index 24343ee1..0c63f968 100644
--- a/clientapi/routing/routing.go
+++ b/clientapi/routing/routing.go
@@ -60,6 +60,7 @@ func Setup(
 	keyAPI keyserverAPI.KeyInternalAPI,
 	extRoomsProvider api.ExtraPublicRoomsProvider,
 ) {
+	rateLimits := newRateLimits(&cfg.RateLimiting)
 	userInteractiveAuth := auth.NewUserInteractive(accountDB.GetAccountByPassword, cfg)
 
 	publicAPIMux.Handle("/versions",
@@ -92,6 +93,9 @@ func Setup(
 	).Methods(http.MethodPost, http.MethodOptions)
 	r0mux.Handle("/join/{roomIDOrAlias}",
 		httputil.MakeAuthAPI(gomatrixserverlib.Join, userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
+			if r := rateLimits.rateLimit(req); r != nil {
+				return *r
+			}
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
 			if err != nil {
 				return util.ErrorResponse(err)
@@ -108,6 +112,9 @@ func Setup(
 	).Methods(http.MethodGet, http.MethodOptions)
 	r0mux.Handle("/rooms/{roomID}/join",
 		httputil.MakeAuthAPI(gomatrixserverlib.Join, userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
+			if r := rateLimits.rateLimit(req); r != nil {
+				return *r
+			}
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
 			if err != nil {
 				return util.ErrorResponse(err)
@@ -119,6 +126,9 @@ func Setup(
 	).Methods(http.MethodPost, http.MethodOptions)
 	r0mux.Handle("/rooms/{roomID}/leave",
 		httputil.MakeAuthAPI("membership", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
+			if r := rateLimits.rateLimit(req); r != nil {
+				return *r
+			}
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
 			if err != nil {
 				return util.ErrorResponse(err)
@@ -139,6 +149,9 @@ func Setup(
 	).Methods(http.MethodPost, http.MethodOptions)
 	r0mux.Handle("/rooms/{roomID}/invite",
 		httputil.MakeAuthAPI("membership", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
+			if r := rateLimits.rateLimit(req); r != nil {
+				return *r
+			}
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
 			if err != nil {
 				return util.ErrorResponse(err)
@@ -253,14 +266,23 @@ func Setup(
 	).Methods(http.MethodPut, http.MethodOptions)
 
 	r0mux.Handle("/register", httputil.MakeExternalAPI("register", func(req *http.Request) util.JSONResponse {
+		if r := rateLimits.rateLimit(req); r != nil {
+			return *r
+		}
 		return Register(req, userAPI, accountDB, cfg)
 	})).Methods(http.MethodPost, http.MethodOptions)
 
 	v1mux.Handle("/register", httputil.MakeExternalAPI("register", func(req *http.Request) util.JSONResponse {
+		if r := rateLimits.rateLimit(req); r != nil {
+			return *r
+		}
 		return LegacyRegister(req, userAPI, cfg)
 	})).Methods(http.MethodPost, http.MethodOptions)
 
 	r0mux.Handle("/register/available", httputil.MakeExternalAPI("registerAvailable", func(req *http.Request) util.JSONResponse {
+		if r := rateLimits.rateLimit(req); r != nil {
+			return *r
+		}
 		return RegisterAvailable(req, cfg, accountDB)
 	})).Methods(http.MethodGet, http.MethodOptions)
 
@@ -332,6 +354,9 @@ func Setup(
 
 	r0mux.Handle("/rooms/{roomID}/typing/{userID}",
 		httputil.MakeAuthAPI("rooms_typing", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
+			if r := rateLimits.rateLimit(req); r != nil {
+				return *r
+			}
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
 			if err != nil {
 				return util.ErrorResponse(err)
@@ -385,6 +410,9 @@ func Setup(
 
 	r0mux.Handle("/account/whoami",
 		httputil.MakeAuthAPI("whoami", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
+			if r := rateLimits.rateLimit(req); r != nil {
+				return *r
+			}
 			return Whoami(req, device)
 		}),
 	).Methods(http.MethodGet, http.MethodOptions)
@@ -393,6 +421,9 @@ func Setup(
 
 	r0mux.Handle("/login",
 		httputil.MakeExternalAPI("login", func(req *http.Request) util.JSONResponse {
+			if r := rateLimits.rateLimit(req); r != nil {
+				return *r
+			}
 			return Login(req, accountDB, userAPI, cfg)
 		}),
 	).Methods(http.MethodGet, http.MethodPost, http.MethodOptions)
@@ -447,6 +478,9 @@ func Setup(
 
 	r0mux.Handle("/profile/{userID}/avatar_url",
 		httputil.MakeAuthAPI("profile_avatar_url", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
+			if r := rateLimits.rateLimit(req); r != nil {
+				return *r
+			}
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
 			if err != nil {
 				return util.ErrorResponse(err)
@@ -469,6 +503,9 @@ func Setup(
 
 	r0mux.Handle("/profile/{userID}/displayname",
 		httputil.MakeAuthAPI("profile_displayname", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
+			if r := rateLimits.rateLimit(req); r != nil {
+				return *r
+			}
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
 			if err != nil {
 				return util.ErrorResponse(err)
@@ -506,6 +543,9 @@ func Setup(
 	// Riot logs get flooded unless this is handled
 	r0mux.Handle("/presence/{userID}/status",
 		httputil.MakeExternalAPI("presence", func(req *http.Request) util.JSONResponse {
+			if r := rateLimits.rateLimit(req); r != nil {
+				return *r
+			}
 			// TODO: Set presence (probably the responsibility of a presence server not clientapi)
 			return util.JSONResponse{
 				Code: http.StatusOK,
@@ -516,6 +556,9 @@ func Setup(
 
 	r0mux.Handle("/voip/turnServer",
 		httputil.MakeAuthAPI("turn_server", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
+			if r := rateLimits.rateLimit(req); r != nil {
+				return *r
+			}
 			return RequestTurnServer(req, device, cfg)
 		}),
 	).Methods(http.MethodGet, http.MethodOptions)
@@ -582,6 +625,9 @@ func Setup(
 
 	r0mux.Handle("/user_directory/search",
 		httputil.MakeAuthAPI("userdirectory_search", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
+			if r := rateLimits.rateLimit(req); r != nil {
+				return *r
+			}
 			postContent := struct {
 				SearchString string `json:"search_term"`
 				Limit        int    `json:"limit"`
@@ -623,6 +669,9 @@ func Setup(
 
 	r0mux.Handle("/rooms/{roomID}/read_markers",
 		httputil.MakeExternalAPI("rooms_read_markers", func(req *http.Request) util.JSONResponse {
+			if r := rateLimits.rateLimit(req); r != nil {
+				return *r
+			}
 			// TODO: return the read_markers.
 			return util.JSONResponse{Code: http.StatusOK, JSON: struct{}{}}
 		}),
@@ -721,6 +770,9 @@ func Setup(
 
 	r0mux.Handle("/capabilities",
 		httputil.MakeAuthAPI("capabilities", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
+			if r := rateLimits.rateLimit(req); r != nil {
+				return *r
+			}
 			return GetCapabilities(req, rsAPI)
 		}),
 	).Methods(http.MethodGet)
diff --git a/cmd/generate-config/main.go b/cmd/generate-config/main.go
index cff376d8..78ed3af6 100644
--- a/cmd/generate-config/main.go
+++ b/cmd/generate-config/main.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"flag"
 	"fmt"
 
 	"github.com/matrix-org/dendrite/internal/config"
@@ -8,6 +9,9 @@ import (
 )
 
 func main() {
+	defaultsForCI := flag.Bool("ci", false, "sane defaults for CI testing")
+	flag.Parse()
+
 	cfg := &config.Dendrite{}
 	cfg.Defaults()
 	cfg.Global.TrustedIDServers = []string{
@@ -56,6 +60,11 @@ func main() {
 		},
 	}
 
+	if *defaultsForCI {
+		cfg.ClientAPI.RateLimiting.Enabled = false
+		cfg.FederationSender.DisableTLSValidation = true
+	}
+
 	j, err := yaml.Marshal(cfg)
 	if err != nil {
 		panic(err)
diff --git a/dendrite-config.yaml b/dendrite-config.yaml
index 23f142a8..570669c1 100644
--- a/dendrite-config.yaml
+++ b/dendrite-config.yaml
@@ -133,6 +133,14 @@ client_api:
     turn_username: ""
     turn_password: ""
 
+  # Settings for rate-limited endpoints. Rate limiting will kick in after the
+  # threshold number of "slots" have been taken by requests from a specific 
+  # host. Each "slot" will be released after the cooloff time in milliseconds.
+  rate_limiting:
+    enabled: true
+    threshold: 5
+    cooloff_ms: 500
+
 # Configuration for the Current State Server.
 current_state_server:
   internal_api:
diff --git a/internal/config/config_clientapi.go b/internal/config/config_clientapi.go
index f7878276..52115491 100644
--- a/internal/config/config_clientapi.go
+++ b/internal/config/config_clientapi.go
@@ -34,6 +34,9 @@ type ClientAPI struct {
 
 	// TURN options
 	TURN TURN `yaml:"turn"`
+
+	// Rate-limiting options
+	RateLimiting RateLimiting `yaml:"rate_limiting"`
 }
 
 func (c *ClientAPI) Defaults() {
@@ -47,6 +50,7 @@ func (c *ClientAPI) Defaults() {
 	c.RecaptchaBypassSecret = ""
 	c.RecaptchaSiteVerifyAPI = ""
 	c.RegistrationDisabled = false
+	c.RateLimiting.Defaults()
 }
 
 func (c *ClientAPI) Verify(configErrs *ConfigErrors, isMonolith bool) {
@@ -61,6 +65,7 @@ func (c *ClientAPI) Verify(configErrs *ConfigErrors, isMonolith bool) {
 		checkNotEmpty(configErrs, "client_api.recaptcha_siteverify_api", string(c.RecaptchaSiteVerifyAPI))
 	}
 	c.TURN.Verify(configErrs)
+	c.RateLimiting.Verify(configErrs)
 }
 
 type TURN struct {
@@ -90,3 +95,29 @@ func (c *TURN) Verify(configErrs *ConfigErrors) {
 		}
 	}
 }
+
+type RateLimiting struct {
+	// Is rate limiting enabled or disabled?
+	Enabled bool `yaml:"enabled"`
+
+	// How many "slots" a user can occupy sending requests to a rate-limited
+	// endpoint before we apply rate-limiting
+	Threshold int64 `yaml:"threshold"`
+
+	// The cooloff period in milliseconds after a request before the "slot"
+	// is freed again
+	CooloffMS int64 `yaml:"cooloff_ms"`
+}
+
+func (r *RateLimiting) Verify(configErrs *ConfigErrors) {
+	if r.Enabled {
+		checkPositive(configErrs, "client_api.rate_limiting.threshold", r.Threshold)
+		checkPositive(configErrs, "client_api.rate_limiting.cooloff_ms", r.CooloffMS)
+	}
+}
+
+func (r *RateLimiting) Defaults() {
+	r.Enabled = true
+	r.Threshold = 5
+	r.CooloffMS = 500
+}