1
0
Fork 0
mirror of https://github.com/binwiederhier/ntfy.git synced 2024-12-14 11:47:33 +00: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

@ -64,7 +64,7 @@ type User struct {
Role Role Role Role
Grants []Grant Grants []Grant
Prefs *UserPrefs Prefs *UserPrefs
Plan *UserPlan Plan *Plan
} }
type UserPrefs struct { type UserPrefs struct {
@ -73,9 +73,18 @@ type UserPrefs struct {
Subscriptions []*UserSubscription `json:"subscriptions,omitempty"` Subscriptions []*UserSubscription `json:"subscriptions,omitempty"`
} }
type UserPlan struct { type PlanCode string
Name string `json:"name"`
MessagesLimit int `json:"messages_limit"` const (
PlanUnlimited = PlanCode("unlimited")
PlanDefault = PlanCode("default")
PlanNone = PlanCode("none")
)
type Plan struct {
Code string `json:"name"`
Upgradable bool `json:"upgradable"`
RequestLimit int `json:"request_limit"`
EmailsLimit int `json:"emails_limit"` EmailsLimit int `json:"emails_limit"`
AttachmentBytesLimit int64 `json:"attachment_bytes_limit"` AttachmentBytesLimit int64 `json:"attachment_bytes_limit"`
} }

View file

@ -23,8 +23,8 @@ const (
BEGIN; BEGIN;
CREATE TABLE IF NOT EXISTS plan ( CREATE TABLE IF NOT EXISTS plan (
id INT NOT NULL, id INT NOT NULL,
name TEXT NOT NULL, code TEXT NOT NULL,
messages_limit INT NOT NULL, request_limit INT NOT NULL,
emails_limit INT NOT NULL, emails_limit INT NOT NULL,
attachment_bytes_limit INT NOT NULL, attachment_bytes_limit INT NOT NULL,
PRIMARY KEY (id) PRIMARY KEY (id)
@ -61,13 +61,13 @@ const (
COMMIT; COMMIT;
` `
selectUserByNameQuery = ` selectUserByNameQuery = `
SELECT u.user, u.pass, u.role, u.settings, p.name, p.messages_limit, p.emails_limit, p.attachment_bytes_limit SELECT u.user, u.pass, u.role, u.settings, p.code, p.request_limit, p.emails_limit, p.attachment_bytes_limit
FROM user u FROM user u
LEFT JOIN plan p on p.id = u.plan_id LEFT JOIN plan p on p.id = u.plan_id
WHERE user = ? WHERE user = ?
` `
selectUserByTokenQuery = ` selectUserByTokenQuery = `
SELECT u.user, u.pass, u.role, u.settings, p.name, p.messages_limit, p.emails_limit, p.attachment_bytes_limit SELECT u.user, u.pass, u.role, u.settings, p.code, p.request_limit, p.emails_limit, p.attachment_bytes_limit
FROM user u FROM user u
JOIN user_token t on u.id = t.user_id JOIN user_token t on u.id = t.user_id
LEFT JOIN plan p on p.id = u.plan_id LEFT JOIN plan p on p.id = u.plan_id
@ -324,13 +324,13 @@ func (a *SQLiteAuthManager) userByToken(token string) (*User, error) {
func (a *SQLiteAuthManager) readUser(rows *sql.Rows) (*User, error) { func (a *SQLiteAuthManager) readUser(rows *sql.Rows) (*User, error) {
defer rows.Close() defer rows.Close()
var username, hash, role string var username, hash, role string
var prefs, planName sql.NullString var prefs, planCode sql.NullString
var messagesLimit, emailsLimit sql.NullInt32 var requestLimit, emailLimit sql.NullInt32
var attachmentBytesLimit sql.NullInt64 var attachmentBytesLimit sql.NullInt64
if !rows.Next() { if !rows.Next() {
return nil, ErrNotFound return nil, ErrNotFound
} }
if err := rows.Scan(&username, &hash, &role, &prefs, &planName, &messagesLimit, &emailsLimit, &attachmentBytesLimit); err != nil { if err := rows.Scan(&username, &hash, &role, &prefs, &planCode, &requestLimit, &emailLimit, &attachmentBytesLimit); err != nil {
return nil, err return nil, err
} else if err := rows.Err(); err != nil { } else if err := rows.Err(); err != nil {
return nil, err return nil, err
@ -351,11 +351,12 @@ func (a *SQLiteAuthManager) readUser(rows *sql.Rows) (*User, error) {
return nil, err return nil, err
} }
} }
if planName.Valid { if planCode.Valid {
user.Plan = &UserPlan{ user.Plan = &Plan{
Name: planName.String, Code: planCode.String,
MessagesLimit: int(messagesLimit.Int32), Upgradable: true, // FIXME
EmailsLimit: int(emailsLimit.Int32), RequestLimit: int(requestLimit.Int32),
EmailsLimit: int(emailLimit.Int32),
AttachmentBytesLimit: attachmentBytesLimit.Int64, AttachmentBytesLimit: attachmentBytesLimit.Int64,
} }
} }

View file

@ -545,6 +545,7 @@ func (s *Server) handlePublishWithoutResponse(r *http.Request, v *visitor) (*mes
return nil, err return nil, err
} }
} }
v.requests.Inc()
s.mu.Lock() s.mu.Lock()
s.messages++ s.messages++
s.mu.Unlock() 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 { if v.user.Plan != nil {
response.Usage.Basis = "account" response.Usage.Basis = "account"
response.Plan = &apiAccountSettingsPlan{ response.Plan = &apiAccountSettingsPlan{
Name: v.user.Plan.Name, Code: v.user.Plan.Code,
MessagesLimit: v.user.Plan.MessagesLimit, RequestLimit: v.user.Plan.RequestLimit,
EmailsLimit: v.user.Plan.EmailsLimit, EmailLimit: v.user.Plan.EmailsLimit,
AttachmentsBytesLimit: v.user.Plan.AttachmentBytesLimit, AttachmentsBytesLimit: v.user.Plan.AttachmentBytesLimit,
} }
} else { } else {
if v.user.Role == auth.RoleAdmin { if v.user.Role == auth.RoleAdmin {
response.Usage.Basis = "account" response.Usage.Basis = "account"
response.Plan = &apiAccountSettingsPlan{ response.Plan = &apiAccountSettingsPlan{
Name: "Unlimited", Code: string(auth.PlanUnlimited),
MessagesLimit: 0, RequestLimit: 0,
EmailsLimit: 0, EmailLimit: 0,
AttachmentsBytesLimit: 0, AttachmentsBytesLimit: 0,
} }
} else { } else {
response.Usage.Basis = "ip" response.Usage.Basis = "ip"
response.Plan = &apiAccountSettingsPlan{ response.Plan = &apiAccountSettingsPlan{
Name: "Free", Code: string(auth.PlanDefault),
MessagesLimit: s.config.VisitorRequestLimitBurst, RequestLimit: s.config.VisitorRequestLimitBurst,
EmailsLimit: s.config.VisitorEmailLimitBurst, EmailLimit: s.config.VisitorEmailLimitBurst,
AttachmentsBytesLimit: s.config.VisitorAttachmentTotalSizeLimit, 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.Role = string(auth.RoleAnonymous)
response.Usage.Basis = "account" response.Usage.Basis = "account"
response.Plan = &apiAccountSettingsPlan{ response.Plan = &apiAccountSettingsPlan{
Name: "Anonymous", Code: string(auth.PlanNone),
MessagesLimit: s.config.VisitorRequestLimitBurst, RequestLimit: s.config.VisitorRequestLimitBurst,
EmailsLimit: s.config.VisitorEmailLimitBurst, EmailLimit: s.config.VisitorEmailLimitBurst,
AttachmentsBytesLimit: s.config.VisitorAttachmentTotalSizeLimit, AttachmentsBytesLimit: s.config.VisitorAttachmentTotalSizeLimit,
} }
} }
response.Usage.Messages = int(v.requests.Tokens()) response.Usage.Requests = v.requests.Value()
response.Usage.AttachmentsBytes = stats.VisitorAttachmentBytesUsed response.Usage.AttachmentsBytes = stats.VisitorAttachmentBytesUsed
if err := json.NewEncoder(w).Encode(response); err != nil { if err := json.NewEncoder(w).Encode(response); err != nil {
return err return err

View file

@ -225,15 +225,16 @@ type apiAccountTokenResponse struct {
} }
type apiAccountSettingsPlan struct { type apiAccountSettingsPlan struct {
Name string `json:"name"` Code string `json:"code"`
MessagesLimit int `json:"messages_limit"` Upgradable bool `json:"upgradable"`
EmailsLimit int `json:"emails_limit"` RequestLimit int `json:"request_limit"`
EmailLimit int `json:"email_limit"`
AttachmentsBytesLimit int64 `json:"attachments_bytes_limit"` AttachmentsBytesLimit int64 `json:"attachments_bytes_limit"`
} }
type apiAccountUsageLimits struct { type apiAccountUsageLimits struct {
Basis string `json:"basis"` // "ip" or "account" Basis string `json:"basis"` // "ip" or "account"
Messages int `json:"messages"` Requests int64 `json:"requests"`
Emails int `json:"emails"` Emails int `json:"emails"`
AttachmentsBytes int64 `json:"attachments_bytes"` 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 // visitor represents an API user, and its associated rate.Limiter used for rate limiting
type visitor struct { type visitor struct {
config *Config config *Config
messageCache *messageCache messageCache *messageCache
ip netip.Addr ip netip.Addr
user *auth.User user *auth.User
requests *rate.Limiter requests *util.AtomicCounter[int64]
emails *rate.Limiter requestLimiter *rate.Limiter
subscriptions util.Limiter emails *rate.Limiter
bandwidth util.Limiter subscriptions util.Limiter
firebase time.Time // Next allowed Firebase message bandwidth util.Limiter
seen time.Time firebase time.Time // Next allowed Firebase message
mu sync.Mutex seen time.Time
mu sync.Mutex
} }
type visitorStats struct { type visitorStats struct {
@ -45,28 +46,29 @@ type visitorStats struct {
} }
func newVisitor(conf *Config, messageCache *messageCache, ip netip.Addr, user *auth.User) *visitor { 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 { 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 { } else {
requests = rate.NewLimiter(rate.Every(conf.VisitorRequestLimitReplenish), conf.VisitorRequestLimitBurst) requestLimiter = rate.NewLimiter(rate.Every(conf.VisitorRequestLimitReplenish), conf.VisitorRequestLimitBurst)
} }
return &visitor{ return &visitor{
config: conf, config: conf,
messageCache: messageCache, messageCache: messageCache,
ip: ip, ip: ip,
user: user, user: user,
requests: requests, requests: util.NewAtomicCounter[int64](0),
emails: rate.NewLimiter(rate.Every(conf.VisitorEmailLimitReplenish), conf.VisitorEmailLimitBurst), requestLimiter: requestLimiter,
subscriptions: util.NewFixedLimiter(int64(conf.VisitorSubscriptionLimit)), emails: rate.NewLimiter(rate.Every(conf.VisitorEmailLimitReplenish), conf.VisitorEmailLimitBurst),
bandwidth: util.NewBytesLimiter(conf.VisitorAttachmentDailyBandwidthLimit, 24*time.Hour), subscriptions: util.NewFixedLimiter(int64(conf.VisitorSubscriptionLimit)),
firebase: time.Unix(0, 0), bandwidth: util.NewBytesLimiter(conf.VisitorAttachmentDailyBandwidthLimit, 24*time.Hour),
seen: time.Now(), firebase: time.Unix(0, 0),
seen: time.Now(),
} }
} }
func (v *visitor) RequestAllowed() error { func (v *visitor) RequestAllowed() error {
if !v.requests.Allow() { if !v.requestLimiter.Allow() {
return errVisitorLimitReached return errVisitorLimitReached
} }
return nil return nil

32
util/atomic_counter.go Normal file
View file

@ -0,0 +1,32 @@
package util
import "sync"
type AtomicCounter[T int | int32 | int64] struct {
value T
mu sync.Mutex
}
func NewAtomicCounter[T int | int32 | int64](value T) *AtomicCounter[T] {
return &AtomicCounter[T]{
value: value,
}
}
func (c *AtomicCounter[T]) Inc() T {
c.mu.Lock()
defer c.mu.Unlock()
c.value++
return c.value
}
func (c *AtomicCounter[T]) Value() T {
c.mu.Lock()
defer c.mu.Unlock()
return c.value
}
func (c *AtomicCounter[T]) Reset() {
c.mu.Lock()
defer c.mu.Unlock()
c.value = 0
}

View file

@ -141,6 +141,10 @@
"subscribe_dialog_login_button_login": "Login", "subscribe_dialog_login_button_login": "Login",
"subscribe_dialog_error_user_not_authorized": "User {{username}} not authorized", "subscribe_dialog_error_user_not_authorized": "User {{username}} not authorized",
"subscribe_dialog_error_user_anonymous": "anonymous", "subscribe_dialog_error_user_anonymous": "anonymous",
"account_type_default": "Default",
"account_type_unlimited": "Unlimited",
"account_type_none": "None",
"account_type_hobbyist": "Hobbyist",
"prefs_notifications_title": "Notifications", "prefs_notifications_title": "Notifications",
"prefs_notifications_sound_title": "Notification sound", "prefs_notifications_sound_title": "Notification sound",
"prefs_notifications_sound_description_none": "Notifications do not play any sound when they arrive", "prefs_notifications_sound_description_none": "Notifications do not play any sound when they arrive",

View file

@ -53,7 +53,9 @@ const Stats = () => {
const { t } = useTranslation(); const { t } = useTranslation();
const { account } = useOutletContext(); const { account } = useOutletContext();
const admin = account?.role === "admin" const admin = account?.role === "admin"
const accountType = account?.plan?.name ?? "Free"; const usage = account?.usage;
const plan = account?.plan;
const accountType = plan?.code ?? "none";
return ( return (
<Card sx={{p: 3}} aria-label={t("xxxxxxxxx")}> <Card sx={{p: 3}} aria-label={t("xxxxxxxxx")}>
<Typography variant="h5" sx={{marginBottom: 2}}> <Typography variant="h5" sx={{marginBottom: 2}}>
@ -64,13 +66,13 @@ const Stats = () => {
<div> <div>
{account?.role === "admin" {account?.role === "admin"
? <>Unlimited <Tooltip title={"You are Admin"}><span style={{cursor: "default"}}>👑</span></Tooltip></> ? <>Unlimited <Tooltip title={"You are Admin"}><span style={{cursor: "default"}}>👑</span></Tooltip></>
: accountType} : t(`account_type_${accountType}`)}
</div> </div>
</Pref> </Pref>
<Pref labelId={"dailyMessages"} title={t("Daily messages")}> <Pref labelId={"dailyMessages"} title={t("Daily messages")}>
<div> <div>
<Typography variant="body2" sx={{float: "left"}}>123</Typography> <Typography variant="body2" sx={{float: "left"}}>{usage?.requests ?? 0}</Typography>
<Typography variant="body2" sx={{float: "right"}}>of 1000</Typography> <Typography variant="body2" sx={{float: "right"}}>{plan?.request_limit > 0 ? t("of {{limit}}", { limit: plan.request_limit }) : t("Unlimited")}</Typography>
</div> </div>
<LinearProgress variant="determinate" value={10} /> <LinearProgress variant="determinate" value={10} />
</Pref> </Pref>