0
0
Fork 0
mirror of https://github.com/matrix-org/dendrite synced 2024-12-15 02:23:43 +01:00
dendrite/internal/httputil/rate_limiting.go
Neil Alexander 6d4bd5d890
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
2022-06-07 14:24:04 +01:00

132 lines
3.6 KiB
Go

package httputil
import (
"net/http"
"sync"
"time"
"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"
)
type RateLimits struct {
limits map[string]chan struct{}
limitsMutex sync.RWMutex
cleanMutex sync.RWMutex
enabled bool
requestThreshold int64
cooloffDuration time.Duration
exemptUserIDs map[string]struct{}
}
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,
exemptUserIDs: map[string]struct{}{},
}
for _, userID := range cfg.ExemptUserIDs {
l.exemptUserIDs[userID] = struct{}{}
}
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) Limit(req *http.Request, device *userapi.Device) *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.
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.
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
}