Rate limiting refactor, race fixes, more tests

This commit is contained in:
binwiederhier 2023-01-27 11:33:51 -05:00
parent ccc2dd1128
commit 62140ec001
8 changed files with 241 additions and 117 deletions

View file

@ -35,27 +35,19 @@ import (
) )
/* /*
TODO
--
- HIGH Rate limiting: Sensitive endpoints (account/login/change-password/...) - HIGH Rate limiting: Sensitive endpoints (account/login/change-password/...)
- HIGH Rate limiting: When ResetStats() is run, reset messagesLimiter (and others)?
- MEDIUM Rate limiting: Test daily message quota read from database initially
- MEDIUM: Races with v.user (see publishSyncEventAsync test) - MEDIUM: Races with v.user (see publishSyncEventAsync test)
- MEDIUM: Test that anonymous user and user without tier are the same visitor
- MEDIUM: Make sure account endpoints make sense for admins
- MEDIUM: Reservation (UI): Show "This topic is reserved" error message when trying to reserve a reserved topic (Thorben) - MEDIUM: Reservation (UI): Show "This topic is reserved" error message when trying to reserve a reserved topic (Thorben)
- MEDIUM: Reservation (UI): Ask for confirmation when removing reservation (deadcade) - MEDIUM: Reservation (UI): Ask for confirmation when removing reservation (deadcade)
- MEDIUM: Reservation table delete button: dialog "keep or delete messages?" - MEDIUM: Reservation table delete button: dialog "keep or delete messages?"
- MEDIUM: Tests for remaining payment endpoints
- LOW: UI: Flickering upgrade banner when logging in - LOW: UI: Flickering upgrade banner when logging in
- LOW: JS constants - LOW: JS constants
- LOW: Payments reconciliation process - LOW: Payments reconciliation process
Limits & rate limiting:
users without tier: should the stats be persisted? are they meaningful? -> test that the visitor is based on the IP address!
Make sure account endpoints make sense for admins
Tests:
- Payment endpoints (make mocks)
*/ */
// Server is the main server, providing the UI and API for ntfy // Server is the main server, providing the UI and API for ntfy
@ -513,7 +505,7 @@ func (s *Server) handleFile(w http.ResponseWriter, r *http.Request, v *visitor)
return errHTTPNotFound return errHTTPNotFound
} }
if r.Method == http.MethodGet { if r.Method == http.MethodGet {
if err := v.BandwidthLimiter().Allow(stat.Size()); err != nil { if !v.BandwidthAllowed(stat.Size()) {
return errHTTPTooManyRequestsLimitAttachmentBandwidth return errHTTPTooManyRequestsLimitAttachmentBandwidth
} }
} }
@ -543,7 +535,7 @@ func (s *Server) handlePublishWithoutResponse(r *http.Request, v *visitor) (*mes
if err != nil { if err != nil {
return nil, err return nil, err
} }
if err := v.MessageAllowed(); err != nil { if !v.MessageAllowed() {
return nil, errHTTPTooManyRequestsLimitMessages return nil, errHTTPTooManyRequestsLimitMessages
} }
body, err := util.Peek(r.Body, s.config.MessageLimit) body, err := util.Peek(r.Body, s.config.MessageLimit)
@ -558,9 +550,7 @@ func (s *Server) handlePublishWithoutResponse(r *http.Request, v *visitor) (*mes
if m.PollID != "" { if m.PollID != "" {
m = newPollRequestMessage(t.ID, m.PollID) m = newPollRequestMessage(t.ID, m.PollID)
} }
if v.user != nil { m.User = v.MaybeUserID()
m.User = v.user.ID
}
m.Expires = time.Now().Add(v.Limits().MessageExpiryDuration).Unix() m.Expires = time.Now().Add(v.Limits().MessageExpiryDuration).Unix()
if err := s.handlePublishBody(r, v, m, body, unifiedpush); err != nil { if err := s.handlePublishBody(r, v, m, body, unifiedpush); err != nil {
return nil, err return nil, err
@ -582,7 +572,6 @@ func (s *Server) handlePublishWithoutResponse(r *http.Request, v *visitor) (*mes
go s.sendToFirebase(v, m) go s.sendToFirebase(v, m)
} }
if s.smtpSender != nil && email != "" { if s.smtpSender != nil && email != "" {
v.IncrementEmails()
go s.sendEmail(v, m, email) go s.sendEmail(v, m, email)
} }
if s.config.UpstreamBaseURL != "" { if s.config.UpstreamBaseURL != "" {
@ -597,8 +586,9 @@ func (s *Server) handlePublishWithoutResponse(r *http.Request, v *visitor) (*mes
return nil, err return nil, err
} }
} }
if s.userManager != nil && v.user != nil { u := v.User()
s.userManager.EnqueueStats(v.user.ID, v.Stats()) // FIXME this makes no sense for tier-less users if s.userManager != nil && u != nil && u.Tier != nil {
s.userManager.EnqueueStats(u.ID, v.Stats())
} }
s.mu.Lock() s.mu.Lock()
s.messages++ s.messages++
@ -704,7 +694,7 @@ func (s *Server) parsePublishParams(r *http.Request, v *visitor, m *message) (ca
} }
email = readParam(r, "x-email", "x-e-mail", "email", "e-mail", "mail", "e") email = readParam(r, "x-email", "x-e-mail", "email", "e-mail", "mail", "e")
if email != "" { if email != "" {
if err := v.EmailAllowed(); err != nil { if !v.EmailAllowed() {
return false, false, "", false, errHTTPTooManyRequestsLimitEmails return false, false, "", false, errHTTPTooManyRequestsLimitEmails
} }
} }
@ -909,7 +899,7 @@ func (s *Server) handleSubscribeRaw(w http.ResponseWriter, r *http.Request, v *v
func (s *Server) handleSubscribeHTTP(w http.ResponseWriter, r *http.Request, v *visitor, contentType string, encoder messageEncoder) error { func (s *Server) handleSubscribeHTTP(w http.ResponseWriter, r *http.Request, v *visitor, contentType string, encoder messageEncoder) error {
log.Debug("%s HTTP stream connection opened", logHTTPPrefix(v, r)) log.Debug("%s HTTP stream connection opened", logHTTPPrefix(v, r))
defer log.Debug("%s HTTP stream connection closed", logHTTPPrefix(v, r)) defer log.Debug("%s HTTP stream connection closed", logHTTPPrefix(v, r))
if err := v.SubscriptionAllowed(); err != nil { if !v.SubscriptionAllowed() {
return errHTTPTooManyRequestsLimitSubscriptions return errHTTPTooManyRequestsLimitSubscriptions
} }
defer v.RemoveSubscription() defer v.RemoveSubscription()
@ -989,7 +979,7 @@ func (s *Server) handleSubscribeWS(w http.ResponseWriter, r *http.Request, v *vi
if strings.ToLower(r.Header.Get("Upgrade")) != "websocket" { if strings.ToLower(r.Header.Get("Upgrade")) != "websocket" {
return errHTTPBadRequestWebSocketsUpgradeHeaderMissing return errHTTPBadRequestWebSocketsUpgradeHeaderMissing
} }
if err := v.SubscriptionAllowed(); err != nil { if !v.SubscriptionAllowed() {
return errHTTPTooManyRequestsLimitSubscriptions return errHTTPTooManyRequestsLimitSubscriptions
} }
defer v.RemoveSubscription() defer v.RemoveSubscription()

View file

@ -23,7 +23,7 @@ func (s *Server) handleAccountCreate(w http.ResponseWriter, r *http.Request, v *
} else if v.user != nil { } else if v.user != nil {
return errHTTPUnauthorized // Cannot create account from user context return errHTTPUnauthorized // Cannot create account from user context
} }
if err := v.AccountCreationAllowed(); err != nil { if !v.AccountCreationAllowed() {
return errHTTPTooManyRequestsLimitAccountCreation return errHTTPTooManyRequestsLimitAccountCreation
} }
} }
@ -428,11 +428,12 @@ func (s *Server) publishSyncEvent(v *visitor) error {
func (s *Server) publishSyncEventAsync(v *visitor) { func (s *Server) publishSyncEventAsync(v *visitor) {
go func() { go func() {
if v.user == nil || v.user.SyncTopic == "" { u := v.User()
if u == nil || u.SyncTopic == "" {
return return
} }
if err := s.publishSyncEvent(v); err != nil { if err := s.publishSyncEvent(v); err != nil {
log.Trace("Error publishing to user %s's sync topic %s: %s", v.user.Name, v.user.SyncTopic, err.Error()) log.Trace("Error publishing to user %s's sync topic %s: %s", u.Name, u.SyncTopic, err.Error())
} }
}() }()
} }

View file

@ -841,23 +841,35 @@ func TestServer_StatsResetter(t *testing.T) {
require.Equal(t, int64(0), account.Stats.Messages) require.Equal(t, int64(0), account.Stats.Messages)
} }
func TestServer_StatsResetter_MessageLimiter(t *testing.T) { func TestServer_StatsResetter_MessageLimiter_EmailsLimiter(t *testing.T) {
// This tests that the messageLimiter (the only fixed limiter) is reset by the stats resetter // This tests that the messageLimiter (the only fixed limiter) and the emailsLimiter (token bucket)
// is reset by the stats resetter
c := newTestConfigWithAuthFile(t) c := newTestConfigWithAuthFile(t)
s := newTestServer(t, c) s := newTestServer(t, c)
s.smtpSender = &testMailer{}
// Publish some messages, and check stats // Publish some messages, and check stats
for i := 0; i < 3; i++ { for i := 0; i < 3; i++ {
response := request(t, s, "PUT", "/mytopic", "test", nil) response := request(t, s, "PUT", "/mytopic", "test", nil)
require.Equal(t, 200, response.Code) require.Equal(t, 200, response.Code)
} }
response := request(t, s, "PUT", "/mytopic", "test", map[string]string{
"Email": "test@email.com",
})
require.Equal(t, 200, response.Code)
rr := request(t, s, "GET", "/v1/account", "", nil) rr := request(t, s, "GET", "/v1/account", "", nil)
require.Equal(t, 200, rr.Code) require.Equal(t, 200, rr.Code)
account, err := util.UnmarshalJSON[apiAccountResponse](io.NopCloser(rr.Body)) account, err := util.UnmarshalJSON[apiAccountResponse](io.NopCloser(rr.Body))
require.Nil(t, err) require.Nil(t, err)
require.Equal(t, int64(3), account.Stats.Messages) require.Equal(t, int64(4), account.Stats.Messages)
require.Equal(t, int64(3), s.visitor(netip.MustParseAddr("9.9.9.9"), nil).messagesLimiter.Value()) require.Equal(t, int64(1), account.Stats.Emails)
v := s.visitor(netip.MustParseAddr("9.9.9.9"), nil)
require.Equal(t, int64(4), v.Stats().Messages)
require.Equal(t, int64(4), v.messagesLimiter.Value())
require.Equal(t, int64(1), v.Stats().Emails)
require.Equal(t, int64(1), v.emailsLimiter.Value())
// Reset stats and check again // Reset stats and check again
s.resetStats() s.resetStats()
@ -866,7 +878,53 @@ func TestServer_StatsResetter_MessageLimiter(t *testing.T) {
account, err = util.UnmarshalJSON[apiAccountResponse](io.NopCloser(rr.Body)) account, err = util.UnmarshalJSON[apiAccountResponse](io.NopCloser(rr.Body))
require.Nil(t, err) require.Nil(t, err)
require.Equal(t, int64(0), account.Stats.Messages) require.Equal(t, int64(0), account.Stats.Messages)
require.Equal(t, int64(0), s.visitor(netip.MustParseAddr("9.9.9.9"), nil).messagesLimiter.Value()) require.Equal(t, int64(0), account.Stats.Emails)
v = s.visitor(netip.MustParseAddr("9.9.9.9"), nil)
require.Equal(t, int64(0), v.Stats().Messages)
require.Equal(t, int64(0), v.messagesLimiter.Value())
require.Equal(t, int64(0), v.Stats().Emails)
require.Equal(t, int64(0), v.emailsLimiter.Value())
}
func TestServer_DailyMessageQuotaFromDatabase(t *testing.T) {
// This tests that the daily message quota is prefilled originally from the database,
// if the visitor is unknown
c := newTestConfigWithAuthFile(t)
s := newTestServer(t, c)
var err error
s.userManager, err = user.NewManagerWithStatsInterval(c.AuthFile, c.AuthStartupQueries, c.AuthDefault, 100*time.Millisecond)
require.Nil(t, err)
// Create user, and update it with some message and email stats
require.Nil(t, s.userManager.CreateTier(&user.Tier{
Code: "test",
}))
require.Nil(t, s.userManager.AddUser("phil", "phil", user.RoleUser))
require.Nil(t, s.userManager.ChangeTier("phil", "test"))
u, err := s.userManager.User("phil")
require.Nil(t, err)
s.userManager.EnqueueStats(u.ID, &user.Stats{
Messages: 123456,
Emails: 999,
})
time.Sleep(400 * time.Millisecond)
// Get account and verify stats are read from the DB, and that the visitor also has these stats
rr := request(t, s, "GET", "/v1/account", "", map[string]string{
"Authorization": util.BasicAuth("phil", "phil"),
})
require.Equal(t, 200, rr.Code)
account, err := util.UnmarshalJSON[apiAccountResponse](io.NopCloser(rr.Body))
require.Nil(t, err)
require.Equal(t, int64(123456), account.Stats.Messages)
require.Equal(t, int64(999), account.Stats.Emails)
v := s.visitor(netip.MustParseAddr("9.9.9.9"), u)
require.Equal(t, int64(123456), v.Stats().Messages)
require.Equal(t, int64(123456), v.messagesLimiter.Value())
require.Equal(t, int64(999), v.Stats().Emails)
require.Equal(t, int64(999), v.emailsLimiter.Value())
} }
type testMailer struct { type testMailer struct {

View file

@ -58,12 +58,11 @@ type visitor struct {
userManager *user.Manager // May be nil userManager *user.Manager // May be nil
ip netip.Addr // Visitor IP address ip netip.Addr // Visitor IP address
user *user.User // Only set if authenticated user, otherwise nil user *user.User // Only set if authenticated user, otherwise nil
emails int64 // Number of emails sent, reset every day
requestLimiter *rate.Limiter // Rate limiter for (almost) all requests (including messages) requestLimiter *rate.Limiter // Rate limiter for (almost) all requests (including messages)
messagesLimiter *util.FixedLimiter // Rate limiter for messages messagesLimiter *util.FixedLimiter // Rate limiter for messages
emailsLimiter *rate.Limiter // Rate limiter for emails emailsLimiter *util.RateLimiter // Rate limiter for emails
subscriptionLimiter util.Limiter // Fixed limiter for active subscriptions (ongoing connections) subscriptionLimiter *util.FixedLimiter // Fixed limiter for active subscriptions (ongoing connections)
bandwidthLimiter util.Limiter // Limiter for attachment bandwidth downloads bandwidthLimiter *util.RateLimiter // Limiter for attachment bandwidth downloads
accountLimiter *rate.Limiter // Rate limiter for account creation, may be nil accountLimiter *rate.Limiter // Rate limiter for account creation, may be nil
firebase time.Time // Next allowed Firebase message firebase time.Time // Next allowed Firebase message
seen time.Time // Last seen time of this visitor (needed for removal of stale visitors) seen time.Time // Last seen time of this visitor (needed for removal of stale visitors)
@ -123,7 +122,6 @@ func newVisitor(conf *Config, messageCache *messageCache, userManager *user.Mana
userManager: userManager, // May be nil userManager: userManager, // May be nil
ip: ip, ip: ip,
user: user, user: user,
emails: emails,
firebase: time.Unix(0, 0), firebase: time.Unix(0, 0),
seen: time.Now(), seen: time.Now(),
subscriptionLimiter: util.NewFixedLimiter(int64(conf.VisitorSubscriptionLimit)), subscriptionLimiter: util.NewFixedLimiter(int64(conf.VisitorSubscriptionLimit)),
@ -133,7 +131,7 @@ func newVisitor(conf *Config, messageCache *messageCache, userManager *user.Mana
bandwidthLimiter: nil, // Set in resetLimiters bandwidthLimiter: nil, // Set in resetLimiters
accountLimiter: nil, // Set in resetLimiters, may be nil accountLimiter: nil, // Set in resetLimiters, may be nil
} }
v.resetLimiters(messages) v.resetLimitersNoLock(messages, emails)
return v return v
} }
@ -153,6 +151,8 @@ func (v *visitor) stringNoLock() string {
} }
func (v *visitor) RequestAllowed() error { func (v *visitor) RequestAllowed() error {
v.mu.Lock() // limiters could be replaced!
defer v.mu.Unlock()
if !v.requestLimiter.Allow() { if !v.requestLimiter.Allow() {
return errVisitorLimitReached return errVisitorLimitReached
} }
@ -174,40 +174,43 @@ func (v *visitor) FirebaseTemporarilyDeny() {
v.firebase = time.Now().Add(v.config.FirebaseQuotaExceededPenaltyDuration) v.firebase = time.Now().Add(v.config.FirebaseQuotaExceededPenaltyDuration)
} }
func (v *visitor) MessageAllowed() error { func (v *visitor) MessageAllowed() bool {
if v.messagesLimiter.Allow(1) != nil { v.mu.Lock() // limiters could be replaced!
return errVisitorLimitReached
}
return nil
}
func (v *visitor) EmailAllowed() error {
if !v.emailsLimiter.Allow() {
return errVisitorLimitReached
}
return nil
}
func (v *visitor) SubscriptionAllowed() error {
v.mu.Lock()
defer v.mu.Unlock() defer v.mu.Unlock()
if err := v.subscriptionLimiter.Allow(1); err != nil { return v.messagesLimiter.Allow()
return errVisitorLimitReached
}
return nil
} }
func (v *visitor) AccountCreationAllowed() error { func (v *visitor) EmailAllowed() bool {
v.mu.Lock() // limiters could be replaced!
defer v.mu.Unlock()
return v.emailsLimiter.Allow()
}
func (v *visitor) SubscriptionAllowed() bool {
v.mu.Lock() // limiters could be replaced!
defer v.mu.Unlock()
return v.subscriptionLimiter.Allow()
}
func (v *visitor) AccountCreationAllowed() bool {
v.mu.Lock() // limiters could be replaced!
defer v.mu.Unlock()
if v.accountLimiter == nil || (v.accountLimiter != nil && !v.accountLimiter.Allow()) { if v.accountLimiter == nil || (v.accountLimiter != nil && !v.accountLimiter.Allow()) {
return errVisitorLimitReached return false
} }
return nil return true
}
func (v *visitor) BandwidthAllowed(bytes int64) bool {
v.mu.Lock() // limiters could be replaced!
defer v.mu.Unlock()
return v.bandwidthLimiter.AllowN(bytes)
} }
func (v *visitor) RemoveSubscription() { func (v *visitor) RemoveSubscription() {
v.mu.Lock() v.mu.Lock()
defer v.mu.Unlock() defer v.mu.Unlock()
v.subscriptionLimiter.Allow(-1) v.subscriptionLimiter.AllowN(-1)
} }
func (v *visitor) Keepalive() { func (v *visitor) Keepalive() {
@ -217,6 +220,8 @@ func (v *visitor) Keepalive() {
} }
func (v *visitor) BandwidthLimiter() util.Limiter { func (v *visitor) BandwidthLimiter() util.Limiter {
v.mu.Lock() // limiters could be replaced!
defer v.mu.Unlock()
return v.bandwidthLimiter return v.bandwidthLimiter
} }
@ -226,26 +231,27 @@ func (v *visitor) Stale() bool {
return time.Since(v.seen) > visitorExpungeAfter return time.Since(v.seen) > visitorExpungeAfter
} }
func (v *visitor) IncrementEmails() {
v.mu.Lock()
defer v.mu.Unlock()
v.emails++
}
func (v *visitor) Stats() *user.Stats { func (v *visitor) Stats() *user.Stats {
v.mu.Lock() v.mu.Lock() // limiters could be replaced!
defer v.mu.Unlock() defer v.mu.Unlock()
return &user.Stats{ return &user.Stats{
Messages: v.messagesLimiter.Value(), Messages: v.messagesLimiter.Value(),
Emails: v.emails, Emails: v.emailsLimiter.Value(),
} }
} }
func (v *visitor) ResetStats() { func (v *visitor) ResetStats() {
v.mu.Lock() // limiters could be replaced!
defer v.mu.Unlock()
v.emailsLimiter.Reset()
v.messagesLimiter.Reset()
}
// User returns the visitor user, or nil if there is none
func (v *visitor) User() *user.User {
v.mu.Lock() v.mu.Lock()
defer v.mu.Unlock() defer v.mu.Unlock()
v.emails = 0 return v.user // May be nil
v.messagesLimiter.Reset()
} }
// SetUser sets the visitors user to the given value // SetUser sets the visitors user to the given value
@ -255,7 +261,7 @@ func (v *visitor) SetUser(u *user.User) {
shouldResetLimiters := v.user.TierID() != u.TierID() // TierID works with nil receiver shouldResetLimiters := v.user.TierID() != u.TierID() // TierID works with nil receiver
v.user = u v.user = u
if shouldResetLimiters { if shouldResetLimiters {
v.resetLimiters(0) v.resetLimitersNoLock(0, 0)
} }
} }
@ -270,12 +276,12 @@ func (v *visitor) MaybeUserID() string {
return "" return ""
} }
func (v *visitor) resetLimiters(messages int64) { func (v *visitor) resetLimitersNoLock(messages, emails int64) {
log.Debug("%s Resetting limiters for visitor", v.stringNoLock()) log.Debug("%s Resetting limiters for visitor", v.stringNoLock())
limits := v.limitsNoLock() limits := v.limitsNoLock()
v.requestLimiter = rate.NewLimiter(limits.RequestLimitReplenish, limits.RequestLimitBurst) v.requestLimiter = rate.NewLimiter(limits.RequestLimitReplenish, limits.RequestLimitBurst)
v.messagesLimiter = util.NewFixedLimiterWithValue(limits.MessageLimit, messages) v.messagesLimiter = util.NewFixedLimiterWithValue(limits.MessageLimit, messages)
v.emailsLimiter = rate.NewLimiter(limits.EmailLimitReplenish, limits.EmailLimitBurst) v.emailsLimiter = util.NewRateLimiterWithValue(limits.EmailLimitReplenish, limits.EmailLimitBurst, emails)
v.bandwidthLimiter = util.NewBytesLimiter(int(limits.AttachmentBandwidthLimit), oneDay) v.bandwidthLimiter = util.NewBytesLimiter(int(limits.AttachmentBandwidthLimit), oneDay)
if v.user == nil { if v.user == nil {
v.accountLimiter = rate.NewLimiter(rate.Every(v.config.VisitorAccountCreationLimitReplenish), v.config.VisitorAccountCreationLimitBurst) v.accountLimiter = rate.NewLimiter(rate.Every(v.config.VisitorAccountCreationLimitReplenish), v.config.VisitorAccountCreationLimitBurst)
@ -340,12 +346,13 @@ func configBasedVisitorLimits(conf *Config) *visitorLimits {
func (v *visitor) Info() (*visitorInfo, error) { func (v *visitor) Info() (*visitorInfo, error) {
v.mu.Lock() v.mu.Lock()
messages := v.messagesLimiter.Value() messages := v.messagesLimiter.Value()
emails := v.emails emails := v.emailsLimiter.Value()
v.mu.Unlock() v.mu.Unlock()
var attachmentsBytesUsed int64 var attachmentsBytesUsed int64
var err error var err error
if v.user != nil { u := v.User()
attachmentsBytesUsed, err = v.messageCache.AttachmentBytesUsedByUser(v.user.ID) if u != nil {
attachmentsBytesUsed, err = v.messageCache.AttachmentBytesUsedByUser(u.ID)
} else { } else {
attachmentsBytesUsed, err = v.messageCache.AttachmentBytesUsedBySender(v.ip.String()) attachmentsBytesUsed, err = v.messageCache.AttachmentBytesUsedBySender(v.ip.String())
} }
@ -353,8 +360,8 @@ func (v *visitor) Info() (*visitorInfo, error) {
return nil, err return nil, err
} }
var reservations int64 var reservations int64
if v.user != nil && v.userManager != nil { if v.userManager != nil && u != nil {
reservations, err = v.userManager.ReservationsCount(v.user.Name) reservations, err = v.userManager.ReservationsCount(u.Name)
if err != nil { if err != nil {
return nil, err return nil, err
} }

View file

@ -301,11 +301,11 @@ var _ Auther = (*Manager)(nil)
// NewManager creates a new Manager instance // NewManager creates a new Manager instance
func NewManager(filename, startupQueries string, defaultAccess Permission) (*Manager, error) { func NewManager(filename, startupQueries string, defaultAccess Permission) (*Manager, error) {
return newManager(filename, startupQueries, defaultAccess, userStatsQueueWriterInterval) return NewManagerWithStatsInterval(filename, startupQueries, defaultAccess, userStatsQueueWriterInterval)
} }
// NewManager creates a new Manager instance // NewManagerWithStatsInterval creates a new Manager instance
func newManager(filename, startupQueries string, defaultAccess Permission, statsWriterInterval time.Duration) (*Manager, error) { func NewManagerWithStatsInterval(filename, startupQueries string, defaultAccess Permission, statsWriterInterval time.Duration) (*Manager, error) {
db, err := sql.Open("sqlite3", filename) db, err := sql.Open("sqlite3", filename)
if err != nil { if err != nil {
return nil, err return nil, err

View file

@ -545,7 +545,7 @@ func TestManager_Token_MaxCount_AutoDelete(t *testing.T) {
} }
func TestManager_EnqueueStats(t *testing.T) { func TestManager_EnqueueStats(t *testing.T) {
a, err := newManager(filepath.Join(t.TempDir(), "db"), "", PermissionReadWrite, 1500*time.Millisecond) a, err := NewManagerWithStatsInterval(filepath.Join(t.TempDir(), "db"), "", PermissionReadWrite, 1500*time.Millisecond)
require.Nil(t, err) require.Nil(t, err)
require.Nil(t, a.AddUser("ben", "ben", RoleUser)) require.Nil(t, a.AddUser("ben", "ben", RoleUser))
@ -575,7 +575,7 @@ func TestManager_EnqueueStats(t *testing.T) {
} }
func TestManager_ChangeSettings(t *testing.T) { func TestManager_ChangeSettings(t *testing.T) {
a, err := newManager(filepath.Join(t.TempDir(), "db"), "", PermissionReadWrite, 1500*time.Millisecond) a, err := NewManagerWithStatsInterval(filepath.Join(t.TempDir(), "db"), "", PermissionReadWrite, 1500*time.Millisecond)
require.Nil(t, err) require.Nil(t, err)
require.Nil(t, a.AddUser("ben", "ben", RoleUser)) require.Nil(t, a.AddUser("ben", "ben", RoleUser))
@ -718,7 +718,7 @@ func newTestManager(t *testing.T, defaultAccess Permission) *Manager {
} }
func newTestManagerFromFile(t *testing.T, filename, startupQueries string, defaultAccess Permission, statsWriterInterval time.Duration) *Manager { func newTestManagerFromFile(t *testing.T, filename, startupQueries string, defaultAccess Permission, statsWriterInterval time.Duration) *Manager {
a, err := newManager(filename, startupQueries, defaultAccess, statsWriterInterval) a, err := NewManagerWithStatsInterval(filename, startupQueries, defaultAccess, statsWriterInterval)
require.Nil(t, err) require.Nil(t, err)
return a return a
} }

View file

@ -13,8 +13,17 @@ var ErrLimitReached = errors.New("limit reached")
// Limiter is an interface that implements a rate limiting mechanism, e.g. based on time or a fixed value // Limiter is an interface that implements a rate limiting mechanism, e.g. based on time or a fixed value
type Limiter interface { type Limiter interface {
// Allow adds n to the limiters internal value, or returns ErrLimitReached if the limit has been reached // Allow adds one to the limiters value, or returns false if the limit has been reached
Allow(n int64) error Allow() bool
// AllowN adds n to the limiters value, or returns false if the limit has been reached
AllowN(n int64) bool
// Value returns the current internal limiter value
Value() int64
// Reset resets the state of the limiter
Reset()
} }
// FixedLimiter is a helper that allows adding values up to a well-defined limit. Once the limit is reached // FixedLimiter is a helper that allows adding values up to a well-defined limit. Once the limit is reached
@ -25,6 +34,8 @@ type FixedLimiter struct {
mu sync.Mutex mu sync.Mutex
} }
var _ Limiter = (*FixedLimiter)(nil)
// NewFixedLimiter creates a new Limiter // NewFixedLimiter creates a new Limiter
func NewFixedLimiter(limit int64) *FixedLimiter { func NewFixedLimiter(limit int64) *FixedLimiter {
return NewFixedLimiterWithValue(limit, 0) return NewFixedLimiterWithValue(limit, 0)
@ -38,16 +49,22 @@ func NewFixedLimiterWithValue(limit, value int64) *FixedLimiter {
} }
} }
// Allow adds n to the limiters internal value, but only if the limit has not been reached. If the limit was // Allow adds one to the limiters internal value, but only if the limit has not been reached. If the limit was
// exceeded after adding n, ErrLimitReached is returned. // exceeded, false is returned.
func (l *FixedLimiter) Allow(n int64) error { func (l *FixedLimiter) Allow() bool {
return l.AllowN(1)
}
// AllowN adds n to the limiters internal value, but only if the limit has not been reached. If the limit was
// exceeded after adding n, false is returned.
func (l *FixedLimiter) AllowN(n int64) bool {
l.mu.Lock() l.mu.Lock()
defer l.mu.Unlock() defer l.mu.Unlock()
if l.value+n > l.limit { if l.value+n > l.limit {
return ErrLimitReached return false
} }
l.value += n l.value += n
return nil return true
} }
// Value returns the current limiter value // Value returns the current limiter value
@ -66,12 +83,29 @@ func (l *FixedLimiter) Reset() {
// RateLimiter is a Limiter that wraps a rate.Limiter, allowing a floating time-based limit. // RateLimiter is a Limiter that wraps a rate.Limiter, allowing a floating time-based limit.
type RateLimiter struct { type RateLimiter struct {
r rate.Limit
b int
value int64
limiter *rate.Limiter limiter *rate.Limiter
mu sync.Mutex
} }
var _ Limiter = (*RateLimiter)(nil)
// NewRateLimiter creates a new RateLimiter // NewRateLimiter creates a new RateLimiter
func NewRateLimiter(r rate.Limit, b int) *RateLimiter { func NewRateLimiter(r rate.Limit, b int) *RateLimiter {
return NewRateLimiterWithValue(r, b, 0)
}
// NewRateLimiterWithValue creates a new RateLimiter with the given starting value.
//
// Note that the starting value only has informational value. It does not impact the underlying
// value of the rate.Limiter.
func NewRateLimiterWithValue(r rate.Limit, b int, value int64) *RateLimiter {
return &RateLimiter{ return &RateLimiter{
r: r,
b: b,
value: value,
limiter: rate.NewLimiter(r, b), limiter: rate.NewLimiter(r, b),
} }
} }
@ -82,16 +116,40 @@ func NewBytesLimiter(bytes int, interval time.Duration) *RateLimiter {
return NewRateLimiter(rate.Limit(bytes)*rate.Every(interval), bytes) return NewRateLimiter(rate.Limit(bytes)*rate.Every(interval), bytes)
} }
// Allow adds n to the limiters internal value, but only if the limit has not been reached. If the limit was // Allow adds one to the limiters internal value, but only if the limit has not been reached. If the limit was
// exceeded after adding n, ErrLimitReached is returned. // exceeded, false is returned.
func (l *RateLimiter) Allow(n int64) error { func (l *RateLimiter) Allow() bool {
return l.AllowN(1)
}
// AllowN adds n to the limiters internal value, but only if the limit has not been reached. If the limit was
// exceeded after adding n, false is returned.
func (l *RateLimiter) AllowN(n int64) bool {
if n <= 0 { if n <= 0 {
return nil // No-op. Can't take back bytes you're written! return false // No-op. Can't take back bytes you're written!
} }
l.mu.Lock()
defer l.mu.Unlock()
if !l.limiter.AllowN(time.Now(), int(n)) { if !l.limiter.AllowN(time.Now(), int(n)) {
return ErrLimitReached return false
} }
return nil l.value += n
return true
}
// Value returns the current limiter value
func (l *RateLimiter) Value() int64 {
l.mu.Lock()
defer l.mu.Unlock()
return l.value
}
// Reset sets the limiter's value back to zero, and resets the underlying rate.Limiter
func (l *RateLimiter) Reset() {
l.mu.Lock()
defer l.mu.Unlock()
l.limiter = rate.NewLimiter(l.r, l.b)
l.value = 0
} }
// LimitWriter implements an io.Writer that will pass through all Write calls to the underlying // LimitWriter implements an io.Writer that will pass through all Write calls to the underlying
@ -117,9 +175,9 @@ func (w *LimitWriter) Write(p []byte) (n int, err error) {
w.mu.Lock() w.mu.Lock()
defer w.mu.Unlock() defer w.mu.Unlock()
for i := 0; i < len(w.limiters); i++ { for i := 0; i < len(w.limiters); i++ {
if err := w.limiters[i].Allow(int64(len(p))); err != nil { if !w.limiters[i].AllowN(int64(len(p))) {
for j := i - 1; j >= 0; j-- { for j := i - 1; j >= 0; j-- {
w.limiters[j].Allow(-int64(len(p))) // Revert limiters limits if allowed w.limiters[j].AllowN(-int64(len(p))) // Revert limiters limits if not allowed
} }
return 0, ErrLimitReached return 0, ErrLimitReached
} }

View file

@ -7,26 +7,31 @@ import (
"time" "time"
) )
func TestFixedLimiter_Add(t *testing.T) { func TestFixedLimiter_AllowValueReset(t *testing.T) {
l := NewFixedLimiter(10) l := NewFixedLimiter(10)
if err := l.Allow(5); err != nil { require.True(t, l.AllowN(5))
t.Fatal(err) require.Equal(t, int64(5), l.Value())
}
if err := l.Allow(5); err != nil { require.True(t, l.AllowN(5))
t.Fatal(err) require.Equal(t, int64(10), l.Value())
}
if err := l.Allow(5); err != ErrLimitReached { require.False(t, l.Allow())
t.Fatalf("expected ErrLimitReached, got %#v", err) require.Equal(t, int64(10), l.Value())
}
l.Reset()
require.Equal(t, int64(0), l.Value())
require.True(t, l.Allow())
require.True(t, l.AllowN(9))
require.False(t, l.Allow())
} }
func TestFixedLimiter_AddSub(t *testing.T) { func TestFixedLimiter_AddSub(t *testing.T) {
l := NewFixedLimiter(10) l := NewFixedLimiter(10)
l.Allow(5) l.AllowN(5)
if l.value != 5 { if l.value != 5 {
t.Fatalf("expected value to be %d, got %d", 5, l.value) t.Fatalf("expected value to be %d, got %d", 5, l.value)
} }
l.Allow(-2) l.AllowN(-2)
if l.value != 3 { if l.value != 3 {
t.Fatalf("expected value to be %d, got %d", 7, l.value) t.Fatalf("expected value to be %d, got %d", 7, l.value)
} }
@ -34,17 +39,22 @@ func TestFixedLimiter_AddSub(t *testing.T) {
func TestBytesLimiter_Add_Simple(t *testing.T) { func TestBytesLimiter_Add_Simple(t *testing.T) {
l := NewBytesLimiter(250*1024*1024, 24*time.Hour) // 250 MB per 24h l := NewBytesLimiter(250*1024*1024, 24*time.Hour) // 250 MB per 24h
require.Nil(t, l.Allow(100*1024*1024)) require.True(t, l.AllowN(100*1024*1024))
require.Nil(t, l.Allow(100*1024*1024)) require.Equal(t, int64(100*1024*1024), l.Value())
require.Equal(t, ErrLimitReached, l.Allow(300*1024*1024))
require.True(t, l.AllowN(100*1024*1024))
require.Equal(t, int64(200*1024*1024), l.Value())
require.False(t, l.AllowN(300*1024*1024))
require.Equal(t, int64(200*1024*1024), l.Value())
} }
func TestBytesLimiter_Add_Wait(t *testing.T) { func TestBytesLimiter_Add_Wait(t *testing.T) {
l := NewBytesLimiter(250*1024*1024, 24*time.Hour) // 250 MB per 24h (~ 303 bytes per 100ms) l := NewBytesLimiter(250*1024*1024, 24*time.Hour) // 250 MB per 24h (~ 303 bytes per 100ms)
require.Nil(t, l.Allow(250*1024*1024)) require.True(t, l.AllowN(250*1024*1024))
require.Equal(t, ErrLimitReached, l.Allow(400)) require.False(t, l.AllowN(400))
time.Sleep(200 * time.Millisecond) time.Sleep(200 * time.Millisecond)
require.Nil(t, l.Allow(400)) require.True(t, l.AllowN(400))
} }
func TestLimitWriter_WriteNoLimiter(t *testing.T) { func TestLimitWriter_WriteNoLimiter(t *testing.T) {