vbatts/ntfy
1
0
Fork 0
Code Pull requests Activity

Subscription limit

Browse source
This commit is contained in:
Philipp Heckel 2021-11-01 15:21:38 -04:00
parent 5f2bb4f876
commit fa7a45902f
3 changed files with 97 additions and 37 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
config/config.go
View file

@ -17,8 +17,9 @@ const (
// Defines the max number of requests, here:
// 50 requests bucket, replenished at a rate of 1 per second
var (
defaultLimit = rate.Every(time.Second)
defaultLimitBurst = 50
defaultRequestLimit = rate.Every(time.Second)
defaultRequestLimitBurst = 50
defaultSubscriptionLimit = 30 // per visitor
)
// Config is the main config struct for the application. Use New to instantiate a default config struct.
@ -28,8 +29,9 @@ type Config struct {
MessageBufferDuration time.Duration
KeepaliveInterval time.Duration
ManagerInterval time.Duration
Limit rate.Limit
LimitBurst int
RequestLimit rate.Limit
RequestLimitBurst int
SubscriptionLimit int
}
// New instantiates a default new config
@ -40,7 +42,8 @@ func New(listenHTTP string) *Config {
MessageBufferDuration: DefaultMessageBufferDuration,
KeepaliveInterval: DefaultKeepaliveInterval,
ManagerInterval: DefaultManagerInterval,
Limit: defaultLimit,
LimitBurst: defaultLimitBurst,
RequestLimit: defaultRequestLimit,
RequestLimitBurst: defaultRequestLimitBurst,
SubscriptionLimit: defaultSubscriptionLimit,
}
}

54
server/server.go
View file

