very very rought draft of reciever rate limiting

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
 	}