Rate limiting changes (#2519)

* Rate limiting changes

This makes the following changes:

* For logged in users, the rate limiting now applies to the device session rather than the remote IP address;
* For non-logged in users, the rate limiting continues to apply to remote address as it does today;
* It is now possible to add user IDs to the `exempt_user_ids` option under `rate_limiting` to exclude bots from rate limiting;
* Admin and appservice users are now exempt from rate limiting by default.

* Fix build with media API

1 files changed, 27 insertions, 4 deletions

diff --git a/internal/httputil/rate_limiting.go b/internal/httputil/rate_limiting.go
index c4f47c7b..dab36481 100644
--- a/internal/httputil/rate_limiting.go
+++ b/internal/httputil/rate_limiting.go
@@ -7,6 +7,7 @@ import (
 
 	"github.com/matrix-org/dendrite/clientapi/jsonerror"
 	"github.com/matrix-org/dendrite/setup/config"
+	userapi "github.com/matrix-org/dendrite/userapi/api"
 	"github.com/matrix-org/util"
 )
 
@@ -17,6 +18,7 @@ type RateLimits struct {
 	enabled          bool
 	requestThreshold int64
 	cooloffDuration  time.Duration
+	exemptUserIDs    map[string]struct{}
 }
 
 func NewRateLimits(cfg *config.RateLimiting) *RateLimits {
@@ -25,6 +27,10 @@ func NewRateLimits(cfg *config.RateLimiting) *RateLimits {
 		enabled:          cfg.Enabled,
 		requestThreshold: cfg.Threshold,
 		cooloffDuration:  time.Duration(cfg.CooloffMS) * time.Millisecond,
+		exemptUserIDs:    map[string]struct{}{},
+	}
+	for _, userID := range cfg.ExemptUserIDs {
+		l.exemptUserIDs[userID] = struct{}{}
 	}
 	if l.enabled {
 		go l.clean()
@@ -52,7 +58,7 @@ func (l *RateLimits) clean() {
 	}
 }
 
-func (l *RateLimits) Limit(req *http.Request) *util.JSONResponse {
+func (l *RateLimits) Limit(req *http.Request, device *userapi.Device) *util.JSONResponse {
 	// If rate limiting is disabled then do nothing.
 	if !l.enabled {
 		return nil
@@ -67,9 +73,26 @@ func (l *RateLimits) Limit(req *http.Request) *util.JSONResponse {
 
 	// 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
+	var caller string
+	if device != nil {
+		switch device.AccountType {
+		case userapi.AccountTypeAdmin:
+			return nil // don't rate-limit server administrators
+		case userapi.AccountTypeAppService:
+			return nil // don't rate-limit appservice users
+		default:
+			if _, ok := l.exemptUserIDs[device.UserID]; ok {
+				// If the user is exempt from rate limiting then do nothing.
+				return nil
+			}
+			caller = device.UserID + device.ID
+		}
+	} else {
+		if forwardedFor := req.Header.Get("X-Forwarded-For"); forwardedFor != "" {
+			caller = forwardedFor
+		} else {
+			caller = req.RemoteAddr
+		}
 	}
 
 	// Look up the caller's channel, if they have one.