very very rought draft of reciever rate limiting

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 {

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
 	}