1
0
Fork 0
mirror of https://github.com/binwiederhier/ntfy.git synced 2025-06-26 06:18:04 +02:00

Limit work

This commit is contained in:
binwiederhier 2022-12-18 14:35:05 -05:00
parent 56ab34a57f
commit 42e46a7c22
9 changed files with 114 additions and 62 deletions

View file

@ -545,6 +545,7 @@ func (s *Server) handlePublishWithoutResponse(r *http.Request, v *visitor) (*mes
return nil, err
}
}
v.requests.Inc()
s.mu.Lock()
s.messages++
s.mu.Unlock()

View file

@ -59,26 +59,26 @@ func (s *Server) handleAccountGet(w http.ResponseWriter, r *http.Request, v *vis
if v.user.Plan != nil {
response.Usage.Basis = "account"
response.Plan = &apiAccountSettingsPlan{
Name: v.user.Plan.Name,
MessagesLimit: v.user.Plan.MessagesLimit,
EmailsLimit: v.user.Plan.EmailsLimit,
Code: v.user.Plan.Code,
RequestLimit: v.user.Plan.RequestLimit,
EmailLimit: v.user.Plan.EmailsLimit,
AttachmentsBytesLimit: v.user.Plan.AttachmentBytesLimit,
}
} else {
if v.user.Role == auth.RoleAdmin {
response.Usage.Basis = "account"
response.Plan = &apiAccountSettingsPlan{
Name: "Unlimited",
MessagesLimit: 0,
EmailsLimit: 0,
Code: string(auth.PlanUnlimited),
RequestLimit: 0,
EmailLimit: 0,
AttachmentsBytesLimit: 0,
}
} else {
response.Usage.Basis = "ip"
response.Plan = &apiAccountSettingsPlan{
Name: "Free",
MessagesLimit: s.config.VisitorRequestLimitBurst,
EmailsLimit: s.config.VisitorEmailLimitBurst,
Code: string(auth.PlanDefault),
RequestLimit: s.config.VisitorRequestLimitBurst,
EmailLimit: s.config.VisitorEmailLimitBurst,
AttachmentsBytesLimit: s.config.VisitorAttachmentTotalSizeLimit,
}
}
@ -88,13 +88,13 @@ func (s *Server) handleAccountGet(w http.ResponseWriter, r *http.Request, v *vis
response.Role = string(auth.RoleAnonymous)
response.Usage.Basis = "account"
response.Plan = &apiAccountSettingsPlan{
Name: "Anonymous",
MessagesLimit: s.config.VisitorRequestLimitBurst,
EmailsLimit: s.config.VisitorEmailLimitBurst,
Code: string(auth.PlanNone),
RequestLimit: s.config.VisitorRequestLimitBurst,
EmailLimit: s.config.VisitorEmailLimitBurst,
AttachmentsBytesLimit: s.config.VisitorAttachmentTotalSizeLimit,
}
}
response.Usage.Messages = int(v.requests.Tokens())
response.Usage.Requests = v.requests.Value()
response.Usage.AttachmentsBytes = stats.VisitorAttachmentBytesUsed
if err := json.NewEncoder(w).Encode(response); err != nil {
return err

View file

@ -225,15 +225,16 @@ type apiAccountTokenResponse struct {
}
type apiAccountSettingsPlan struct {
Name string `json:"name"`
MessagesLimit int `json:"messages_limit"`
EmailsLimit int `json:"emails_limit"`
Code string `json:"code"`
Upgradable bool `json:"upgradable"`
RequestLimit int `json:"request_limit"`
EmailLimit int `json:"email_limit"`
AttachmentsBytesLimit int64 `json:"attachments_bytes_limit"`
}
type apiAccountUsageLimits struct {
Basis string `json:"basis"` // "ip" or "account"
Messages int `json:"messages"`
Requests int64 `json:"requests"`
Emails int `json:"emails"`
AttachmentsBytes int64 `json:"attachments_bytes"`
}

View file

@ -24,17 +24,18 @@ var (
// visitor represents an API user, and its associated rate.Limiter used for rate limiting
type visitor struct {
config *Config
messageCache *messageCache
ip netip.Addr
user *auth.User
requests *rate.Limiter
emails *rate.Limiter
subscriptions util.Limiter
bandwidth util.Limiter
firebase time.Time // Next allowed Firebase message
seen time.Time
mu sync.Mutex
config *Config
messageCache *messageCache
ip netip.Addr
user *auth.User
requests *util.AtomicCounter[int64]
requestLimiter *rate.Limiter
emails *rate.Limiter
subscriptions util.Limiter
bandwidth util.Limiter
firebase time.Time // Next allowed Firebase message
seen time.Time
mu sync.Mutex
}
type visitorStats struct {
@ -45,28 +46,29 @@ type visitorStats struct {
}
func newVisitor(conf *Config, messageCache *messageCache, ip netip.Addr, user *auth.User) *visitor {
var requests *rate.Limiter
var requestLimiter *rate.Limiter
if user != nil && user.Plan != nil {
requests = rate.NewLimiter(rate.Limit(user.Plan.MessagesLimit)*rate.Every(24*time.Hour), user.Plan.MessagesLimit)
requestLimiter = rate.NewLimiter(rate.Limit(user.Plan.RequestLimit)*rate.Every(24*time.Hour), conf.VisitorRequestLimitBurst)
} else {
requests = rate.NewLimiter(rate.Every(conf.VisitorRequestLimitReplenish), conf.VisitorRequestLimitBurst)
requestLimiter = rate.NewLimiter(rate.Every(conf.VisitorRequestLimitReplenish), conf.VisitorRequestLimitBurst)
}
return &visitor{
config: conf,
messageCache: messageCache,
ip: ip,
user: user,
requests: requests,
emails: rate.NewLimiter(rate.Every(conf.VisitorEmailLimitReplenish), conf.VisitorEmailLimitBurst),
subscriptions: util.NewFixedLimiter(int64(conf.VisitorSubscriptionLimit)),
bandwidth: util.NewBytesLimiter(conf.VisitorAttachmentDailyBandwidthLimit, 24*time.Hour),
firebase: time.Unix(0, 0),
seen: time.Now(),
config: conf,
messageCache: messageCache,
ip: ip,
user: user,
requests: util.NewAtomicCounter[int64](0),
requestLimiter: requestLimiter,
emails: rate.NewLimiter(rate.Every(conf.VisitorEmailLimitReplenish), conf.VisitorEmailLimitBurst),
subscriptions: util.NewFixedLimiter(int64(conf.VisitorSubscriptionLimit)),
bandwidth: util.NewBytesLimiter(conf.VisitorAttachmentDailyBandwidthLimit, 24*time.Hour),
firebase: time.Unix(0, 0),
seen: time.Now(),
}
}
func (v *visitor) RequestAllowed() error {
if !v.requests.Allow() {
if !v.requestLimiter.Allow() {
return errVisitorLimitReached
}
return nil