1
0
Fork 0
forked from mirrors/ntfy
ntfy/server/visitor.go
Karmanyaah Malhotra c2382d29a1 refactor visitor IPs and allow exempting IP Ranges
Use netip.Addr instead of storing addresses as strings. This requires
conversions at the database level and in tests, but is more memory
efficient otherwise, and facilitates the following.

Parse rate limit exemptions as netip.Prefix. This allows storing IP
ranges in the exemption list. Regular IP addresses (entered explicitly
or resolved from hostnames) are IPV4/32, denoting a range of one
address.
2022-10-05 16:04:42 -05:00

134 lines
3.6 KiB
Go

package server
import (
"errors"
"net/netip"
"sync"
"time"
"golang.org/x/time/rate"
"heckel.io/ntfy/util"
)
const (
// visitorExpungeAfter defines how long a visitor is active before it is removed from memory. This number
// has to be very high to prevent e-mail abuse, but it doesn't really affect the other limits anyway, since
// they are replenished faster (typically).
visitorExpungeAfter = 24 * time.Hour
)
var (
errVisitorLimitReached = errors.New("limit reached")
)
// visitor represents an API user, and its associated rate.Limiter used for rate limiting
type visitor struct {
config *Config
messageCache *messageCache
ip netip.Addr
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
}
type visitorStats struct {
AttachmentFileSizeLimit int64 `json:"attachmentFileSizeLimit"`
VisitorAttachmentBytesTotal int64 `json:"visitorAttachmentBytesTotal"`
VisitorAttachmentBytesUsed int64 `json:"visitorAttachmentBytesUsed"`
VisitorAttachmentBytesRemaining int64 `json:"visitorAttachmentBytesRemaining"`
}
func newVisitor(conf *Config, messageCache *messageCache, ip netip.Addr) *visitor {
return &visitor{
config: conf,
messageCache: messageCache,
ip: ip,
requests: rate.NewLimiter(rate.Every(conf.VisitorRequestLimitReplenish), conf.VisitorRequestLimitBurst),
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() {
return errVisitorLimitReached
}
return nil
}
func (v *visitor) FirebaseAllowed() error {
v.mu.Lock()
defer v.mu.Unlock()
if time.Now().Before(v.firebase) {
return errVisitorLimitReached
}
return nil
}
func (v *visitor) FirebaseTemporarilyDeny() {
v.mu.Lock()
defer v.mu.Unlock()
v.firebase = time.Now().Add(v.config.FirebaseQuotaExceededPenaltyDuration)
}
func (v *visitor) EmailAllowed() error {
if !v.emails.Allow() {
return errVisitorLimitReached
}
return nil
}
func (v *visitor) SubscriptionAllowed() error {
v.mu.Lock()
defer v.mu.Unlock()
if err := v.subscriptions.Allow(1); err != nil {
return errVisitorLimitReached
}
return nil
}
func (v *visitor) RemoveSubscription() {
v.mu.Lock()
defer v.mu.Unlock()
v.subscriptions.Allow(-1)
}
func (v *visitor) Keepalive() {
v.mu.Lock()
defer v.mu.Unlock()
v.seen = time.Now()
}
func (v *visitor) BandwidthLimiter() util.Limiter {
return v.bandwidth
}
func (v *visitor) Stale() bool {
v.mu.Lock()
defer v.mu.Unlock()
return time.Since(v.seen) > visitorExpungeAfter
}
func (v *visitor) Stats() (*visitorStats, error) {
attachmentsBytesUsed, err := v.messageCache.AttachmentBytesUsed(v.ip.String())
if err != nil {
return nil, err
}
attachmentsBytesRemaining := v.config.VisitorAttachmentTotalSizeLimit - attachmentsBytesUsed
if attachmentsBytesRemaining < 0 {
attachmentsBytesRemaining = 0
}
return &visitorStats{
AttachmentFileSizeLimit: v.config.AttachmentFileSizeLimit,
VisitorAttachmentBytesTotal: v.config.VisitorAttachmentTotalSizeLimit,
VisitorAttachmentBytesUsed: attachmentsBytesUsed,
VisitorAttachmentBytesRemaining: attachmentsBytesRemaining,
}, nil
}