Ratelimit requests to /media/r0/download|upload (#2020)

* Add /media/r0/config handler

Signed-off-by: Till Faelligen <tfaelligen@gmail.com>

* Add rate limiting to media api

* Rename variable

* Add passing tests

* Don't send multiple headers

Co-authored-by: Neil Alexander <neilalexander@users.noreply.github.com>

2 files changed, 23 insertions, 132 deletions

diff --git a/clientapi/routing/rate_limiting.go b/clientapi/routing/rate_limiting.go
deleted file mode 100644
index 5291caba..00000000
--- a/clientapi/routing/rate_limiting.go
+++ /dev/null
@@ -1,109 +0,0 @@
-package routing
-
-import (
-	"net/http"
-	"sync"
-	"time"
-
-	"github.com/matrix-org/dendrite/clientapi/jsonerror"
-	"github.com/matrix-org/dendrite/setup/config"
-	"github.com/matrix-org/util"
-)
-
-type rateLimits struct {
-	limits           map[string]chan struct{}
-	limitsMutex      sync.RWMutex
-	cleanMutex       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.cleanMutex.Lock()
-		l.limitsMutex.Lock()
-		for k, c := range l.limits {
-			if len(c) == 0 {
-				close(c)
-				delete(l.limits, k)
-			}
-		}
-		l.limitsMutex.Unlock()
-		l.cleanMutex.Unlock()
-	}
-}
-
-func (l *rateLimits) rateLimit(req *http.Request) *util.JSONResponse {
-	// If rate limiting is disabled then do nothing.
-	if !l.enabled {
-		return nil
-	}
-
-	// Take a read lock out on the cleaner mutex. The cleaner expects to
-	// be able to take a write lock, which isn't possible while there are
-	// readers, so this has the effect of blocking the cleaner goroutine
-	// from doing its work until there are no requests in flight.
-	l.cleanMutex.RLock()
-	defer l.cleanMutex.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.
-	l.limitsMutex.RLock()
-	rateLimit, ok := l.limits[caller]
-	l.limitsMutex.RUnlock()
-
-	// If the caller doesn't have a channel, create one and write it
-	// back to the map.
-	if !ok {
-		rateLimit = make(chan struct{}, l.requestThreshold)
-
-		l.limitsMutex.Lock()
-		l.limits[caller] = rateLimit
-		l.limitsMutex.Unlock()
-	}
-
-	// 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 813d9d16..9263c66b 100644
--- a/clientapi/routing/routing.go
+++ b/clientapi/routing/routing.go
@@ -61,7 +61,7 @@ func Setup(
 	extRoomsProvider api.ExtraPublicRoomsProvider,
 	mscCfg *config.MSCs,
 ) {
-	rateLimits := newRateLimits(&cfg.RateLimiting)
+	rateLimits := httputil.NewRateLimits(&cfg.RateLimiting)
 	userInteractiveAuth := auth.NewUserInteractive(accountDB.GetAccountByPassword, cfg)
 
 	unstableFeatures := map[string]bool{
@@ -127,7 +127,7 @@ 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 {
+			if r := rateLimits.Limit(req); r != nil {
 				return *r
 			}
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
@@ -143,7 +143,7 @@ func Setup(
 	if mscCfg.Enabled("msc2753") {
 		r0mux.Handle("/peek/{roomIDOrAlias}",
 			httputil.MakeAuthAPI(gomatrixserverlib.Peek, userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
-				if r := rateLimits.rateLimit(req); r != nil {
+				if r := rateLimits.Limit(req); r != nil {
 					return *r
 				}
 				vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
@@ -163,7 +163,7 @@ 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 {
+			if r := rateLimits.Limit(req); r != nil {
 				return *r
 			}
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
@@ -177,7 +177,7 @@ 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 {
+			if r := rateLimits.Limit(req); r != nil {
 				return *r
 			}
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
@@ -211,7 +211,7 @@ 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 {
+			if r := rateLimits.Limit(req); r != nil {
 				return *r
 			}
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
@@ -329,14 +329,14 @@ 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 {
+		if r := rateLimits.Limit(req); r != nil {
 			return *r
 		}
 		return Register(req, userAPI, accountDB, 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 {
+		if r := rateLimits.Limit(req); r != nil {
 			return *r
 		}
 		return RegisterAvailable(req, cfg, accountDB)
@@ -410,7 +410,7 @@ 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 {
+			if r := rateLimits.Limit(req); r != nil {
 				return *r
 			}
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
@@ -466,7 +466,7 @@ 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 {
+			if r := rateLimits.Limit(req); r != nil {
 				return *r
 			}
 			return Whoami(req, device)
@@ -475,7 +475,7 @@ func Setup(
 
 	r0mux.Handle("/account/password",
 		httputil.MakeAuthAPI("password", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
-			if r := rateLimits.rateLimit(req); r != nil {
+			if r := rateLimits.Limit(req); r != nil {
 				return *r
 			}
 			return Password(req, userAPI, accountDB, device, cfg)
@@ -484,7 +484,7 @@ func Setup(
 
 	r0mux.Handle("/account/deactivate",
 		httputil.MakeAuthAPI("deactivate", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
-			if r := rateLimits.rateLimit(req); r != nil {
+			if r := rateLimits.Limit(req); r != nil {
 				return *r
 			}
 			return Deactivate(req, userInteractiveAuth, userAPI, device)
@@ -495,7 +495,7 @@ func Setup(
 
 	r0mux.Handle("/login",
 		httputil.MakeExternalAPI("login", func(req *http.Request) util.JSONResponse {
-			if r := rateLimits.rateLimit(req); r != nil {
+			if r := rateLimits.Limit(req); r != nil {
 				return *r
 			}
 			return Login(req, accountDB, userAPI, cfg)
@@ -552,7 +552,7 @@ 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 {
+			if r := rateLimits.Limit(req); r != nil {
 				return *r
 			}
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
@@ -577,7 +577,7 @@ 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 {
+			if r := rateLimits.Limit(req); r != nil {
 				return *r
 			}
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
@@ -617,7 +617,7 @@ func Setup(
 	// Element 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 {
+			if r := rateLimits.Limit(req); r != nil {
 				return *r
 			}
 			// TODO: Set presence (probably the responsibility of a presence server not clientapi)
@@ -630,7 +630,7 @@ 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 {
+			if r := rateLimits.Limit(req); r != nil {
 				return *r
 			}
 			return RequestTurnServer(req, device, cfg)
@@ -709,7 +709,7 @@ func Setup(
 
 	r0mux.Handle("/user/{userID}/openid/request_token",
 		httputil.MakeAuthAPI("openid_request_token", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
-			if r := rateLimits.rateLimit(req); r != nil {
+			if r := rateLimits.Limit(req); r != nil {
 				return *r
 			}
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
@@ -722,7 +722,7 @@ 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 {
+			if r := rateLimits.Limit(req); r != nil {
 				return *r
 			}
 			postContent := struct {
@@ -767,7 +767,7 @@ func Setup(
 
 	r0mux.Handle("/rooms/{roomID}/read_markers",
 		httputil.MakeAuthAPI("rooms_read_markers", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
-			if r := rateLimits.rateLimit(req); r != nil {
+			if r := rateLimits.Limit(req); r != nil {
 				return *r
 			}
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
@@ -780,7 +780,7 @@ func Setup(
 
 	r0mux.Handle("/rooms/{roomID}/forget",
 		httputil.MakeAuthAPI("rooms_forget", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
-			if r := rateLimits.rateLimit(req); r != nil {
+			if r := rateLimits.Limit(req); r != nil {
 				return *r
 			}
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
@@ -884,7 +884,7 @@ 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 {
+			if r := rateLimits.Limit(req); r != nil {
 				return *r
 			}
 			return GetCapabilities(req, rsAPI)
@@ -1100,7 +1100,7 @@ func Setup(
 	).Methods(http.MethodPost, http.MethodOptions)
 	r0mux.Handle("/rooms/{roomId}/receipt/{receiptType}/{eventId}",
 		httputil.MakeAuthAPI(gomatrixserverlib.Join, userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
-			if r := rateLimits.rateLimit(req); r != nil {
+			if r := rateLimits.Limit(req); r != nil {
 				return *r
 			}
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))