@ -9,7 +9,6 @@ import (
firebase "firebase.google.com/go"
"firebase.google.com/go/messaging"
"fmt"
"golang.org/x/time/rate"
"google.golang.org/api/option"
"heckel.io/ntfy/config"
"io"
@ -23,9 +22,8 @@ import (
"time"
)
// TODO add "max connections open" limit
// TODO add "max messages in a topic" limit
// TODO add "max topics" limit
// TODO implement persistence
// Server is the main server
type Server struct {
@ -37,12 +35,6 @@ type Server struct {
mu sync.Mutex
}
// visitor represents an API user, and its associated rate.Limiter used for rate limiting
type visitor struct {
limiter *rate.Limiter
seen time.Time
}
// errHTTP is a generic HTTP error for any non-200 HTTP error
type errHTTP struct {
Code int
@ -54,8 +46,7 @@ func (e errHTTP) Error() string {
}
const (
messageLimit = 1024
visitorExpungeAfter = 30 * time.Minute
messageLimit = 1024
)
var (
@ -147,8 +138,8 @@ func (s *Server) handle(w http.ResponseWriter, r *http.Request) {
func (s *Server) handleInternal(w http.ResponseWriter, r *http.Request) error {
v := s.visitor(r.RemoteAddr)
if !v.limiter.Allow() {
return errHTTPTooManyRequests
if err := v.RequestAllowed(); err != nil {
return err
}
if r.Method == http.MethodGet && r.URL.Path == "/" {
return s.handleHome(w, r)
@ -157,11 +148,11 @@ func (s *Server) handleInternal(w http.ResponseWriter, r *http.Request) error {
} else if (r.Method == http.MethodPut || r.Method == http.MethodPost) && topicRegex.MatchString(r.URL.Path) {
return s.handlePublish(w, r)
} else if r.Method == http.MethodGet && jsonRegex.MatchString(r.URL.Path) {
return s.handleSubscribeJSON(w, r)
return s.handleSubscribeJSON(w, r, v)
} else if r.Method == http.MethodGet && sseRegex.MatchString(r.URL.Path) {
return s.handleSubscribeSSE(w, r)
return s.handleSubscribeSSE(w, r, v)
} else if r.Method == http.MethodGet && rawRegex.MatchString(r.URL.Path) {
return s.handleSubscribeRaw(w, r)
return s.handleSubscribeRaw(w, r, v)
} else if r.Method == http.MethodOptions {
return s.handleOptions(w, r)
}
@ -195,7 +186,7 @@ func (s *Server) handlePublish(w http.ResponseWriter, r *http.Request) error {
return nil
}
func (s *Server) handleSubscribeJSON(w http.ResponseWriter, r *http.Request) error {
func (s *Server) handleSubscribeJSON(w http.ResponseWriter, r *http.Request, v *visitor) error {
encoder := func(msg *message) (string, error) {
var buf bytes.Buffer
if err := json.NewEncoder(&buf).Encode(&msg); err != nil {
@ -203,10 +194,10 @@ func (s *Server) handleSubscribeJSON(w http.ResponseWriter, r *http.Request) err
}
return buf.String(), nil
}
return s.handleSubscribe(w, r, "json", "application/stream+json", encoder)
return s.handleSubscribe(w, r, v, "json", "application/stream+json", encoder)
}
func (s *Server) handleSubscribeSSE(w http.ResponseWriter, r *http.Request) error {
func (s *Server) handleSubscribeSSE(w http.ResponseWriter, r *http.Request, v *visitor) error {
encoder := func(msg *message) (string, error) {
var buf bytes.Buffer
if err := json.NewEncoder(&buf).Encode(&msg); err != nil {
@ -217,20 +208,24 @@ func (s *Server) handleSubscribeSSE(w http.ResponseWriter, r *http.Request) erro
}
return fmt.Sprintf("data: %s\n", buf.String()), nil
}
return s.handleSubscribe(w, r, "sse", "text/event-stream", encoder)
return s.handleSubscribe(w, r, v, "sse", "text/event-stream", encoder)
}
func (s *Server) handleSubscribeRaw(w http.ResponseWriter, r *http.Request) error {
func (s *Server) handleSubscribeRaw(w http.ResponseWriter, r *http.Request, v *visitor) error {
encoder := func(msg *message) (string, error) {
if msg.Event == "" { // only handle default events
return strings.ReplaceAll(msg.Message, "\n", " ") + "\n", nil
}
return "\n", nil // "keepalive" and "open" events just send an empty line
}
return s.handleSubscribe(w, r, "raw", "text/plain", encoder)
return s.handleSubscribe(w, r, v, "raw", "text/plain", encoder)
}
func (s *Server) handleSubscribe(w http.ResponseWriter, r *http.Request, format string, contentType string, encoder messageEncoder) error {
func (s *Server) handleSubscribe(w http.ResponseWriter, r *http.Request, v *visitor, format string, contentType string, encoder messageEncoder) error {
if err := v.AddSubscription(); err != nil {
return err
}
defer v.RemoveSubscription()
t := s.createTopic(strings.TrimSuffix(r.URL.Path[1:], "/"+format)) // Hack
since, err := parseSince(r)
if err != nil {
@ -270,6 +265,7 @@ func (s *Server) handleSubscribe(w http.ResponseWriter, r *http.Request, format
case <-r.Context().Done():
return nil
case <-time.After(s.config.KeepaliveInterval):
v.Keepalive()
if err := sub(newKeepaliveMessage(t.id)); err != nil { // Send keepalive message
return err
}
@ -326,12 +322,12 @@ func (s *Server) updateStatsAndExpire() {
// Expire visitors from rate visitors map
for ip, v := range s.visitors {
if time.Since(v.seen) > visitorExpungeAfter {
if v.Stale() {
delete(s.visitors, ip)
}
}
// Prune old messages, remove topics without subscribers
// Prune old messages, remove subscriptions without subscribers
for _, t := range s.topics {
t.Prune(s.config.MessageBufferDuration)
subs, msgs := t.Stats()
@ -362,12 +358,8 @@ func (s *Server) visitor(remoteAddr string) *visitor {
}
v, exists := s.visitors[ip]
if !exists {
v = &visitor{
rate.NewLimiter(s.config.Limit, s.config.LimitBurst),
time.Now(),
}
s.visitors[ip] = v
return v
s.visitors[ip] = newVisitor(s.config)
return s.visitors[ip]
}
v.seen = time.Now()
return v

65
server/visitor.go Normal file
View file

@ -0,0 +1,65 @@
package server
import (
"golang.org/x/time/rate"
"heckel.io/ntfy/config"
"sync"
"time"
)
const (
visitorExpungeAfter = 30 * time.Minute
)
// visitor represents an API user, and its associated rate.Limiter used for rate limiting
type visitor struct {
config *config.Config
limiter *rate.Limiter
subscriptions int
seen time.Time
mu sync.Mutex
}
func newVisitor(conf *config.Config) *visitor {
return &visitor{
config: conf,
limiter: rate.NewLimiter(conf.RequestLimit, conf.RequestLimitBurst),
seen: time.Now(),
}
}
func (v *visitor) RequestAllowed() error {
if !v.limiter.Allow() {
return errHTTPTooManyRequests
}
return nil
}
func (v *visitor) AddSubscription() error {
v.mu.Lock()
defer v.mu.Unlock()
if v.subscriptions >= v.config.SubscriptionLimit {
return errHTTPTooManyRequests
}
v.subscriptions++
return nil
}
func (v *visitor) RemoveSubscription() {
v.mu.Lock()
defer v.mu.Unlock()
v.subscriptions--
}
func (v *visitor) Keepalive() {
v.mu.Lock()
defer v.mu.Unlock()
v.seen = time.Now()
}
func (v *visitor) Stale() bool {
v.mu.Lock()
defer v.mu.Unlock()
v.seen = time.Now()
return time.Since(v.seen) > visitorExpungeAfter
}