From 84e2ee9b7c80dfb753bf54764de3c12f2e3f1dae Mon Sep 17 00:00:00 2001 From: Karmanyaah Malhotra Date: Mon, 16 Jan 2023 16:09:10 -0600 Subject: [PATCH] very very rought draft of reciever rate limiting --- server/server.go | 37 ++++++++++++++++++++++++++++--------- server/topic.go | 48 ++++++++++++++++++++++++++++++++++++++++-------- 2 files changed, 68 insertions(+), 17 deletions(-) diff --git a/server/server.go b/server/server.go index 64891313..18bdae23 100644 --- a/server/server.go +++ b/server/server.go @@ -137,12 +137,14 @@ var ( ) const ( - firebaseControlTopic = "~control" // See Android if changed - firebasePollTopic = "~poll" // See iOS if changed - emptyMessageBody = "triggered" // Used if message body is empty - newMessageBody = "New message" // Used in poll requests as generic message - defaultAttachmentMessage = "You received a file: %s" // Used if message body is empty, and there is an attachment - encodingBase64 = "base64" + firebaseControlTopic = "~control" // See Android if changed + firebasePollTopic = "~poll" // See iOS if changed + emptyMessageBody = "triggered" // Used if message body is empty + newMessageBody = "New message" // Used in poll requests as generic message + defaultAttachmentMessage = "You received a file: %s" // Used if message body is empty, and there is an attachment + encodingBase64 = "base64" + unifiedpushTopicPrefix = "up" + unifiedPushSubscriptionDuration = 12 * time.Hour ) // WebSocket constants @@ -544,6 +546,17 @@ func (s *Server) handlePublishWithoutResponse(r *http.Request, v *visitor) (*mes if err != nil { return nil, err } + + var v_billing *visitor + if strings.HasPrefix(t.ID, unifiedpushTopicPrefix) { + v_billing := t.getBillee() + if v_billing != nil { + // instant reject and won't even store it if there's no one registered for a UP topic in the past some time + // need to find error code for device not available try again later + return nil, errHTTPInternalError + } + } + if err := v.MessageAllowed(); err != nil { return nil, errHTTPTooManyRequestsLimitMessages } @@ -569,6 +582,7 @@ func (s *Server) handlePublishWithoutResponse(r *http.Request, v *visitor) (*mes if m.Message == "" { m.Message = emptyMessageBody } + // we do not need to handle delays, because 1. for UP delays are not needed 2. if the up app server is adding a delay it is shooting itself in the foot delayed := m.Time > time.Now().Unix() log.Debug("%s Received message: event=%s, user=%s, body=%d byte(s), delayed=%t, firebase=%t, cache=%t, up=%t, email=%s", logMessagePrefix(v, m), m.Event, m.User, len(m.Message), delayed, firebase, cache, unifiedpush, email) @@ -582,6 +596,7 @@ func (s *Server) handlePublishWithoutResponse(r *http.Request, v *visitor) (*mes if s.firebaseClient != nil && firebase { go s.sendToFirebase(v, m) } + // same as delays, it should count against app servers if s.smtpSender != nil && email != "" { v.IncrementEmails() go s.sendEmail(v, m, email) @@ -598,7 +613,11 @@ func (s *Server) handlePublishWithoutResponse(r *http.Request, v *visitor) (*mes return nil, err } } - v.IncrementMessages() + if v_billing != nil { + v_billing.IncrementMessages() + } else { + v.IncrementMessages() + } if s.userManager != nil && v.user != nil { s.userManager.EnqueueStats(v.user) } @@ -961,7 +980,7 @@ func (s *Server) handleSubscribeHTTP(w http.ResponseWriter, r *http.Request, v * } subscriberIDs := make([]int, 0) for _, t := range topics { - subscriberIDs = append(subscriberIDs, t.Subscribe(sub)) + subscriberIDs = append(subscriberIDs, t.Subscribe(sub, v)) } defer func() { for i, subscriberID := range subscriberIDs { @@ -1076,7 +1095,7 @@ func (s *Server) handleSubscribeWS(w http.ResponseWriter, r *http.Request, v *vi } subscriberIDs := make([]int, 0) for _, t := range topics { - subscriberIDs = append(subscriberIDs, t.Subscribe(sub)) + subscriberIDs = append(subscriberIDs, t.Subscribe(sub, v)) } defer func() { for i, subscriberID := range subscriberIDs { diff --git a/server/topic.go b/server/topic.go index 3bc74736..3f9288d9 100644 --- a/server/topic.go +++ b/server/topic.go @@ -1,36 +1,64 @@ package server import ( - "heckel.io/ntfy/log" "math/rand" "sync" + "time" + + "heckel.io/ntfy/log" ) // topic represents a channel to which subscribers can subscribe, and publishers // can publish a message type topic struct { ID string - subscribers map[int]subscriber + subscribers map[int]topicSubscription + lastUnsub topicSubscription mu sync.Mutex } // subscriber is a function that is called for every new message on a topic +type topicSubscription struct { + f subscriber + v *visitor + unsubTime time.Time +} type subscriber func(v *visitor, msg *message) error // newTopic creates a new topic func newTopic(id string) *topic { return &topic{ ID: id, - subscribers: make(map[int]subscriber), + subscribers: make(map[int]topicSubscription), } } +// need a better name for bill? +// Returns nil, nil for non-UP topics +// returns visitor, nil for active UP topics +// returns nil, err for inactive UP topics +func (t *topic) getBillee() *visitor { + //get a pseudo random visitor??? + for _, this_subscriber := range t.subscribers { + return this_subscriber.v + } + // what if someone unsubscribed and DOESNT want their sub to count against them anymore, maybe the app server lost sync and will keep on sending stuff + // I guess they suffer for unifiedPushSubscriptionDuration? + + // if lastunsub v exists, and the time since it was unsubbed is longer than our limit, it should not exist + if t.lastUnsub.v != nil && time.Since(t.lastUnsub.unsubTime) > unifiedPushSubscriptionDuration { + t.lastUnsub.v = nil + } + + return t.lastUnsub.v +} + // Subscribe subscribes to this topic -func (t *topic) Subscribe(s subscriber) int { +func (t *topic) Subscribe(s subscriber, v *visitor) int { t.mu.Lock() defer t.mu.Unlock() subscriberID := rand.Int() - t.subscribers[subscriberID] = s + t.subscribers[subscriberID] = topicSubscription{f: s, v: v} return subscriberID } @@ -38,6 +66,10 @@ func (t *topic) Subscribe(s subscriber) int { func (t *topic) Unsubscribe(id int) { t.mu.Lock() defer t.mu.Unlock() + if len(t.subscribers) == 1 { + t.lastUnsub = t.subscribers[id] + t.lastUnsub.unsubTime = time.Now() + } delete(t.subscribers, id) } @@ -56,7 +88,7 @@ func (t *topic) Publish(v *visitor, m *message) error { if err := s(v, m); err != nil { log.Warn("%s Error forwarding to subscriber", logMessagePrefix(v, m)) } - }(s) + }(s.f) } } else { log.Trace("%s No stream or WebSocket subscribers, not forwarding", logMessagePrefix(v, m)) @@ -73,10 +105,10 @@ func (t *topic) SubscribersCount() int { } // subscribersCopy returns a shallow copy of the subscribers map -func (t *topic) subscribersCopy() map[int]subscriber { +func (t *topic) subscribersCopy() map[int]topicSubscription { t.mu.Lock() defer t.mu.Unlock() - subscribers := make(map[int]subscriber) + subscribers := make(map[int]topicSubscription) for k, v := range t.subscribers { subscribers[k] = v }