1
0
Fork 0
forked from mirrors/ntfy

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
Grants []Grant
Prefs *UserPrefs
Plan *UserPlan
Plan *Plan
}
type UserPrefs struct {
@ -73,9 +73,18 @@ type UserPrefs struct {
Subscriptions []*UserSubscription `json:"subscriptions,omitempty"`
}
type UserPlan struct {
Name string `json:"name"`
MessagesLimit int `json:"messages_limit"`
type PlanCode string
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"`
AttachmentBytesLimit int64 `json:"attachment_bytes_limit"`
}

View file

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

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

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_error_user_not_authorized": "User {{username}} not authorized",
"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_sound_title": "Notification sound",
"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 { account } = useOutletContext();
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 (
<Card sx={{p: 3}} aria-label={t("xxxxxxxxx")}>
<Typography variant="h5" sx={{marginBottom: 2}}>
@ -64,13 +66,13 @@ const Stats = () => {
<div>
{account?.role === "admin"
? <>Unlimited <Tooltip title={"You are Admin"}><span style={{cursor: "default"}}>👑</span></Tooltip></>
: accountType}
: t(`account_type_${accountType}`)}
</div>
</Pref>
<Pref labelId={"dailyMessages"} title={t("Daily messages")}>
<div>
<Typography variant="body2" sx={{float: "left"}}>123</Typography>
<Typography variant="body2" sx={{float: "right"}}>of 1000</Typography>
<Typography variant="body2" sx={{float: "left"}}>{usage?.requests ?? 0}</Typography>
<Typography variant="body2" sx={{float: "right"}}>{plan?.request_limit > 0 ? t("of {{limit}}", { limit: plan.request_limit }) : t("Unlimited")}</Typography>
</div>
<LinearProgress variant="determinate" value={10} />
</Pref>