Update rate limits to hopefully be self-cleaning

clientapi/routing/rate_limiting.go

@@ -9,36 +9,63 @@ import (
 	"github.com/matrix-org/util"
 )
 
-var clientRateLimits sync.Map // device ID -> chan bool buffered
+type rateLimits struct {
-var clientRateLimitMaxRequests = 10
+	limits       map[string]chan struct{}
-var clientRateLimitTimeIntervalMS = time.Millisecond * 500
+	limitsMutex  sync.RWMutex
+	maxRequests  int
+	timeInterval time.Duration
+}
 
-func rateLimit(req *http.Request) *util.JSONResponse {
+func newRateLimits() *rateLimits {
+	l := &rateLimits{
+		limits:       make(map[string]chan struct{}),
+		maxRequests:  10,
+		timeInterval: time.Millisecond * 500,
+	}
+	go l.clean()
+	return l
+}
+
+func (l *rateLimits) clean() {
+	for {
+		time.Sleep(time.Second * 10)
+		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 {
 	// Check if the user has got free resource slots for this request.
 	// If they don't then we'll return an error.
-	rateLimit, _ := clientRateLimits.LoadOrStore(req.RemoteAddr, make(chan struct{}, clientRateLimitMaxRequests))
+	l.limitsMutex.RLock()
-	rateChan := rateLimit.(chan struct{})
+	defer l.limitsMutex.RUnlock()
+
+	rateLimit, ok := l.limits[req.RemoteAddr]
+	if !ok {
+		l.limits[req.RemoteAddr] = make(chan struct{}, l.maxRequests)
+		rateLimit = l.limits[req.RemoteAddr]
+	}
 	select {
-	case rateChan <- struct{}{}:
+	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!", clientRateLimitTimeIntervalMS.Milliseconds()),
+			JSON: jsonerror.LimitExceeded("You are sending too many requests too quickly!", l.timeInterval.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(clientRateLimitTimeIntervalMS)
+		<-time.After(l.timeInterval)
-		<-rateChan
+		<-rateLimit
-
-		// TODO: racy?
-		if len(rateChan) == 0 {
-			close(rateChan)
-			clientRateLimits.Delete(req.RemoteAddr)
-		}
 	}()
 	return nil
 }

clientapi/routing/routing.go

@@ -60,6 +60,7 @@ func Setup(
 	keyAPI keyserverAPI.KeyInternalAPI,
 	extRoomsProvider api.ExtraPublicRoomsProvider,
 ) {
+	rateLimits := newRateLimits()
 	userInteractiveAuth := auth.NewUserInteractive(accountDB.GetAccountByPassword, cfg)
 
 	publicAPIMux.Handle("/versions",
@@ -92,7 +93,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 := rateLimit(req); r != nil {
+			if r := rateLimits.rateLimit(req); r != nil {
 				return *r
 			}
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
@@ -111,7 +112,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 := rateLimit(req); r != nil {
+			if r := rateLimits.rateLimit(req); r != nil {
 				return *r
 			}
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
@@ -125,7 +126,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 := rateLimit(req); r != nil {
+			if r := rateLimits.rateLimit(req); r != nil {
 				return *r
 			}
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
@@ -148,7 +149,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 := rateLimit(req); r != nil {
+			if r := rateLimits.rateLimit(req); r != nil {
 				return *r
 			}
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
@@ -265,21 +266,21 @@ func Setup(
 	).Methods(http.MethodPut, http.MethodOptions)
 
 	r0mux.Handle("/register", httputil.MakeExternalAPI("register", func(req *http.Request) util.JSONResponse {
-		if r := rateLimit(req); r != nil {
+		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 := rateLimit(req); r != nil {
+		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 := rateLimit(req); r != nil {
+		if r := rateLimits.rateLimit(req); r != nil {
 			return *r
 		}
 		return RegisterAvailable(req, cfg, accountDB)
@@ -353,7 +354,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 := rateLimit(req); r != nil {
+			if r := rateLimits.rateLimit(req); r != nil {
 				return *r
 			}
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
@@ -409,7 +410,7 @@ func Setup(
 
 	r0mux.Handle("/account/whoami",
 		httputil.MakeAuthAPI("whoami", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
-			if r := rateLimit(req); r != nil {
+			if r := rateLimits.rateLimit(req); r != nil {
 				return *r
 			}
 			return Whoami(req, device)
@@ -420,7 +421,7 @@ func Setup(
 
 	r0mux.Handle("/login",
 		httputil.MakeExternalAPI("login", func(req *http.Request) util.JSONResponse {
-			if r := rateLimit(req); r != nil {
+			if r := rateLimits.rateLimit(req); r != nil {
 				return *r
 			}
 			return Login(req, accountDB, userAPI, cfg)
@@ -477,7 +478,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 := rateLimit(req); r != nil {
+			if r := rateLimits.rateLimit(req); r != nil {
 				return *r
 			}
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
@@ -502,7 +503,7 @@ func Setup(
 
 	r0mux.Handle("/profile/{userID}/displayname",
 		httputil.MakeAuthAPI("profile_displayname", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
-			if r := rateLimit(req); r != nil {
+			if r := rateLimits.rateLimit(req); r != nil {
 				return *r
 			}
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
@@ -542,7 +543,7 @@ 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 := rateLimit(req); r != nil {
+			if r := rateLimits.rateLimit(req); r != nil {
 				return *r
 			}
 			// TODO: Set presence (probably the responsibility of a presence server not clientapi)
@@ -555,7 +556,7 @@ func Setup(
 
 	r0mux.Handle("/voip/turnServer",
 		httputil.MakeAuthAPI("turn_server", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
-			if r := rateLimit(req); r != nil {
+			if r := rateLimits.rateLimit(req); r != nil {
 				return *r
 			}
 			return RequestTurnServer(req, device, cfg)
@@ -624,7 +625,7 @@ func Setup(
 
 	r0mux.Handle("/user_directory/search",
 		httputil.MakeAuthAPI("userdirectory_search", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
-			if r := rateLimit(req); r != nil {
+			if r := rateLimits.rateLimit(req); r != nil {
 				return *r
 			}
 			postContent := struct {
@@ -668,7 +669,7 @@ func Setup(
 
 	r0mux.Handle("/rooms/{roomID}/read_markers",
 		httputil.MakeExternalAPI("rooms_read_markers", func(req *http.Request) util.JSONResponse {
-			if r := rateLimit(req); r != nil {
+			if r := rateLimits.rateLimit(req); r != nil {
 				return *r
 			}
 			// TODO: return the read_markers.
@@ -769,7 +770,7 @@ func Setup(
 
 	r0mux.Handle("/capabilities",
 		httputil.MakeAuthAPI("capabilities", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
-			if r := rateLimit(req); r != nil {
+			if r := rateLimits.rateLimit(req); r != nil {
 				return *r
 			}
 			return GetCapabilities(req, rsAPI)

internal/httputil/httpapi.go

@@ -55,7 +55,6 @@ func MakeAuthAPI(
 		if err != nil {
 			return *err
 		}
-
 		// add the user ID to the logger
 		logger := util.GetLogger((req.Context()))
 		logger = logger.WithField("user_id", device.UserID)