1
0
Fork 0
mirror of https://github.com/binwiederhier/ntfy.git synced 2025-06-10 06:54:30 +02:00

Add limiters and database changes

This commit is contained in:
binwiederhier 2023-05-07 11:59:15 -04:00
parent 113b7c8a08
commit f9e2d6ddcb
11 changed files with 173 additions and 32 deletions

View file

@ -56,6 +56,8 @@ type visitor struct {
requestLimiter *rate.Limiter // Rate limiter for (almost) all requests (including messages)
messagesLimiter *util.FixedLimiter // Rate limiter for messages
emailsLimiter *util.RateLimiter // Rate limiter for emails
smsLimiter *util.FixedLimiter // Rate limiter for SMS
callsLimiter *util.FixedLimiter // Rate limiter for calls
subscriptionLimiter *util.FixedLimiter // Fixed limiter for active subscriptions (ongoing connections)
bandwidthLimiter *util.RateLimiter // Limiter for attachment bandwidth downloads
accountLimiter *rate.Limiter // Rate limiter for account creation, may be nil
@ -79,6 +81,8 @@ type visitorLimits struct {
EmailLimit int64
EmailLimitBurst int
EmailLimitReplenish rate.Limit
SMSLimit int64
CallLimit int64
ReservationsLimit int64
AttachmentTotalSizeLimit int64
AttachmentFileSizeLimit int64
@ -91,6 +95,10 @@ type visitorStats struct {
MessagesRemaining int64
Emails int64
EmailsRemaining int64
SMS int64
SMSRemaining int64
Calls int64
CallsRemaining int64
Reservations int64
ReservationsRemaining int64
AttachmentTotalSize int64
@ -107,10 +115,12 @@ const (
)
func newVisitor(conf *Config, messageCache *messageCache, userManager *user.Manager, ip netip.Addr, user *user.User) *visitor {
var messages, emails int64
var messages, emails, sms, calls int64
if user != nil {
messages = user.Stats.Messages
emails = user.Stats.Emails
sms = user.Stats.SMS
calls = user.Stats.Calls
}
v := &visitor{
config: conf,
@ -124,11 +134,13 @@ func newVisitor(conf *Config, messageCache *messageCache, userManager *user.Mana
requestLimiter: nil, // Set in resetLimiters
messagesLimiter: nil, // Set in resetLimiters, may be nil
emailsLimiter: nil, // Set in resetLimiters
smsLimiter: nil, // Set in resetLimiters, may be nil
callsLimiter: nil, // Set in resetLimiters, may be nil
bandwidthLimiter: nil, // Set in resetLimiters
accountLimiter: nil, // Set in resetLimiters, may be nil
authLimiter: nil, // Set in resetLimiters, may be nil
}
v.resetLimitersNoLock(messages, emails, false)
v.resetLimitersNoLock(messages, emails, sms, calls, false)
return v
}
@ -147,12 +159,22 @@ func (v *visitor) contextNoLock() log.Context {
"visitor_messages": info.Stats.Messages,
"visitor_messages_limit": info.Limits.MessageLimit,
"visitor_messages_remaining": info.Stats.MessagesRemaining,
"visitor_emails": info.Stats.Emails,
"visitor_emails_limit": info.Limits.EmailLimit,
"visitor_emails_remaining": info.Stats.EmailsRemaining,
"visitor_request_limiter_limit": v.requestLimiter.Limit(),
"visitor_request_limiter_tokens": v.requestLimiter.Tokens(),
}
if v.config.SMTPSenderFrom != "" {
fields["visitor_emails"] = info.Stats.Emails
fields["visitor_emails_limit"] = info.Limits.EmailLimit
fields["visitor_emails_remaining"] = info.Stats.EmailsRemaining
}
if v.config.TwilioAccount != "" {
fields["visitor_sms"] = info.Stats.SMS
fields["visitor_sms_limit"] = info.Limits.SMSLimit
fields["visitor_sms_remaining"] = info.Stats.SMSRemaining
fields["visitor_calls"] = info.Stats.Calls
fields["visitor_calls_limit"] = info.Limits.CallLimit
fields["visitor_calls_remaining"] = info.Stats.CallsRemaining
}
if v.authLimiter != nil {
fields["visitor_auth_limiter_limit"] = v.authLimiter.Limit()
fields["visitor_auth_limiter_tokens"] = v.authLimiter.Tokens()
@ -216,6 +238,18 @@ func (v *visitor) EmailAllowed() bool {
return v.emailsLimiter.Allow()
}
func (v *visitor) SMSAllowed() bool {
v.mu.RLock() // limiters could be replaced!
defer v.mu.RUnlock()
return v.smsLimiter.Allow()
}
func (v *visitor) CallAllowed() bool {
v.mu.RLock() // limiters could be replaced!
defer v.mu.RUnlock()
return v.callsLimiter.Allow()
}
func (v *visitor) SubscriptionAllowed() bool {
v.mu.RLock() // limiters could be replaced!
defer v.mu.RUnlock()
@ -296,6 +330,8 @@ func (v *visitor) Stats() *user.Stats {
return &user.Stats{
Messages: v.messagesLimiter.Value(),
Emails: v.emailsLimiter.Value(),
SMS: v.smsLimiter.Value(),
Calls: v.callsLimiter.Value(),
}
}
@ -304,6 +340,8 @@ func (v *visitor) ResetStats() {
defer v.mu.RUnlock()
v.emailsLimiter.Reset()
v.messagesLimiter.Reset()
v.smsLimiter.Reset()
v.callsLimiter.Reset()
}
// User returns the visitor user, or nil if there is none
@ -334,11 +372,11 @@ func (v *visitor) SetUser(u *user.User) {
shouldResetLimiters := v.user.TierID() != u.TierID() // TierID works with nil receiver
v.user = u // u may be nil!
if shouldResetLimiters {
var messages, emails int64
var messages, emails, sms, calls int64
if u != nil {
messages, emails = u.Stats.Messages, u.Stats.Emails
messages, emails, sms, calls = u.Stats.Messages, u.Stats.Emails, u.Stats.SMS, u.Stats.Calls
}
v.resetLimitersNoLock(messages, emails, true)
v.resetLimitersNoLock(messages, emails, sms, calls, true)
}
}
@ -353,11 +391,13 @@ func (v *visitor) MaybeUserID() string {
return ""
}
func (v *visitor) resetLimitersNoLock(messages, emails int64, enqueueUpdate bool) {
func (v *visitor) resetLimitersNoLock(messages, emails, sms, calls int64, enqueueUpdate bool) {
limits := v.limitsNoLock()
v.requestLimiter = rate.NewLimiter(limits.RequestLimitReplenish, limits.RequestLimitBurst)
v.messagesLimiter = util.NewFixedLimiterWithValue(limits.MessageLimit, messages)
v.emailsLimiter = util.NewRateLimiterWithValue(limits.EmailLimitReplenish, limits.EmailLimitBurst, emails)
v.smsLimiter = util.NewFixedLimiterWithValue(limits.SMSLimit, sms)
v.callsLimiter = util.NewFixedLimiterWithValue(limits.CallLimit, calls)
v.bandwidthLimiter = util.NewBytesLimiter(int(limits.AttachmentBandwidthLimit), oneDay)
if v.user == nil {
v.accountLimiter = rate.NewLimiter(rate.Every(v.config.VisitorAccountCreationLimitReplenish), v.config.VisitorAccountCreationLimitBurst)
@ -370,6 +410,8 @@ func (v *visitor) resetLimitersNoLock(messages, emails int64, enqueueUpdate bool
go v.userManager.EnqueueUserStats(v.user.ID, &user.Stats{
Messages: messages,
Emails: emails,
SMS: sms,
Calls: calls,
})
}
log.Fields(v.contextNoLock()).Debug("Rate limiters reset for visitor") // Must be after function, because contextNoLock() describes rate limiters
@ -398,6 +440,8 @@ func tierBasedVisitorLimits(conf *Config, tier *user.Tier) *visitorLimits {
EmailLimit: tier.EmailLimit,
EmailLimitBurst: util.MinMax(int(float64(tier.EmailLimit)*visitorEmailLimitBurstRate), conf.VisitorEmailLimitBurst, visitorEmailLimitBurstMax),
EmailLimitReplenish: dailyLimitToRate(tier.EmailLimit),
SMSLimit: tier.SMSLimit,
CallLimit: tier.CallLimit,
ReservationsLimit: tier.ReservationLimit,
AttachmentTotalSizeLimit: tier.AttachmentTotalSizeLimit,
AttachmentFileSizeLimit: tier.AttachmentFileSizeLimit,
@ -420,6 +464,8 @@ func configBasedVisitorLimits(conf *Config) *visitorLimits {
EmailLimit: replenishDurationToDailyLimit(conf.VisitorEmailLimitReplenish), // Approximation!
EmailLimitBurst: conf.VisitorEmailLimitBurst,
EmailLimitReplenish: rate.Every(conf.VisitorEmailLimitReplenish),
SMSLimit: int64(conf.VisitorSMSDailyLimit),
CallLimit: int64(conf.VisitorCallDailyLimit),
ReservationsLimit: visitorDefaultReservationsLimit,
AttachmentTotalSizeLimit: conf.VisitorAttachmentTotalSizeLimit,
AttachmentFileSizeLimit: conf.AttachmentFileSizeLimit,
@ -465,12 +511,18 @@ func (v *visitor) Info() (*visitorInfo, error) {
func (v *visitor) infoLightNoLock() *visitorInfo {
messages := v.messagesLimiter.Value()
emails := v.emailsLimiter.Value()
sms := v.smsLimiter.Value()
calls := v.callsLimiter.Value()
limits := v.limitsNoLock()
stats := &visitorStats{
Messages: messages,
MessagesRemaining: zeroIfNegative(limits.MessageLimit - messages),
Emails: emails,
EmailsRemaining: zeroIfNegative(limits.EmailLimit - emails),
SMS: sms,
SMSRemaining: zeroIfNegative(limits.SMSLimit - sms),
Calls: calls,
CallsRemaining: zeroIfNegative(limits.CallLimit - calls),
}
return &visitorInfo{
Limits: limits,