diff --git a/limit/rate_visitor.go b/limit/rate_visitor.go deleted file mode 100644 index 9c386920..00000000 --- a/limit/rate_visitor.go +++ /dev/null @@ -1,169 +0,0 @@ -package limit - -import ( - "context" - "time" - - "heckel.io/ntfy/v2/log" -) - -const ( - // rateVisitorTouchInterval is how often a touch of a topic's rate visitor is written at most; an open - // stream is touched about this often, so other nodes may see it up to this much older - rateVisitorTouchInterval = 5 * time.Minute - - // rateVisitorPruneAfter is how long a rate visitor row is kept after it was last written. It must be well - // above the server's 404 threshold for an abandoned UnifiedPush topic (7 days), so the row - // that starts the countdown is still there when the threshold is reached. - rateVisitorPruneAfter = 30 * oneDay -) - -// SetRateVisitor makes p the topic's rate visitor; the latest subscriber wins. p.SeenAt is the node's clock. -func (s *Service) SetRateVisitor(topic string, p *RateVisitor) { - s.rateVisitors.mu.Lock() - defer s.rateVisitors.mu.Unlock() - s.rateVisitors.entries[topic] = &rateVisitorEntry{ - key: p.Key, - userID: p.UserID, - seenAt: p.SeenAt.Unix(), - write: rateVisitorWriteSet, - } -} - -// TouchRateVisitor marks subscriber activity on the topic now (an open or closed stream), keeping its -// rate visitor. It does nothing if the topic was seen within rateVisitorTouchInterval. -func (s *Service) TouchRateVisitor(topic string) { - now := s.now().Unix() - s.rateVisitors.mu.Lock() - defer s.rateVisitors.mu.Unlock() - e, ok := s.rateVisitors.entries[topic] - if !ok { - s.rateVisitors.addNoLock(topic, now) - return - } else if now-e.seenAt < int64(rateVisitorTouchInterval/time.Second) { - return - } - e.seenAt = now - e.write = max(e.write, rateVisitorWriteTouch) -} - -// NoteTopic remembers a topic this node does not know yet, as seen now and without a rate visitor, e.g. -// on a refused push, so that the topic's abandonment can be told apart from a new topic. It is -// written as a touch: a row this node has not read yet (written in the last seconds, or while its -// load failed) keeps its rate visitor, and its seen_at moves to now at most. -func (s *Service) NoteTopic(topic string) { - now := s.now().Unix() - s.rateVisitors.mu.Lock() - defer s.rateVisitors.mu.Unlock() - if _, ok := s.rateVisitors.entries[topic]; !ok { - s.rateVisitors.addNoLock(topic, now) - } -} - -// RateVisitor returns the topic's rate visitor, or nil if the topic is unknown -func (s *Service) RateVisitor(topic string) *RateVisitor { - s.rateVisitors.mu.Lock() - defer s.rateVisitors.mu.Unlock() - e, ok := s.rateVisitors.entries[topic] - if !ok { - return nil - } - return &RateVisitor{ - Key: e.key, - UserID: e.userID, - SeenAt: time.Unix(e.seenAt, 0).UTC(), - } -} - -// writeRateVisitors writes the rate visitors changed here. A failed write marks them changed again, for the -// next round; rewriting the ones that were written is harmless. -func (s *Service) writeRateVisitors(ctx context.Context) { - rows := s.unwrittenRateVisitors() - if len(rows) == 0 { - return - } - if err := s.store.writeRateVisitors(ctx, rows, s.now().Unix()); err != nil { - s.markRateVisitorsUnwritten(rows) - log.Tag(tag).Err(err).Warn("Cannot write %d rate visitor(s); retrying in the next round", len(rows)) - } -} - -// unwrittenRateVisitors returns the rows of the rate visitors changed here, and marks them written -func (s *Service) unwrittenRateVisitors() []rateVisitorRow { - s.rateVisitors.mu.Lock() - defer s.rateVisitors.mu.Unlock() - rows := make([]rateVisitorRow, 0) - for topic, e := range s.rateVisitors.entries { - if e.write == rateVisitorWriteNone { - continue - } - r := rateVisitorRow{ - topic: topic, - seenAt: e.seenAt, - } - if e.write == rateVisitorWriteSet { - r.key, r.userID = e.key, e.userID - } - rows = append(rows, r) - e.write = rateVisitorWriteNone - } - return rows -} - -// markRateVisitorsUnwritten marks the rate visitors of rows that failed to write as changed again -func (s *Service) markRateVisitorsUnwritten(rows []rateVisitorRow) { - s.rateVisitors.mu.Lock() - defer s.rateVisitors.mu.Unlock() - for _, r := range rows { - e, ok := s.rateVisitors.entries[r.topic] - if !ok { - continue - } - write := rateVisitorWriteTouch - if r.key != "" { - write = rateVisitorWriteSet - } - e.write = max(e.write, write) - } -} - -// readRateVisitors applies the rate visitor rows written since the given time, except to entries with a change -// not written yet (that change is newer). It returns how many rows it applied. -func (s *Service) readRateVisitors(ctx context.Context, since int64) (int, error) { - applied := 0 - err := s.store.readRateVisitorsSince(ctx, since, func(r *rateVisitorRow) { - s.rateVisitors.mu.Lock() - defer s.rateVisitors.mu.Unlock() - e, ok := s.rateVisitors.entries[r.topic] - if ok && e.write != rateVisitorWriteNone { - return - } else if !ok { - e = &rateVisitorEntry{} - s.rateVisitors.entries[r.topic] = e - } - e.key, e.userID, e.seenAt = r.key, r.userID, r.seenAt - applied++ - }) - return applied, err -} - -// pruneRateVisitors forgets the rate visitors not seen for rateVisitorPruneAfter -func (s *Service) pruneRateVisitors(now time.Time) { - cutoff := now.Add(-rateVisitorPruneAfter).Unix() - s.rateVisitors.mu.Lock() - defer s.rateVisitors.mu.Unlock() - for topic, e := range s.rateVisitors.entries { - if e.seenAt < cutoff { - delete(s.rateVisitors.entries, topic) - } - } -} - -// addNoLock adds the topic as seen at now without a rate visitor, to be written as a touch. The caller -// must hold pt.mu. -func (pt *rateVisitorTable) addNoLock(topic string, now int64) { - pt.entries[topic] = &rateVisitorEntry{ - seenAt: now, - write: rateVisitorWriteTouch, - } -} diff --git a/limit/rate_visitor_test.go b/limit/rate_visitor_test.go deleted file mode 100644 index f977db72..00000000 --- a/limit/rate_visitor_test.go +++ /dev/null @@ -1,183 +0,0 @@ -package limit - -import ( - "context" - "errors" - "testing" - "time" - - "github.com/stretchr/testify/require" -) - -func TestService_RateVisitor_SetTouchNote(t *testing.T) { - s, clock := newTestServiceWithClock(t) - require.Nil(t, s.RateVisitor("upAAAAAAAAAAAA")) - - // A refused push makes the topic known, without a rate visitor; a second one does not move its time - s.NoteTopic("upAAAAAAAAAAAA") - noted := clock.now() - clock.advance(time.Hour) - s.NoteTopic("upAAAAAAAAAAAA") - require.Equal(t, &RateVisitor{ - SeenAt: noted, - }, s.RateVisitor("upAAAAAAAAAAAA")) - - // A subscriber becomes the rate visitor; the latest one wins - s.SetRateVisitor("upAAAAAAAAAAAA", &RateVisitor{ - Key: "ip:1.2.3.4", - SeenAt: clock.now(), - }) - s.SetRateVisitor("upAAAAAAAAAAAA", &RateVisitor{ - Key: "user:u_1", - UserID: "u_1", - SeenAt: clock.now(), - }) - set := clock.now() - require.Equal(t, userRateVisitor(set), s.RateVisitor("upAAAAAAAAAAAA")) - - // A touch keeps the rate visitor, and moves its time only once rateVisitorTouchInterval has passed - clock.advance(rateVisitorTouchInterval - time.Second) - s.TouchRateVisitor("upAAAAAAAAAAAA") - require.Equal(t, userRateVisitor(set), s.RateVisitor("upAAAAAAAAAAAA")) - clock.advance(time.Second) - s.TouchRateVisitor("upAAAAAAAAAAAA") - require.Equal(t, userRateVisitor(clock.now()), s.RateVisitor("upAAAAAAAAAAAA")) - - // A touch of an unknown topic makes it known, without a rate visitor - s.TouchRateVisitor("upBBBBBBBBBBBB") - require.Equal(t, &RateVisitor{ - SeenAt: clock.now(), - }, s.RateVisitor("upBBBBBBBBBBBB")) -} - -func TestService_RateVisitor_Prune(t *testing.T) { - s, clock := newTestServiceWithClock(t) - s.NoteTopic("upAAAAAAAAAAAA") - clock.advance(time.Hour) - s.NoteTopic("upBBBBBBBBBBBB") - clock.advance(rateVisitorPruneAfter - 30*time.Minute) - s.Prune(time.Hour) - require.Nil(t, s.RateVisitor("upAAAAAAAAAAAA")) - require.NotNil(t, s.RateVisitor("upBBBBBBBBBBBB")) -} - -func TestService_Persist_RateVisitorSharedAcrossNodes(t *testing.T) { - // A rate visitor set on one node is the rate visitor on the other after a sync; a touch on the other node - // moves its time for both, and keeps the rate visitor - forEachBackend(t, func(t *testing.T, newStore func() *Store) { - a := startTestService(t, newStore(), time.Now) - b := startTestService(t, newStore(), time.Now) - aged := time.Unix(time.Now().Add(-25*time.Hour).Unix(), 0).UTC() - a.SetRateVisitor("upAAAAAAAAAAAA", &RateVisitor{ - Key: "user:u_1", - UserID: "u_1", - SeenAt: aged, - }) - a.sync() - b.sync() - require.Equal(t, userRateVisitor(aged), b.RateVisitor("upAAAAAAAAAAAA")) - b.TouchRateVisitor("upAAAAAAAAAAAA") - require.NotNil(t, b.RateVisitor("upAAAAAAAAAAAA")) - touched := b.RateVisitor("upAAAAAAAAAAAA").SeenAt - require.True(t, touched.After(aged)) - b.sync() - a.sync() - require.Equal(t, userRateVisitor(touched), a.RateVisitor("upAAAAAAAAAAAA")) - }) -} - -func TestService_Persist_RateVisitorSurvivesRestart(t *testing.T) { - forEachBackend(t, func(t *testing.T, newStore func() *Store) { - s := startTestService(t, newStore(), time.Now) - seen := time.Unix(time.Now().Unix(), 0).UTC() - s.SetRateVisitor("upAAAAAAAAAAAA", &RateVisitor{ - Key: "ip:1.2.3.4", - SeenAt: seen, - }) - s.Close() // Writes what is unwritten - s = startTestService(t, newStore(), time.Now) - require.Equal(t, ipRateVisitor(seen), s.RateVisitor("upAAAAAAAAAAAA")) - }) -} - -func TestService_Persist_NoteKeepsStoredRateVisitor(t *testing.T) { - // A refused push on a node that has not read the rate visitor yet must not wipe it - forEachBackend(t, func(t *testing.T, newStore func() *Store) { - a := startTestService(t, newStore(), time.Now) - b := startTestService(t, newStore(), time.Now) - a.SetRateVisitor("upAAAAAAAAAAAA", &RateVisitor{ - Key: "ip:1.2.3.4", - SeenAt: time.Now(), - }) - a.sync() - b.NoteTopic("upAAAAAAAAAAAA") - b.sync() - a.sync() - for _, s := range []*Service{a, b} { - require.NotNil(t, s.RateVisitor("upAAAAAAAAAAAA")) - require.Equal(t, "ip:1.2.3.4", s.RateVisitor("upAAAAAAAAAAAA").Key) - } - }) -} - -func TestService_Persist_ReadBackKeepsUnwrittenRateVisitor(t *testing.T) { - // A rate visitor set here and not written yet is not overwritten by an older row read back - forEachBackend(t, func(t *testing.T, newStore func() *Store) { - a := startTestService(t, newStore(), time.Now) - b := startTestService(t, newStore(), time.Now) - a.SetRateVisitor("upAAAAAAAAAAAA", &RateVisitor{ - Key: "ip:1.2.3.4", - SeenAt: time.Now(), - }) - a.sync() - b.SetRateVisitor("upAAAAAAAAAAAA", &RateVisitor{ - Key: "ip:5.6.7.8", - SeenAt: time.Now(), - }) - _, err := b.readRateVisitors(context.Background(), 0) - require.Nil(t, err) - require.NotNil(t, b.RateVisitor("upAAAAAAAAAAAA")) - require.Equal(t, "ip:5.6.7.8", b.RateVisitor("upAAAAAAAAAAAA").Key) - b.sync() - a.sync() - require.NotNil(t, a.RateVisitor("upAAAAAAAAAAAA")) - require.Equal(t, "ip:5.6.7.8", a.RateVisitor("upAAAAAAAAAAAA").Key) - }) -} - -func TestService_Persist_RateVisitorWriteFailureRetries(t *testing.T) { - forEachBackend(t, func(t *testing.T, newStore func() *Store) { - store := newStore() - s := startTestService(t, store, time.Now) - s.SetRateVisitor("upAAAAAAAAAAAA", &RateVisitor{ - Key: "ip:1.2.3.4", - SeenAt: time.Now(), - }) - writeRateVisitorChunk := store.writeRateVisitorChunk - store.writeRateVisitorChunk = func(ctx context.Context, rows []rateVisitorRow, now int64) error { - return errors.New("database is down") - } - s.sync() - require.Empty(t, readAllRateVisitors(t, store, 0)) - store.writeRateVisitorChunk = writeRateVisitorChunk - s.sync() - require.Equal(t, "ip:1.2.3.4", readAllRateVisitors(t, store, 0)["upAAAAAAAAAAAA"].key) - }) -} - -// userRateVisitor returns the rate visitor of user u_1, seen at the given time -func userRateVisitor(seen time.Time) *RateVisitor { - return &RateVisitor{ - Key: "user:u_1", - UserID: "u_1", - SeenAt: seen, - } -} - -// ipRateVisitor returns the rate visitor of the anonymous visitor 1.2.3.4, seen at the given time -func ipRateVisitor(seen time.Time) *RateVisitor { - return &RateVisitor{ - Key: "ip:1.2.3.4", - SeenAt: seen, - } -} diff --git a/limit/service.go b/limit/service.go index 127206bb..f1ee6f9b 100644 --- a/limit/service.go +++ b/limit/service.go @@ -15,9 +15,6 @@ // current rows. Store failures never reach the request path: counts stay unflushed and are // written once the store is back (fail open, per-node limits meanwhile). // -// The service also keeps the rate visitors of the UnifiedPush topics (see SetRateVisitor), synced through the -// topic table by the same worker, so a push on any node is billed to the topic's subscriber. -// // Periods, updated_at and the read-back watermark use the node's clock, so the nodes must // keep their clocks in sync (NTP); a skew below a second or so is harmless, a skew beyond // readOverlap may skip rows in the read-back. The database's clock is never used. @@ -63,19 +60,18 @@ const ( // Service decides on usage from memory; see the package comment type Service struct { - conf *Config - now func() time.Time // Injectable clock, for tests - counters [shardCount]*counterShard // Daily counts, kept until their period ends - buckets [shardCount]*bucketShard // Token buckets, per node, evicted when idle and full - rateVisitors *rateVisitorTable // Rate visitors of the UnifiedPush topics, kept for rateVisitorPruneAfter - store *Store // Nil: memory only - watermark int64 // Time of the last read, unix seconds; the sync worker's only - prunedAt int64 // Time of the last prune, unix seconds; the sync worker's only - interval time.Duration // How often the worker syncs (syncInterval); tests change it - ctx context.Context // Done on Close; the sync worker's store calls derive from it, so Close never waits on a stuck one - cancel context.CancelFunc - closeOnce sync.Once - wg sync.WaitGroup + conf *Config + now func() time.Time // Injectable clock, for tests + counters [shardCount]*counterShard // Daily counts, kept until their period ends + buckets [shardCount]*bucketShard // Token buckets, per node, evicted when idle and full + store *Store // Nil: memory only + watermark int64 // Time of the last read, unix seconds; the sync worker's only + prunedAt int64 // Time of the last prune, unix seconds; the sync worker's only + interval time.Duration // How often the worker syncs (syncInterval); tests change it + ctx context.Context // Done on Close; the sync worker's store calls derive from it, so Close never waits on a stuck one + cancel context.CancelFunc + closeOnce sync.Once + wg sync.WaitGroup } // New creates the service @@ -89,9 +85,6 @@ func newService(conf *Config, now func() time.Time) *Service { now: now, store: conf.Store, interval: syncInterval, - rateVisitors: &rateVisitorTable{ - entries: make(map[string]*rateVisitorEntry), - }, } s.ctx, s.cancel = context.WithCancel(context.Background()) for i := 0; i < shardCount; i++ { @@ -172,7 +165,6 @@ func (s *Service) Close() { ctx, cancel := context.WithTimeout(context.Background(), statementTimeout) defer cancel() s.write(ctx) - s.writeRateVisitors(ctx) }) } @@ -271,7 +263,6 @@ func (s *Service) Prune(idle time.Duration) int { } cs.mu.Unlock() } - s.pruneRateVisitors(now) return pruned } @@ -353,38 +344,24 @@ func (s *Service) load() bool { log.Tag(tag).Err(err).Warn("Cannot load usage; retrying in %v", s.interval) return false } - rateVisitors, err := s.readRateVisitors(ctx, now-int64(rateVisitorPruneAfter/time.Second)) - if err != nil { - log.Tag(tag).Err(err).Warn("Cannot load rate visitors; retrying in %v", s.interval) - return false - } s.watermark, s.prunedAt = now, now - log.Tag(tag).Info("Loaded usage of %d key(s) and %d rate visitor(s) in %v", n, rateVisitors, time.Since(started).Round(time.Millisecond)) + log.Tag(tag).Info("Loaded usage of %d key(s) in %v", n, time.Since(started).Round(time.Millisecond)) return true } -// sync is one round of the sync worker: write the unflushed counts of the promoted keys and the -// changed rate visitors, read back the rows changed since the last round, and prune once in a while. -// Every error is logged and leaves the state as it was, for the next round. +// sync is one round of the sync worker: write the unflushed counts of the promoted keys, read +// back the rows changed since the last round, and prune once in a while. Every error is logged +// and leaves the state as it was, for the next round. func (s *Service) sync() { now := s.now().Unix() ctx, cancel := context.WithTimeout(s.ctx, statementTimeout) s.write(ctx) cancel() ctx, cancel = context.WithTimeout(s.ctx, statementTimeout) - s.writeRateVisitors(ctx) + _, err := s.read(ctx, s.watermark-int64(readOverlap/time.Second)) cancel() - since := s.watermark - int64(readOverlap/time.Second) - ctx, cancel = context.WithTimeout(s.ctx, statementTimeout) - _, err := s.read(ctx, since) - cancel() - if err == nil { - ctx, cancel = context.WithTimeout(s.ctx, statementTimeout) - _, err = s.readRateVisitors(ctx, since) - cancel() - } if err != nil { - log.Tag(tag).Err(err).Warn("Cannot read usage or rate visitors; those of other nodes are stale until the next round") + log.Tag(tag).Err(err).Warn("Cannot read usage; counts of other nodes are stale until the next round") } else { s.watermark = now } @@ -398,14 +375,6 @@ func (s *Service) sync() { } else { log.Tag(tag).Debug("Pruned %d usage row(s)", deleted) } - ctx, cancel = context.WithTimeout(s.ctx, statementTimeout) - deleted, err = s.store.pruneRateVisitors(ctx, now-int64(rateVisitorPruneAfter/time.Second)) - cancel() - if err != nil { - log.Tag(tag).Err(err).Warn("Cannot prune rate visitors") - } else { - log.Tag(tag).Debug("Pruned %d rate visitor row(s)", deleted) - } } } diff --git a/limit/store.go b/limit/store.go index 60c3f3f4..5e2804b4 100644 --- a/limit/store.go +++ b/limit/store.go @@ -32,33 +32,20 @@ type usageRow struct { counts [numDaily]int64 } -// rateVisitorRow is one topic's row of the topic table. On the way to the store, an empty key -// keeps the stored rate visitor and only moves seen_at forward (a touch). -type rateVisitorRow struct { - topic string - key string - userID string - seenAt int64 // Unix seconds -} - // Store reads and writes the usage table, on SQLite (the user database) or Postgres. Rows are // stamped (updated_at) with the writing node's clock, like the reader's watermark, so the // database's clock does not matter. Only the service's sync worker calls it. type Store struct { - db *db.DB - queries *queries - writeChunk func(ctx context.Context, rows []usageRow, now int64) error // Dialect specific: unnest on Postgres, a prepared statement per row on SQLite - writeRateVisitorChunk func(ctx context.Context, rows []rateVisitorRow, now int64) error // Dialect specific, as writeChunk + db *db.DB + queries *queries + writeChunk func(ctx context.Context, rows []usageRow, now int64) error // Dialect specific: unnest on Postgres, a prepared statement per row on SQLite } // queries are the dialect's SQL statements type queries struct { - selectSince string - upsert string - prune string - selectRateVisitorsSince string - upsertRateVisitor string - pruneRateVisitors string + selectSince string + upsert string + prune string } // write adds the rows' counts to their keys' rows if the periods match, or replaces them if the @@ -106,15 +93,9 @@ func (s *Store) readSince(ctx context.Context, since int64, fn func(r *usageRow) // prune deletes the rows not written since the given time (unix seconds), in chunks, and returns how // many it deleted func (s *Store) prune(ctx context.Context, before int64) (int64, error) { - return s.pruneChunked(ctx, s.queries.prune, before) -} - -// pruneChunked runs the prune query (age, chunk size) until a chunk deletes fewer rows than the -// chunk size, and returns how many rows it deleted -func (s *Store) pruneChunked(ctx context.Context, query string, before int64) (int64, error) { var deleted int64 for { - result, err := s.db.ExecContext(ctx, query, before, pruneChunkRows) + result, err := s.db.ExecContext(ctx, s.queries.prune, before, pruneChunkRows) if err != nil { return deleted, err } @@ -128,42 +109,3 @@ func (s *Store) pruneChunked(ctx context.Context, query string, before int64) (i } } } - -// writeRateVisitors writes the rate visitor rows, sorted by topic in chunks (see write), stamped with now -// (unix seconds). A set replaces the stored rate visitor and seen_at; a touch keeps the rate visitor and moves -// seen_at forward. Rows are state, not deltas, so a failed write is simply written again. -func (s *Store) writeRateVisitors(ctx context.Context, rows []rateVisitorRow, now int64) error { - slices.SortFunc(rows, func(a, b rateVisitorRow) int { - return strings.Compare(a.topic, b.topic) - }) - for chunk := range slices.Chunk(rows, writeChunkRows) { - if err := s.writeRateVisitorChunk(ctx, chunk, now); err != nil { - return err - } - } - return nil -} - -// readRateVisitorsSince calls fn for every rate visitor row written at or after the given time (unix seconds), -// streaming them from the primary -func (s *Store) readRateVisitorsSince(ctx context.Context, since int64, fn func(r *rateVisitorRow)) error { - rows, err := s.db.QueryContext(ctx, s.queries.selectRateVisitorsSince, since) - if err != nil { - return err - } - defer rows.Close() - var r rateVisitorRow - for rows.Next() { - if err := rows.Scan(&r.topic, &r.key, &r.userID, &r.seenAt); err != nil { - return err - } - fn(&r) - } - return rows.Err() -} - -// pruneRateVisitors deletes the rate visitor rows not written since the given time (unix seconds), in chunks, -// and returns how many it deleted -func (s *Store) pruneRateVisitors(ctx context.Context, before int64) (int64, error) { - return s.pruneChunked(ctx, s.queries.pruneRateVisitors, before) -} diff --git a/limit/store_postgres.go b/limit/store_postgres.go index a25e1b45..b8154138 100644 --- a/limit/store_postgres.go +++ b/limit/store_postgres.go @@ -28,39 +28,18 @@ const ( postgresPruneQuery = `DELETE FROM usage WHERE updated_at < $1 AND key IN (SELECT key FROM usage WHERE updated_at < $1 LIMIT $2)` // The outer age check is rechecked on a row updated while the delete waits for it ) -// The rate visitor upsert, one fixed statement per chunk like the usage upsert. A set (a rate visitor key) -// replaces the rate visitor and seen_at; a touch (no rate visitor key) keeps the rate visitor and moves seen_at forward. -const ( - postgresSelectRateVisitorsSinceQuery = `SELECT topic, rate_visitor_key, rate_visitor_user_id, seen_at FROM topic WHERE updated_at >= $1` - postgresUpsertRateVisitorQuery = ` - INSERT INTO topic AS p (topic, rate_visitor_key, rate_visitor_user_id, seen_at, updated_at) - SELECT t, k, u, s, $5::BIGINT - FROM unnest($1::TEXT[], $2::TEXT[], $3::TEXT[], $4::BIGINT[]) AS d(t, k, u, s) - ON CONFLICT (topic) DO UPDATE SET - rate_visitor_key = CASE WHEN EXCLUDED.rate_visitor_key = '' THEN p.rate_visitor_key ELSE EXCLUDED.rate_visitor_key END, - rate_visitor_user_id = CASE WHEN EXCLUDED.rate_visitor_key = '' THEN p.rate_visitor_user_id ELSE EXCLUDED.rate_visitor_user_id END, - seen_at = CASE WHEN EXCLUDED.rate_visitor_key = '' THEN GREATEST(p.seen_at, EXCLUDED.seen_at) ELSE EXCLUDED.seen_at END, - updated_at = EXCLUDED.updated_at - ` - postgresPruneRateVisitorsQuery = `DELETE FROM topic WHERE updated_at < $1 AND topic IN (SELECT topic FROM topic WHERE updated_at < $1 LIMIT $2)` // As the usage prune -) - // NewPostgresStore returns the store on the shared pool; the usage table is the user schema's // (see package user, which creates and migrates it) func NewPostgresStore(d *db.DB) *Store { s := &Store{ db: d, queries: &queries{ - selectSince: postgresSelectSinceQuery, - upsert: postgresUpsertQuery, - prune: postgresPruneQuery, - selectRateVisitorsSince: postgresSelectRateVisitorsSinceQuery, - upsertRateVisitor: postgresUpsertRateVisitorQuery, - pruneRateVisitors: postgresPruneRateVisitorsQuery, + selectSince: postgresSelectSinceQuery, + upsert: postgresUpsertQuery, + prune: postgresPruneQuery, }, } s.writeChunk = s.writePostgres - s.writeRateVisitorChunk = s.writeRateVisitorsPostgres return s } @@ -86,17 +65,6 @@ func (s *Store) writePostgres(ctx context.Context, rows []usageRow, now int64) e return err } -// writeRateVisitorsPostgres writes the rate visitor rows in one statement, as arrays -func (s *Store) writeRateVisitorsPostgres(ctx context.Context, rows []rateVisitorRow, now int64) error { - topics, keys, userIDs := make([]string, len(rows)), make([]string, len(rows)), make([]string, len(rows)) - seenAts := make([]int64, len(rows)) - for i, r := range rows { - topics[i], keys[i], userIDs[i], seenAts[i] = r.topic, r.key, r.userID, r.seenAt - } - _, err := s.db.ExecContext(ctx, s.queries.upsertRateVisitor, topics, keys, userIDs, seenAts, now) - return err -} - // mayHaveCommittedPostgres reports whether a failed write may have committed anyway: a timeout // after the statement was sent, e.g. a commit acknowledged after the deadline. An error from the // server, or one before the statement was sent, means nothing was written. diff --git a/limit/store_sqlite.go b/limit/store_sqlite.go index 189c8d23..dd9b1d44 100644 --- a/limit/store_sqlite.go +++ b/limit/store_sqlite.go @@ -24,37 +24,18 @@ const ( sqlitePruneQuery = `DELETE FROM usage WHERE updated_at < ?1 AND key IN (SELECT key FROM usage WHERE updated_at < ?1 LIMIT ?2)` ) -// Same as the Postgres rate visitor queries (see there), one row per statement -const ( - sqliteSelectRateVisitorsSinceQuery = `SELECT topic, rate_visitor_key, rate_visitor_user_id, seen_at FROM topic WHERE updated_at >= ?` - sqliteUpsertRateVisitorQuery = ` - INSERT INTO topic (topic, rate_visitor_key, rate_visitor_user_id, seen_at, updated_at) - VALUES (?, ?, ?, ?, ?) - ON CONFLICT (topic) DO UPDATE SET - rate_visitor_key = CASE WHEN excluded.rate_visitor_key = '' THEN rate_visitor_key ELSE excluded.rate_visitor_key END, - rate_visitor_user_id = CASE WHEN excluded.rate_visitor_key = '' THEN rate_visitor_user_id ELSE excluded.rate_visitor_user_id END, - seen_at = CASE WHEN excluded.rate_visitor_key = '' THEN MAX(seen_at, excluded.seen_at) ELSE excluded.seen_at END, - updated_at = excluded.updated_at - ` - sqlitePruneRateVisitorsQuery = `DELETE FROM topic WHERE updated_at < ?1 AND topic IN (SELECT topic FROM topic WHERE updated_at < ?1 LIMIT ?2)` -) - // NewSQLiteStore returns the store on the user database (auth-file), which holds the usage // table (see package user, which creates and migrates it) func NewSQLiteStore(d *db.DB) *Store { s := &Store{ db: d, queries: &queries{ - selectSince: sqliteSelectSinceQuery, - upsert: sqliteUpsertQuery, - prune: sqlitePruneQuery, - selectRateVisitorsSince: sqliteSelectRateVisitorsSinceQuery, - upsertRateVisitor: sqliteUpsertRateVisitorQuery, - pruneRateVisitors: sqlitePruneRateVisitorsQuery, + selectSince: sqliteSelectSinceQuery, + upsert: sqliteUpsertQuery, + prune: sqlitePruneQuery, }, } s.writeChunk = s.writeSQLite - s.writeRateVisitorChunk = s.writeRateVisitorsSQLite return s } @@ -74,20 +55,3 @@ func (s *Store) writeSQLite(ctx context.Context, rows []usageRow, now int64) err return nil }) } - -// writeRateVisitorsSQLite writes the rate visitor rows in one transaction, one prepared statement per row -func (s *Store) writeRateVisitorsSQLite(ctx context.Context, rows []rateVisitorRow, now int64) error { - return db.ExecTx(s.db, func(tx *sql.Tx) error { - stmt, err := tx.PrepareContext(ctx, s.queries.upsertRateVisitor) - if err != nil { - return err - } - defer stmt.Close() - for _, r := range rows { - if _, err := stmt.ExecContext(ctx, r.topic, r.key, r.userID, r.seenAt, now); err != nil { - return err - } - } - return nil - }) -} diff --git a/limit/store_test.go b/limit/store_test.go index e6c837ec..9693b318 100644 --- a/limit/store_test.go +++ b/limit/store_test.go @@ -5,7 +5,6 @@ import ( "fmt" "path/filepath" "slices" - "strings" "testing" "time" @@ -214,83 +213,3 @@ func keysOf(rows map[string]usageRow) []string { slices.Sort(keys) return keys } - -func TestStore_WriteRateVisitors_SetReplacesTouchKeeps(t *testing.T) { - forEachBackend(t, func(t *testing.T, newStore func() *Store) { - s := newStore() - write := func(topic, key string, seenAt, now int64) { - require.Nil(t, s.writeRateVisitors(context.Background(), []rateVisitorRow{newRateVisitorRow(topic, key, seenAt)}, now)) - } - write("up1", "ip:1.2.3.4", 100, 101) - write("up1", "", 200, 201) // A touch keeps the rate visitor - require.Equal(t, newRateVisitorRow("up1", "ip:1.2.3.4", 200), readAllRateVisitors(t, s, 0)["up1"]) - write("up1", "", 150, 202) // An older touch does not move seen_at back - require.Equal(t, newRateVisitorRow("up1", "ip:1.2.3.4", 200), readAllRateVisitors(t, s, 0)["up1"]) - write("up1", "user:u_1", 50, 203) // A set replaces both - require.Equal(t, newRateVisitorRow("up1", "user:u_1", 50), readAllRateVisitors(t, s, 0)["up1"]) - write("up2", "", 300, 301) // A touch of an unknown topic: no rate visitor - require.Equal(t, newRateVisitorRow("up2", "", 300), readAllRateVisitors(t, s, 0)["up2"]) - - // The read-back is by write time - require.Equal(t, []string{"up2"}, rateVisitorTopicsOf(readAllRateVisitors(t, s, 301))) - }) -} - -func TestStore_WriteRateVisitors_Chunked(t *testing.T) { - forEachBackend(t, func(t *testing.T, newStore func() *Store) { - s := newStore() - rows := make([]rateVisitorRow, 0) - for i := 0; i < writeChunkRows+5; i++ { - rows = append(rows, newRateVisitorRow(fmt.Sprintf("up%012d", i), "ip:1.2.3.4", 100)) - } - require.Nil(t, s.writeRateVisitors(context.Background(), rows, 100)) - require.Len(t, readAllRateVisitors(t, s, 0), writeChunkRows+5) - }) -} - -func TestStore_PruneRateVisitors(t *testing.T) { - forEachBackend(t, func(t *testing.T, newStore func() *Store) { - s := newStore() - for i := 0; i < 5; i++ { - require.Nil(t, s.writeRateVisitors(context.Background(), []rateVisitorRow{newRateVisitorRow(fmt.Sprintf("up%d", i), "", 1)}, int64(100+i))) - } - deleted, err := s.pruneRateVisitors(context.Background(), 103) - require.Nil(t, err) - require.Equal(t, int64(3), deleted) - require.Equal(t, []string{"up3", "up4"}, rateVisitorTopicsOf(readAllRateVisitors(t, s, 0))) - }) -} - -// readAllRateVisitors returns the rate visitor rows written at or after since, by topic -func readAllRateVisitors(t *testing.T, s *Store, since int64) map[string]rateVisitorRow { - t.Helper() - rows := make(map[string]rateVisitorRow) - require.Nil(t, s.readRateVisitorsSince(context.Background(), since, func(r *rateVisitorRow) { - rows[r.topic] = *r - })) - return rows -} - -// rateVisitorTopicsOf returns the map's topics, sorted -func rateVisitorTopicsOf(rows map[string]rateVisitorRow) []string { - topics := make([]string, 0, len(rows)) - for topic := range rows { - topics = append(topics, topic) - } - slices.Sort(topics) - return topics -} - -// newRateVisitorRow returns a rate visitor row; a user key ("user:") has its user ID -func newRateVisitorRow(topic, key string, seenAt int64) rateVisitorRow { - userID, _ := strings.CutPrefix(key, "user:") - if userID == key { - userID = "" - } - return rateVisitorRow{ - topic: topic, - key: key, - userID: userID, - seenAt: seenAt, - } -} diff --git a/limit/types.go b/limit/types.go index d2a23821..f89af1d4 100644 --- a/limit/types.go +++ b/limit/types.go @@ -53,14 +53,6 @@ type Usage struct { counts [numKinds]int64 } -// RateVisitor is who pays for the pushes to a UnifiedPush topic (its subscriber), and when the topic last -// had subscriber activity (see Service.SetRateVisitor) -type RateVisitor struct { - Key string // The subscriber's limiter key; empty for a topic with subscriber activity or a refused push, but no rate visitor - UserID string // The subscriber's user ID, empty if anonymous - SeenAt time.Time // Last subscribe, poll or open stream, or the refused push that made the topic known; node clock -} - // Config configures the service type Config struct { Store *Store // Where the daily counts are shared and kept across restarts; nil for memory only @@ -95,29 +87,6 @@ type bucketShard struct { mu sync.Mutex // Protects entries and everything inside them } -// rateVisitorWrite is what a rate visitor entry still has to write to the store; a set outranks a touch -type rateVisitorWrite uint8 - -const ( - rateVisitorWriteNone rateVisitorWrite = iota - rateVisitorWriteTouch // Only seen_at: keep the stored rate visitor (TouchRateVisitor, NoteTopic) - rateVisitorWriteSet // RateVisitor and seen_at (SetRateVisitor) -) - -// rateVisitorEntry is one topic's rate visitor, as last set here or read from the store -type rateVisitorEntry struct { - key string - userID string - seenAt int64 // Unix seconds - write rateVisitorWrite // Unwritten change; the read-back leaves an entry with one alone -} - -// rateVisitorTable is the rate visitors of the UnifiedPush topics, by topic -type rateVisitorTable struct { - entries map[string]*rateVisitorEntry - mu sync.Mutex // Protects entries and everything inside them -} - // bucket is a token bucket. It is created full: last is zero until the first use, and tokens // then start at the rule's burst. Tokens may go negative (debt) when a burst shrinks. type bucket struct { diff --git a/server/rate_visitor.go b/server/rate_visitor.go index 48f30e45..22cbe6f2 100644 --- a/server/rate_visitor.go +++ b/server/rate_visitor.go @@ -1,133 +1,179 @@ package server import ( - "errors" + "net/http" + "net/netip" "strings" "sync" "time" + "heckel.io/ntfy/v2/log" + tp "heckel.io/ntfy/v2/topic" "heckel.io/ntfy/v2/user" + vis "heckel.io/ntfy/v2/visitor" ) const ( // unifiedPushAbandonedAfter is how long a UnifiedPush topic may go without subscriber activity // (subscribe, poll, open stream) before a push to it answers 404 instead of 507: RFC 8030 // section 7.3 signals an expired subscription with 404, and app servers like Mastodon then - // delete it (#664). It is long enough for a phone that is offline for a few days. + // delete it (#664). It is long enough for a phone that is offline for a few days. A topic + // that never had a subscriber counts from its first refused push. unifiedPushAbandonedAfter = 7 * oneDay - // rateVisitorUserCacheTTL is how long a rate visitor's user is cached (see rateVisitorUserCache) - rateVisitorUserCacheTTL = time.Minute + // rateVisitorLiveFor is how long a topic's rate visitor pays for its pushes after the + // subscriber was last active on any node (see Server.rateVisitor); a stream open here keeps + // it live regardless + rateVisitorLiveFor = oneDay ) -// rateVisitorUserCache caches the users of the topics' rate visitors by ID, so a push does not read the user -// database; a change to the user (a new tier) applies to pushes within rateVisitorUserCacheTTL +// rateVisitorUserCache holds the users of the live rate visitors by ID, so a push never reads the +// user database: the server refreshes them in the background (see Server.refreshRateVisitorUsers), +// and a subscribe puts the subscriber's user in right away type rateVisitorUserCache struct { - lookup func(id string) (*user.User, error) // userManager.UserByID, nil without a user manager; replaced in tests - entries map[string]*rateVisitorUserEntry - mu sync.Mutex // Protects entries + lookup func(ids []string) ([]*user.User, error) // userManager.UsersByIDs, nil without a user manager; replaced in tests + users map[string]*user.User // By ID; a deleted user is not in here + mu sync.Mutex // Protects users } -// rateVisitorUserEntry is a cached user; nil if it does not exist (anymore) -type rateVisitorUserEntry struct { - user *user.User - fetched time.Time -} - -func newRateVisitorUserCache(lookup func(id string) (*user.User, error)) *rateVisitorUserCache { +func newRateVisitorUserCache(lookup func(ids []string) ([]*user.User, error)) *rateVisitorUserCache { return &rateVisitorUserCache{ - lookup: lookup, - entries: make(map[string]*rateVisitorUserEntry), + lookup: lookup, + users: make(map[string]*user.User), } } -// rateVisitor returns the visitor of the topic's live rate visitor, or nil if it has none. A rate visitor is -// live while the topic has a stream open here, or for visitorExpungeAfter after the subscriber was -// last seen on any node (the manager of a node with an open stream touches it). A deleted -// user pays for nothing; a failed user lookup is returned, so the publisher retries. -func (s *Server) rateVisitor(t *topic) (*visitor, error) { +// rateVisitor returns the visitor the topic's pushes are charged to, or nil if the topic has no +// live rate visitor. A rate visitor is live while the topic has a stream open here, or for +// rateVisitorLiveFor after the subscriber was last active on any node (the manager of a node +// with an open stream keeps it active). It is charged like the subscriber would be now: on its +// user key with its tier's limits if the user has a tier (as of the last refresh), else on its +// IP key with the anonymous limits. A deleted user, or one not looked up yet, is charged on the IP. +func (s *Server) rateVisitor(t *topic) *visitor { if !s.config.VisitorSubscriberRateLimiting || !isUnifiedPushTopic(t.ID) { - return nil, nil + return nil } - p := s.limiter.RateVisitor(t.ID) - if p == nil || p.Key == "" { - return nil, nil - } else if t.SubscribersCount() == 0 && time.Since(p.SeenAt) > visitorExpungeAfter { - return nil, nil + state := s.topicService.Get(t.ID) + if state == nil || state.RateVisitor == nil { + return nil + } else if t.SubscribersCount() == 0 && time.Since(state.SubscribedAt) > rateVisitorLiveFor { + return nil + } + ip, err := netip.ParseAddr(state.RateVisitor.IP) + if err != nil { + return nil } var u *user.User - if p.UserID != "" { - var err error - if u, err = s.rateVisitorUsers.User(p.UserID); err != nil { - return nil, err - } else if u == nil { - return nil, nil - } + if state.RateVisitor.UserID != "" { + u = s.rateVisitorUsers.User(state.RateVisitor.UserID) } - return newRateVisitor(s.visitors, p.Key, u), nil + return newVisitor(s.visitors, ip, u) } // unifiedPushTopicAbandoned returns true if the topic has had no subscriber activity, on any node, -// for unifiedPushAbandonedAfter, counted from its first refused push if it never had any -func (s *Server) unifiedPushTopicAbandoned(t *topic) bool { +// for unifiedPushAbandonedAfter, counted from its first refused push if it never had any. Before +// saying so, it reads the topic's row from the database, since a subscriber may have come back +// on another node and the row be fresher than what this node has read back; if that read fails, +// the topic is not abandoned (the publisher retries). +func (s *Server) unifiedPushTopicAbandoned(r *http.Request, t *topic) bool { if !isUnifiedPushTopic(t.ID) || t.SubscribersCount() > 0 { return false } - p := s.limiter.RateVisitor(t.ID) - return p != nil && time.Since(p.SeenAt) > unifiedPushAbandonedAfter + state := s.topicService.Get(t.ID) + if state == nil || !topicAbandoned(state) { + return false + } + state, err := s.topicService.Refresh(r.Context(), t.ID) + if err != nil { + logr(r).Tag(tagPublish).With(t).Err(err).Warn("Cannot read topic %s to confirm it is abandoned; answering as if it were not", t.ID) + return false + } + return state != nil && topicAbandoned(state) } -// touchRateVisitors marks subscriber activity on the UnifiedPush topics now, keeping their rate visitors -func (s *Server) touchRateVisitors(topics []*topic) { +// touchTopics records subscriber activity on the UnifiedPush topics now, e.g. a stream that +// closed, keeping their rate visitors +func (s *Server) touchTopics(topics []*topic) { if !s.config.VisitorSubscriberRateLimiting { return } for _, t := range topics { if isUnifiedPushTopic(t.ID) { - s.limiter.TouchRateVisitor(t.ID) + s.topicService.Touch(t.ID) } } } -// User returns the user with the given ID, from the cache if it was looked up within -// rateVisitorUserCacheTTL, or nil if it does not exist or is deleted -func (c *rateVisitorUserCache) User(id string) (*user.User, error) { - c.mu.Lock() - e, ok := c.entries[id] - c.mu.Unlock() - if ok && time.Since(e.fetched) < rateVisitorUserCacheTTL { - return e.user, nil +// keepaliveTopics records that the UnifiedPush topics have a stream open here, so the other nodes +// see their subscribers as active too +func (s *Server) keepaliveTopics(topics []*topic) { + if !s.config.VisitorSubscriberRateLimiting { + return } - var u *user.User - if c.lookup != nil { - var err error - u, err = c.lookup(id) - if errors.Is(err, user.ErrUserNotFound) { - u = nil - } else if err != nil { - return nil, err - } else if u.Deleted { - u = nil + for _, t := range topics { + if isUnifiedPushTopic(t.ID) { + s.topicService.Keepalive(t.ID) } } - c.mu.Lock() - defer c.mu.Unlock() - c.entries[id] = &rateVisitorUserEntry{ - user: u, - fetched: time.Now(), - } - return u, nil } -// Prune forgets the users looked up longer than rateVisitorUserCacheTTL ago -func (c *rateVisitorUserCache) Prune() { +// refreshRateVisitorUsers looks up the users of the live rate visitors again, so a tier change (or +// a deleted account) applies to the pushes within a manager interval +func (s *Server) refreshRateVisitorUsers() { + if !s.config.VisitorSubscriberRateLimiting || s.userManager == nil { + return + } + ids := s.topicService.RateVisitorUserIDs(time.Now().Add(-rateVisitorLiveFor)) + if err := s.rateVisitorUsers.Refresh(ids); err != nil { + log.Tag(tagManager).Err(err).Warn("Cannot refresh the users of %d rate visitor(s); keeping the cached ones", len(ids)) + } +} + +// User returns the cached user with the given ID, or nil if it is not cached (not a live rate +// visitor's, not looked up yet, or deleted) +func (c *rateVisitorUserCache) User(id string) *user.User { c.mu.Lock() defer c.mu.Unlock() - for id, e := range c.entries { - if time.Since(e.fetched) >= rateVisitorUserCacheTTL { - delete(c.entries, id) + return c.users[id] +} + +// Put caches the user, e.g. the subscriber's on a subscribe, so its pushes are charged with its +// tier before the next refresh +func (c *rateVisitorUserCache) Put(u *user.User) { + c.mu.Lock() + defer c.mu.Unlock() + c.users[u.ID] = u +} + +// Refresh replaces the cached users with fresh lookups of the given IDs, dropping the others. On +// an error the cache is left as it is: stale users are better than none. +func (c *rateVisitorUserCache) Refresh(ids []string) error { + users := make(map[string]*user.User, len(ids)) + if len(ids) > 0 && c.lookup != nil { + looked, err := c.lookup(ids) + if err != nil { + return err + } + for _, u := range looked { + if !u.Deleted { + users[u.ID] = u + } } } + c.mu.Lock() + defer c.mu.Unlock() + c.users = users + return nil +} + +// topicAbandoned returns true if the topic's newest observation (a subscriber, or its first +// refused push) is older than unifiedPushAbandonedAfter +func topicAbandoned(state *tp.Topic) bool { + newest := state.FirstSeenAt + if state.SubscribedAt.After(newest) { + newest = state.SubscribedAt + } + return time.Since(newest) > unifiedPushAbandonedAfter } // isUnifiedPushTopic returns true if the topic may have a rate visitor: an "up..." topic of the @@ -135,3 +181,9 @@ func (c *rateVisitorUserCache) Prune() { func isUnifiedPushTopic(id string) bool { return strings.HasPrefix(id, unifiedPushTopicPrefix) && len(id) == unifiedPushTopicLength } + +// maskedIP returns the visitor's IP masked to the configured prefix bits, as a string, which is +// what is kept of a subscriber (see topic.RateVisitor) +func (v *visitor) maskedIP() string { + return vis.Mask(v.ip, v.env.limitsConfig).String() +} diff --git a/server/rate_visitor_test.go b/server/rate_visitor_test.go index aa213f42..dff8c337 100644 --- a/server/rate_visitor_test.go +++ b/server/rate_visitor_test.go @@ -6,13 +6,14 @@ import ( "fmt" "net/http" "net/netip" + "strings" "testing" "time" "github.com/stretchr/testify/require" "heckel.io/ntfy/v2/db" "heckel.io/ntfy/v2/db/pg" - "heckel.io/ntfy/v2/limit" + tp "heckel.io/ntfy/v2/topic" "heckel.io/ntfy/v2/user" "heckel.io/ntfy/v2/util" ) @@ -30,7 +31,7 @@ func TestServer_UnifiedPush_RateVisitorOnOtherNode(t *testing.T) { }) require.Equal(t, 200, rr.Code) waitForWithMaxWait(t, 3*syncWait, func() bool { - return rateVisitorKey(s2, "upAAAAAAAAAAAA") == "ip:1.2.3.4" + return s2.topicService.Get("upAAAAAAAAAAAA") != nil }) response := request(t, s2, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil) require.Equal(t, 200, response.Code) @@ -39,9 +40,41 @@ func TestServer_UnifiedPush_RateVisitorOnOtherNode(t *testing.T) { }) } +func TestServer_UnifiedPush_RateVisitorOnOtherNodeUsesTierAfterRefresh(t *testing.T) { + // The other node knows the subscriber's user ID from the row, but its tier only once the + // background refresh has looked it up: until then, the subscriber's IP is charged + forEachBackend(t, func(t *testing.T, databaseURL string) { + c := newTestConfigWithAuthFile(t, databaseURL) + c.VisitorSubscriberRateLimiting = true + s1 := newTestServer(t, c) + s2 := newTestServer(t, c) + require.Nil(t, s1.userManager.AddTier(&user.Tier{ + Code: "pro", + MessageLimit: 10, + })) + require.Nil(t, s1.userManager.AddUser("phil", "phil", user.RoleUser, false)) + require.Nil(t, s1.userManager.ChangeTier("phil", "pro")) + phil, err := s1.userManager.User("phil") + require.Nil(t, err) + rr := request(t, s1, "GET", "/upAAAAAAAAAAAA/json?poll=1", "", map[string]string{ + "Authorization": util.BasicAuth("phil", "phil"), + }, func(r *http.Request) { + r.RemoteAddr = "1.2.3.4:1234" + }) + require.Equal(t, 200, rr.Code) + require.Equal(t, "user:"+phil.ID, rateVisitorKey(s1, "upAAAAAAAAAAAA")) // The subscribe primes the cache + waitForWithMaxWait(t, 3*syncWait, func() bool { + return s2.topicService.Get("upAAAAAAAAAAAA") != nil + }) + require.Equal(t, "ip:1.2.3.4", rateVisitorKey(s2, "upAAAAAAAAAAAA")) + s2.execManager() // Refreshes the rate visitors' users + require.Equal(t, "user:"+phil.ID, rateVisitorKey(s2, "upAAAAAAAAAAAA")) + }) +} + func TestServer_UnifiedPush_StreamOnOtherNodeKeepsRateVisitor(t *testing.T) { - // A stream open on one node keeps the topic's rate visitor live on the other nodes too: the manager - // touches it, so it does not expire a day after the subscribe + // A stream open on one node keeps the topic's rate visitor live on the other nodes too: the + // manager records it as subscriber activity, so it does not expire a day after the subscribe forEachBackend(t, func(t *testing.T, databaseURL string) { c := newTestConfigWithAuthFile(t, databaseURL) c.VisitorSubscriberRateLimiting = true @@ -59,13 +92,17 @@ func TestServer_UnifiedPush_StreamOnOtherNodeKeepsRateVisitor(t *testing.T) { to := lockedTopic(s1, "upAAAAAAAAAAAA") return to != nil && to.SubscribersCount() == 1 }) + waitForWithMaxWait(t, 3*syncWait, func() bool { + return s2.topicService.Get("upAAAAAAAAAAAA") != nil + }) // A day later, the stream is still open on s1 - ageRateVisitor(s1, "upAAAAAAAAAAAA", 25*time.Hour) + ageTopic(s1, "upAAAAAAAAAAAA", 25*time.Hour) + ageTopic(s2, "upAAAAAAAAAAAA", 25*time.Hour) + require.Equal(t, 507, request(t, s2, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil).Code) s1.execManager() waitForWithMaxWait(t, 3*syncWait, func() bool { - p := s2.limiter.RateVisitor("upAAAAAAAAAAAA") - return p != nil && p.Key == "ip:1.2.3.4" && time.Since(p.SeenAt) < time.Hour + return time.Since(s2.topicService.Get("upAAAAAAAAAAAA").SubscribedAt) < time.Hour }) response := request(t, s2, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil) require.Equal(t, 200, response.Code) @@ -75,50 +112,92 @@ func TestServer_UnifiedPush_StreamOnOtherNodeKeepsRateVisitor(t *testing.T) { func TestServer_UnifiedPush_AbandonedTopicAnswers404(t *testing.T) { // A UnifiedPush push to a topic without subscriber activity for unifiedPushAbandonedAfter is // told to go away (404, RFC 8030 section 7.3); before that, it may be retried (507) - c := newTestConfig(t, "") - c.VisitorSubscriberRateLimiting = true - s := newTestServer(t, c) - push := func() *errHTTP { - response := request(t, s, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil) - if response.Code == 200 { - return nil + forEachBackend(t, func(t *testing.T, databaseURL string) { + c := newTestConfigWithAuthFile(t, databaseURL) + c.VisitorSubscriberRateLimiting = true + s := newTestServer(t, c) + push := func() *errHTTP { + response := request(t, s, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil) + if response.Code == 200 { + return nil + } + return toHTTPError(t, response.Body.String()) } - return toHTTPError(t, response.Body.String()) - } - // Never subscribed: 507 for a week from the first push, then 404 - require.Equal(t, 50701, push().Code) - ageRateVisitor(s, "upAAAAAAAAAAAA", unifiedPushAbandonedAfter-time.Hour) - require.Equal(t, 50701, push().Code) - ageRateVisitor(s, "upAAAAAAAAAAAA", unifiedPushAbandonedAfter+time.Hour) - require.Equal(t, 40402, push().Code) - require.Equal(t, 404, push().HTTPCode) + // Never subscribed: 507 for a week from the first push, then 404 + require.Equal(t, 50701, push().Code) + ageTopic(s, "upAAAAAAAAAAAA", unifiedPushAbandonedAfter-time.Hour) + require.Equal(t, 50701, push().Code) + ageTopic(s, "upAAAAAAAAAAAA", 2*time.Hour) + require.Equal(t, 40402, push().Code) + require.Equal(t, 404, push().HTTPCode) - // A subscriber comes (back): it pays again; a week after it is gone, 404 - rr := request(t, s, "GET", "/upAAAAAAAAAAAA/json?poll=1", "", nil, func(r *http.Request) { - r.RemoteAddr = "1.2.3.4:1234" - }) - require.Equal(t, 200, rr.Code) - require.Nil(t, push()) - ageRateVisitor(s, "upAAAAAAAAAAAA", 2*oneDay) - require.Equal(t, 50701, push().Code) - ageRateVisitor(s, "upAAAAAAAAAAAA", unifiedPushAbandonedAfter+time.Hour) - require.Equal(t, 40402, push().Code) - - // An open stream keeps the topic, however long ago it was seen otherwise - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - go func() { - request(t, s, "GET", "/upAAAAAAAAAAAA/json", "", nil, func(r *http.Request) { + // A subscriber comes (back): it pays again; a week after it is gone, 404 + rr := request(t, s, "GET", "/upAAAAAAAAAAAA/json?poll=1", "", nil, func(r *http.Request) { r.RemoteAddr = "1.2.3.4:1234" + }) + require.Equal(t, 200, rr.Code) + require.Nil(t, push()) + ageTopic(s, "upAAAAAAAAAAAA", 2*oneDay) + require.Equal(t, 50701, push().Code) + ageTopic(s, "upAAAAAAAAAAAA", unifiedPushAbandonedAfter) + require.Equal(t, 40402, push().Code) + + // An open stream keeps the topic, however long ago it was seen otherwise + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go func() { + request(t, s, "GET", "/upAAAAAAAAAAAA/json", "", nil, func(r *http.Request) { + r.RemoteAddr = "1.2.3.4:1234" + *r = *r.WithContext(ctx) + }) + }() + waitFor(t, func() bool { + return lockedTopic(s, "upAAAAAAAAAAAA").SubscribersCount() == 1 + }) + ageTopic(s, "upAAAAAAAAAAAA", unifiedPushAbandonedAfter+time.Hour) + require.Nil(t, push()) + }) +} + +func TestServer_UnifiedPush_AbandonedTopicReadsRowBeforeAnswering404(t *testing.T) { + // Before a 404, the node reads the topic's row: a subscriber may have come back on another + // node, and the row may be fresher than what this node has read back so far + forEachBackend(t, func(t *testing.T, databaseURL string) { + c := newTestConfigWithAuthFile(t, databaseURL) + c.VisitorSubscriberRateLimiting = true + s := newTestServer(t, c) + ageTopic(s, "upAAAAAAAAAAAA", unifiedPushAbandonedAfter+time.Hour) + now := time.Now().Unix() + execUserDB(t, c, fmt.Sprintf(`INSERT INTO topic (topic, first_seen_at, subscribed_at, rate_visitor_user_id, rate_visitor_ip, updated_at) VALUES ('upAAAAAAAAAAAA', %d, %d, '', '1.2.3.4', 1)`, now-2*int64(oneDay/time.Second), now-2*int64(oneDay/time.Second))) + response := request(t, s, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil) + require.Equal(t, 507, response.Code) + require.WithinDuration(t, time.Now().Add(-2*oneDay), s.topicService.Get("upAAAAAAAAAAAA").SubscribedAt, time.Minute) + + // Known as abandoned, and so is the row: 404 + ageTopic(s, "upAAAAAAAAAAAA", 6*oneDay) + execUserDB(t, c, fmt.Sprintf(`UPDATE topic SET subscribed_at = %d, first_seen_at = %d`, now-8*int64(oneDay/time.Second), now-8*int64(oneDay/time.Second))) + response = request(t, s, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil) + require.Equal(t, 404, response.Code) + }) +} + +func TestServer_UnifiedPush_AbandonedTopicAnswers507IfRowReadFails(t *testing.T) { + // If the row cannot be read, the push is not told to go away: 507, and the app server retries + forEachBackend(t, func(t *testing.T, databaseURL string) { + c := newTestConfigWithAuthFile(t, databaseURL) + c.VisitorSubscriberRateLimiting = true + s := newTestServer(t, c) + ageTopic(s, "upAAAAAAAAAAAA", unifiedPushAbandonedAfter+time.Hour) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + response := request(t, s, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil, func(r *http.Request) { *r = *r.WithContext(ctx) }) - }() - waitFor(t, func() bool { - return lockedTopic(s, "upAAAAAAAAAAAA").SubscribersCount() == 1 + require.Equal(t, 507, response.Code) + response = request(t, s, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil) + require.Equal(t, 404, response.Code) }) - ageRateVisitor(s, "upAAAAAAAAAAAA", unifiedPushAbandonedAfter+time.Hour) - require.Nil(t, push()) } func TestServer_UnifiedPush_AbandonedTopicOnOtherNode(t *testing.T) { @@ -131,7 +210,7 @@ func TestServer_UnifiedPush_AbandonedTopicOnOtherNode(t *testing.T) { response := request(t, s1, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil) require.Equal(t, 507, response.Code) s1.Stop() // Writes the topic - execUserDB(t, c, fmt.Sprintf(`UPDATE topic SET seen_at = seen_at - %d`, int64((unifiedPushAbandonedAfter+time.Hour)/time.Second))) + execUserDB(t, c, fmt.Sprintf(`UPDATE topic SET first_seen_at = first_seen_at - %d`, int64((unifiedPushAbandonedAfter+time.Hour)/time.Second))) s2 := newTestServer(t, c) response = request(t, s2, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil) require.Equal(t, 404, response.Code) @@ -139,6 +218,30 @@ func TestServer_UnifiedPush_AbandonedTopicOnOtherNode(t *testing.T) { }) } +func TestServer_UnifiedPush_RefusedPushDoesNotReviveRateVisitor(t *testing.T) { + // A refused push to a topic whose rate visitor expired is not subscriber activity: the rate + // visitor stays expired, here and in the row + forEachBackend(t, func(t *testing.T, databaseURL string) { + c := newTestConfigWithAuthFile(t, databaseURL) + c.VisitorSubscriberRateLimiting = true + s := newTestServer(t, c) + rr := request(t, s, "GET", "/upAAAAAAAAAAAA/json?poll=1", "", nil, func(r *http.Request) { + r.RemoteAddr = "1.2.3.4:1234" + }) + require.Equal(t, 200, rr.Code) + ageTopic(s, "upAAAAAAAAAAAA", 25*time.Hour) + subscribed := s.topicService.Get("upAAAAAAAAAAAA").SubscribedAt + for i := 0; i < 3; i++ { + require.Equal(t, 507, request(t, s, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil).Code) + } + require.Equal(t, subscribed, s.topicService.Get("upAAAAAAAAAAAA").SubscribedAt) + s.Stop() // Writes the topic + var subscribedAt int64 + queryUserDB(t, c, `SELECT subscribed_at FROM topic WHERE topic = 'upAAAAAAAAAAAA'`, &subscribedAt) + require.Equal(t, subscribed.Unix(), subscribedAt) + }) +} + func TestServer_UnifiedPush_NonUnifiedPushTopicIsNotAbandoned(t *testing.T) { // A topic that cannot have a rate visitor (not an "up..." topic) keeps answering 507 c := newTestConfig(t, "") @@ -146,21 +249,25 @@ func TestServer_UnifiedPush_NonUnifiedPushTopicIsNotAbandoned(t *testing.T) { s := newTestServer(t, c) response := request(t, s, "PUT", "/mytopic?up=1", "some message", nil) require.Equal(t, 507, response.Code) - require.Nil(t, s.limiter.RateVisitor("mytopic")) + require.Nil(t, s.topicService.Get("mytopic")) } func TestServer_MatrixGateway_AbandonedTopicRejectsPushKey(t *testing.T) { - // The Matrix gateway answers an abandoned topic with a rejected pushkey, not with the 404 + // The Matrix gateway answers an abandoned topic with a rejected pushkey, not with the 404; a + // topic without a subscriber for less than that is a 507, however long the node has known it c := newTestConfig(t, "") c.VisitorSubscriberRateLimiting = true s := newTestServer(t, c) notification := `{"notification":{"devices":[{"pushkey":"http://127.0.0.1:12345/upAAAAAAAAAAAA?up=1"}]}}` response := request(t, s, "POST", "/_matrix/push/v1/notify", notification, nil) require.Equal(t, 507, response.Code) - ageRateVisitor(s, "upAAAAAAAAAAAA", unifiedPushAbandonedAfter+time.Hour) + lockedTopic(s, "upAAAAAAAAAAAA").lastAccess = time.Now().Add(-13 * time.Hour) + response = request(t, s, "POST", "/_matrix/push/v1/notify", notification, nil) + require.Equal(t, 507, response.Code) + ageTopic(s, "upAAAAAAAAAAAA", unifiedPushAbandonedAfter+time.Hour) response = request(t, s, "POST", "/_matrix/push/v1/notify", notification, nil) require.Equal(t, 200, response.Code) - require.Equal(t, `{"rejected":["http://127.0.0.1:12345/upAAAAAAAAAAAA?up=1"]}`+"\n", response.Body.String()) + require.Equal(t, `{"rejected":["http://127.0.0.1:12345/upAAAAAAAAAAAA?up=1"]}`, strings.TrimSpace(response.Body.String())) } func TestServer_UnifiedPush_RateVisitorLivenessIsTheSubscribers(t *testing.T) { @@ -175,18 +282,16 @@ func TestServer_UnifiedPush_RateVisitorLivenessIsTheSubscribers(t *testing.T) { }) require.Equal(t, 200, rr.Code) s.limiter.Prune(0) // The key is gone, the subscriber is not - rateVisitor, err := s.rateVisitor(lockedTopic(s, "upAAAAAAAAAAAA")) - require.Nil(t, err) + rateVisitor := s.rateVisitor(lockedTopic(s, "upAAAAAAAAAAAA")) require.NotNil(t, rateVisitor) - ageRateVisitor(s, "upAAAAAAAAAAAA", 2*visitorExpungeAfter) // The subscriber is gone - require.True(t, rateVisitor.MessageAllowed()) // A publisher charges the key - rateVisitor, err = s.rateVisitor(lockedTopic(s, "upAAAAAAAAAAAA")) - require.Nil(t, err) - require.Nil(t, rateVisitor) + ageTopic(s, "upAAAAAAAAAAAA", 2*rateVisitorLiveFor) // The subscriber is gone + require.True(t, rateVisitor.MessageAllowed()) // A publisher charges the key + require.Nil(t, s.rateVisitor(lockedTopic(s, "upAAAAAAAAAAAA"))) } -func TestServer_UnifiedPush_TierChangeAppliesAfterUserCache(t *testing.T) { - // Pushes are billed with the subscriber's current tier, as of the cached user lookup +func TestServer_UnifiedPush_TierChangeAppliesAfterRefresh(t *testing.T) { + // Pushes are billed with the subscriber's current tier, as of the last background refresh of + // the rate visitors' users; a push never looks a user up forEachBackend(t, func(t *testing.T, databaseURL string) { c := newTestConfigWithAuthFile(t, databaseURL) c.VisitorSubscriberRateLimiting = true @@ -205,135 +310,217 @@ func TestServer_UnifiedPush_TierChangeAppliesAfterUserCache(t *testing.T) { "Authorization": util.BasicAuth("phil", "phil"), }) require.Equal(t, 200, rr.Code) + lookups := 0 + s.rateVisitorUsers.lookup = func(ids []string) ([]*user.User, error) { + lookups++ + return s.userManager.UsersByIDs(ids) + } push := func() int { return request(t, s, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil).Code } require.Equal(t, 200, push()) require.Equal(t, 429, push()) - // The new tier applies once the cached user expires + // The new tier applies once the users were refreshed require.Nil(t, s.userManager.ChangeTier("phil", "large")) require.Equal(t, 429, push()) - expireRateVisitorUsers(s) + require.Equal(t, 0, lookups) + s.execManager() + require.Equal(t, 1, lookups) require.Equal(t, 200, push()) + require.Equal(t, 1, lookups) }) } -func TestServer_UnifiedPush_RateVisitorUserIsCached(t *testing.T) { - // Pushes do not read the subscriber's user every time - c := newTestConfigWithAuthFile(t, "") - c.VisitorSubscriberRateLimiting = true - s := newTestServer(t, c) - require.Nil(t, s.userManager.AddUser("phil", "phil", user.RoleUser, false)) - rr := request(t, s, "GET", "/upAAAAAAAAAAAA/json?poll=1", "", map[string]string{ - "Authorization": util.BasicAuth("phil", "phil"), - }) - require.Equal(t, 200, rr.Code) - lookups := 0 - s.rateVisitorUsers.lookup = func(id string) (*user.User, error) { - lookups++ - return s.userManager.UserByID(id) - } - for i := 0; i < 3; i++ { - require.Equal(t, 200, request(t, s, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil).Code) - } - require.Equal(t, 1, lookups) -} - -func TestServer_UnifiedPush_DeletedSubscriberIsNoRateVisitor(t *testing.T) { - // A deleted subscriber pays for nothing, and its IP does not pay in its place: the push gets - // the 507 of a topic without subscriber +func TestServer_UnifiedPush_TierDowngradeChargesIP(t *testing.T) { + // A subscriber that loses its tier is charged on its IP like an anonymous one, with the + // anonymous limits, from the next refresh; the user key stops being charged forEachBackend(t, func(t *testing.T, databaseURL string) { c := newTestConfigWithAuthFile(t, databaseURL) c.VisitorSubscriberRateLimiting = true s := newTestServer(t, c) + require.Nil(t, s.userManager.AddTier(&user.Tier{ + Code: "pro", + MessageLimit: 10, + })) + require.Nil(t, s.userManager.AddUser("phil", "phil", user.RoleUser, false)) + require.Nil(t, s.userManager.ChangeTier("phil", "pro")) + phil, err := s.userManager.User("phil") + require.Nil(t, err) + rr := request(t, s, "GET", "/upAAAAAAAAAAAA/json?poll=1", "", map[string]string{ + "Authorization": util.BasicAuth("phil", "phil"), + }, func(r *http.Request) { + r.RemoteAddr = "1.2.3.4:1234" + }) + require.Equal(t, 200, rr.Code) + require.Equal(t, 200, request(t, s, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil).Code) + require.Equal(t, int64(1), s.visitor(netip.MustParseAddr("9.9.9.9"), phil).Stats().Messages) + require.Equal(t, int64(0), s.visitor(netip.MustParseAddr("1.2.3.4"), nil).Stats().Messages) + require.Nil(t, s.userManager.ResetTier("phil")) + s.execManager() + require.Equal(t, "ip:1.2.3.4", rateVisitorKey(s, "upAAAAAAAAAAAA")) + require.Equal(t, 200, request(t, s, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil).Code) + require.Equal(t, int64(1), s.visitor(netip.MustParseAddr("9.9.9.9"), phil).Stats().Messages) + require.Equal(t, int64(1), s.visitor(netip.MustParseAddr("1.2.3.4"), nil).Stats().Messages) + }) +} + +func TestServer_UnifiedPush_DeletedSubscriberChargesIP(t *testing.T) { + // A deleted subscriber's account pays for nothing: its IP pays, like an anonymous subscriber's + forEachBackend(t, func(t *testing.T, databaseURL string) { + c := newTestConfigWithAuthFile(t, databaseURL) + c.VisitorSubscriberRateLimiting = true + s := newTestServer(t, c) + require.Nil(t, s.userManager.AddTier(&user.Tier{ + Code: "pro", + MessageLimit: 10, + })) require.Nil(t, s.userManager.AddUser("alice", "alice", user.RoleUser, false)) require.Nil(t, s.userManager.AddUser("bob", "bob", user.RoleUser, false)) + require.Nil(t, s.userManager.ChangeTier("alice", "pro")) + require.Nil(t, s.userManager.ChangeTier("bob", "pro")) + alice, err := s.userManager.User("alice") + require.Nil(t, err) rr := request(t, s, "GET", "/upAAAAAAAAAAAA/json?poll=1", "", map[string]string{ "Authorization": util.BasicAuth("alice", "alice"), + }, func(r *http.Request) { + r.RemoteAddr = "1.2.3.4:1234" }) require.Equal(t, 200, rr.Code) rr = request(t, s, "GET", "/upBBBBBBBBBBBB/json?poll=1", "", map[string]string{ "Authorization": util.BasicAuth("bob", "bob"), + }, func(r *http.Request) { + r.RemoteAddr = "5.6.7.8:1234" }) require.Equal(t, 200, rr.Code) // Alice deletes her account (soft delete), bob is removed by an admin (hard delete) - alice, err := s.userManager.User("alice") - require.Nil(t, err) require.Nil(t, s.userManager.MarkUserRemoved(alice)) require.Nil(t, s.userManager.RemoveUser("bob")) + s.execManager() + require.Equal(t, "ip:1.2.3.4", rateVisitorKey(s, "upAAAAAAAAAAAA")) + require.Equal(t, "ip:5.6.7.8", rateVisitorKey(s, "upBBBBBBBBBBBB")) for _, topic := range []string{"upAAAAAAAAAAAA", "upBBBBBBBBBBBB"} { - response := request(t, s, "PUT", "/"+topic+"?up=1", "some message", nil) - require.Equal(t, 507, response.Code, topic) - require.Equal(t, 50701, toHTTPError(t, response.Body.String()).Code) + require.Equal(t, 200, request(t, s, "PUT", "/"+topic+"?up=1", "some message", nil).Code, topic) } + require.Equal(t, int64(1), s.visitor(netip.MustParseAddr("1.2.3.4"), nil).Stats().Messages) + require.Equal(t, int64(1), s.visitor(netip.MustParseAddr("5.6.7.8"), nil).Stats().Messages) }) } -func TestServer_UnifiedPush_RateVisitorUserLookupErrorIsRetryable(t *testing.T) { - // If the subscriber cannot be looked up, the push fails with a 5xx the publisher retries: it is - // not a topic without subscriber, which would make the Matrix gateway reject the pushkey +func TestServer_UnifiedPush_RateVisitorUserRefreshFailureKeepsCachedUsers(t *testing.T) { + // If the users cannot be refreshed, pushes keep using what was cached, rather than failing or + // falling back to the IP c := newTestConfigWithAuthFile(t, "") c.VisitorSubscriberRateLimiting = true s := newTestServer(t, c) + require.Nil(t, s.userManager.AddTier(&user.Tier{ + Code: "pro", + MessageLimit: 10, + })) require.Nil(t, s.userManager.AddUser("alice", "alice", user.RoleUser, false)) + require.Nil(t, s.userManager.ChangeTier("alice", "pro")) + alice, err := s.userManager.User("alice") + require.Nil(t, err) rr := request(t, s, "GET", "/upAAAAAAAAAAAA/json?poll=1", "", map[string]string{ "Authorization": util.BasicAuth("alice", "alice"), }) require.Equal(t, 200, rr.Code) - s.rateVisitorUsers.lookup = func(id string) (*user.User, error) { + s.rateVisitorUsers.lookup = func(ids []string) ([]*user.User, error) { return nil, errors.New("database is down") } + s.execManager() + require.Equal(t, "user:"+alice.ID, rateVisitorKey(s, "upAAAAAAAAAAAA")) response := request(t, s, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil) - require.Equal(t, 500, response.Code) + require.Equal(t, 200, response.Code) + require.Equal(t, int64(1), s.visitor(netip.MustParseAddr("9.9.9.9"), alice).Stats().Messages) } -// syncWait is how long a rate visitor takes at most to reach another node in healthy operation: one +func TestServer_UnifiedPush_RateVisitorUsersAreOnlyTheLiveOnes(t *testing.T) { + // The refresh looks up the users of the live rate visitors only, and drops the others + c := newTestConfigWithAuthFile(t, "") + c.VisitorSubscriberRateLimiting = true + s := newTestServer(t, c) + require.Nil(t, s.userManager.AddUser("alice", "alice", user.RoleUser, false)) + require.Nil(t, s.userManager.AddUser("bob", "bob", user.RoleUser, false)) + alice, err := s.userManager.User("alice") + require.Nil(t, err) + bob, err := s.userManager.User("bob") + require.Nil(t, err) + for name, topic := range map[string]string{"alice": "upAAAAAAAAAAAA", "bob": "upBBBBBBBBBBBB"} { + rr := request(t, s, "GET", "/"+topic+"/json?poll=1", "", map[string]string{ + "Authorization": util.BasicAuth(name, name), + }) + require.Equal(t, 200, rr.Code) + } + ageTopic(s, "upBBBBBBBBBBBB", 2*rateVisitorLiveFor) + var looked []string + s.rateVisitorUsers.lookup = func(ids []string) ([]*user.User, error) { + looked = ids + return s.userManager.UsersByIDs(ids) + } + s.execManager() + require.Equal(t, []string{alice.ID}, looked) + require.NotNil(t, s.rateVisitorUsers.User(alice.ID)) + require.Nil(t, s.rateVisitorUsers.User(bob.ID)) +} + +// syncWait is how long a topic takes at most to reach another node in healthy operation: one // sync round to be written, one to be read const syncWait = 10 * time.Second -// rateVisitorKey returns the key of the topic's rate visitor, or "" if it has none +// rateVisitorKey returns the limiter key the topic's pushes are charged to, or "" if the topic +// has no live rate visitor func rateVisitorKey(s *Server, topic string) string { - if p := s.limiter.RateVisitor(topic); p != nil { - return p.Key + to, err := s.topicFromID(nil, topic) + if err != nil { + return "" + } + if v := s.rateVisitor(to); v != nil { + return v.key } return "" } -// ageRateVisitor makes the topic's rate visitor (or the topic, if it has none) look as if it was last seen d ago -func ageRateVisitor(s *Server, topic string, d time.Duration) { - p := s.limiter.RateVisitor(topic) - if p == nil { - p = &limit.RateVisitor{} // Unknown: make it known, so the caller's assertion fails, not the helper +// ageTopic makes the topic look as if its subscriber activity and first observation were d +// earlier than they are; an unknown topic becomes one first seen d ago +func ageTopic(s *Server, topic string, d time.Duration) { + t := s.topicService.Get(topic) + if t == nil { + t = &tp.Topic{FirstSeenAt: time.Now()} } - p.SeenAt = time.Now().Add(-d) - s.limiter.SetRateVisitor(topic, p) -} - -// expireRateVisitorUsers makes the cached users of the rate visitors look older than rateVisitorUserCacheTTL -func expireRateVisitorUsers(s *Server) { - s.rateVisitorUsers.mu.Lock() - defer s.rateVisitorUsers.mu.Unlock() - for _, e := range s.rateVisitorUsers.entries { - e.fetched = e.fetched.Add(-2 * rateVisitorUserCacheTTL) + t.FirstSeenAt = t.FirstSeenAt.Add(-d) + if !t.SubscribedAt.IsZero() { + t.SubscribedAt = t.SubscribedAt.Add(-d) } + s.topicService.SetForTest(topic, t) } // execUserDB runs the query on the config's user database: the auth-file, or the Postgres schema func execUserDB(t *testing.T, c *Config, query string) { t.Helper() - var d *db.DB - if c.DatabaseURL != "" { - host, err := pg.Open(c.DatabaseURL) - require.Nil(t, err) - d = db.New(host, nil) - } else { - var err error - d, err = user.OpenSQLite(c.AuthFile) - require.Nil(t, err) - } + d := openUserDB(t, c) defer d.Close() _, err := d.Exec(query) require.Nil(t, err) } + +// queryUserDB runs the single-row query on the config's user database, scanning into dest +func queryUserDB(t *testing.T, c *Config, query string, dest ...any) { + t.Helper() + d := openUserDB(t, c) + defer d.Close() + require.Nil(t, d.QueryRow(query).Scan(dest...)) +} + +func openUserDB(t *testing.T, c *Config) *db.DB { + t.Helper() + if c.DatabaseURL != "" { + host, err := pg.Open(c.DatabaseURL) + require.Nil(t, err) + return db.New(host, nil) + } + d, err := user.OpenSQLite(c.AuthFile) + require.Nil(t, err) + return d +} diff --git a/server/server.go b/server/server.go index 846cb575..8e39ff55 100644 --- a/server/server.go +++ b/server/server.go @@ -43,6 +43,7 @@ import ( "heckel.io/ntfy/v2/metrics" "heckel.io/ntfy/v2/model" "heckel.io/ntfy/v2/payments" + tp "heckel.io/ntfy/v2/topic" "heckel.io/ntfy/v2/twilio" "heckel.io/ntfy/v2/user" "heckel.io/ntfy/v2/util" @@ -67,7 +68,8 @@ type Server struct { topics map[string]*topic limiter *limit.Service // Decides on and counts the visitors' usage, keyed by visitor key visitors *visitorEnv // Shared by all visitors, which are built per request - rateVisitorUsers *rateVisitorUserCache // Users of the UnifiedPush topics' rate visitors, by ID + topicService *tp.Service // Cluster-wide topic state: subscriber activity and rate visitors of the UnifiedPush topics + rateVisitorUsers *rateVisitorUserCache // Users of the live rate visitors, by ID, refreshed by the manager ban *ban.Service // Abuse ban-feed; nil when the feature is disabled (no ban file) firebaseClient *firebaseClient twilio *twilio.Client @@ -361,9 +363,12 @@ func New(conf *Config) (*Server, error) { Store: createUsageStore(pool, userDB), }) s.visitors = newVisitorEnv(conf, s.limiter, messageCache, userManager) + s.topicService = tp.New(&tp.Config{ + Store: createTopicStore(pool, userDB), + }) s.rateVisitorUsers = newRateVisitorUserCache(nil) if userManager != nil { - s.rateVisitorUsers.lookup = userManager.UserByID + s.rateVisitorUsers.lookup = userManager.UsersByIDs } s.firebaseClient = firebaseClient if pool != nil { @@ -395,6 +400,10 @@ func New(conf *Config) (*Server, error) { if err := s.limiter.Start(context.Background()); err != nil { return nil, err } + if err := s.topicService.Start(context.Background()); err != nil { + return nil, err + } + s.refreshRateVisitorUsers() // The loaded rate visitors' tiers apply from the start, not from the first manager run return s, nil } @@ -424,6 +433,17 @@ func createUsageStore(pool, userDB *db.DB) *limit.Store { return nil } +// createTopicStore returns the store of the topic table, which is part of the user schema like +// the usage table (see createUsageStore). Without a user database, the topics are kept in memory only. +func createTopicStore(pool, userDB *db.DB) *tp.Store { + if pool != nil { + return tp.NewPostgresStore(pool) + } else if userDB != nil { + return tp.NewSQLiteStore(userDB) + } + return nil +} + func createMessageCache(conf *Config, pool *db.DB) (*message.Cache, error) { if conf.CacheDuration == 0 { return message.NewNopStore() @@ -646,7 +666,8 @@ func (s *Server) stop() { if s.cluster != nil { s.cluster.Close() // Before the databases close, so it can deregister the node and its heartbeat stops } - s.limiter.Close() // Writes the unflushed daily counts, before the databases close under it + s.limiter.Close() // Writes the unflushed daily counts, before the databases close under it + s.topicService.Close() // Likewise, the changed topics s.closeDatabases() if s.ban != nil { s.ban.Close() @@ -1110,10 +1131,10 @@ func (s *Server) handlePublishInternal(r *http.Request, v *visitor) (*model.Mess // the subscription as invalid if any 400-499 code (except 429/408) is returned. // See https://github.com/mastodon/mastodon/blob/730bb3e211a84a2f30e3e2bbeae3f77149824a68/app/workers/web/push_notification_worker.rb#L35-L46 // That is what a topic abandoned for unifiedPushAbandonedAfter gets: a 404, so the subscription is removed. - if s.unifiedPushTopicAbandoned(t) { + if s.unifiedPushTopicAbandoned(r, t) { return nil, errHTTPNotFoundUnifiedPushTopicAbandoned.With(t) } else if isUnifiedPushTopic(t.ID) { - s.limiter.NoteTopic(t.ID) // Starts the countdown to the 404 of a topic nobody ever subscribed to + s.topicService.Note(t.ID) // Starts the countdown to the 404 of a topic nobody ever subscribed to } return nil, errHTTPInsufficientStorageUnifiedPush.With(t) } else if !util.ContainsIP(s.config.VisitorRequestExemptPrefixes, v.ip) && !vrate.MessageAllowed() { @@ -1227,24 +1248,15 @@ func (s *Server) handlePublishMatrix(w http.ResponseWriter, r *http.Request, v * if err != nil { metrics.MessagesPublishedFailure.Inc() metrics.MatrixPublishedFailure.Inc() + // An abandoned topic (see unifiedPushTopicAbandoned) gets its pushkey rejected, so the Matrix + // server stops sending to it; everything else, including a topic nobody subscribed to yet, is + // the error as it is (a 507 is transient to the Matrix server) if e, ok := err.(*errHTTP); ok && e.Code == errHTTPNotFoundUnifiedPushTopicAbandoned.Code { pushKey, err := fromContext[string](r, contextMatrixPushKey) if err != nil { return err } return writeMatrixResponse(w, pushKey) - } else if ok && e.HTTPCode == errHTTPInsufficientStorageUnifiedPush.HTTPCode { - topic, err := fromContext[*topic](r, contextTopic) - if err != nil { - return err - } - pushKey, err := fromContext[string](r, contextMatrixPushKey) - if err != nil { - return err - } - if time.Since(topic.LastAccess()) > matrixRejectPushKeyForUnifiedPushTopicWithoutRateVisitorAfter { - return writeMatrixResponse(w, pushKey) - } } return err } @@ -1755,7 +1767,7 @@ func (s *Server) handleSubscribeHTTP(w http.ResponseWriter, r *http.Request, v * for i, subscriberID := range subscriberIDs { topics[i].Unsubscribe(subscriberID) // Order! } - s.touchRateVisitors(topics) // The subscriber was last seen now + s.touchTopics(topics) // The subscriber was last seen now }() if err := s.authorizeSubscribed(v, topics); err != nil { return err @@ -1915,7 +1927,7 @@ func (s *Server) handleSubscribeWS(w http.ResponseWriter, r *http.Request, v *vi for i, subscriberID := range subscriberIDs { topics[i].Unsubscribe(subscriberID) // Order! } - s.touchRateVisitors(topics) // The subscriber was last seen now + s.touchTopics(topics) // The subscriber was last seen now }() if err := s.authorizeSubscribed(v, topics); err != nil { return err @@ -1951,7 +1963,7 @@ func parseSubscribeParams(r *http.Request) (poll bool, since model.SinceMarker, return } -// maybeSetRateVisitors makes v the rate visitor of a topic (limit.Service.SetRateVisitor), indicating that all messages published +// maybeSetRateVisitors makes v the rate visitor of a topic (topic.Service.SetRateVisitor), indicating that all messages published // to that topic will be rate limited against the rate visitor instead of the publishing visitor. // // Setting the rate visitor is ony allowed if the `visitor-subscriber-rate-limiting` setting is enabled, AND @@ -1976,7 +1988,7 @@ func (s *Server) maybeSetRateVisitors(r *http.Request, v *visitor, topics []*top if len(eligibleRateTopics) == 0 { return nil } - s.touchRateVisitors(eligibleRateTopics) + s.touchTopics(eligibleRateTopics) // If access controls are turned off, v has access to everything, and we can set the rate visitor if s.userManager == nil { @@ -2012,11 +2024,13 @@ func (s *Server) setRateVisitors(r *http.Request, v *visitor, rateTopics []*topi Tag(tagSubscribe). With(t). Debug("Setting visitor as rate visitor for topic %s", t.ID) - s.limiter.SetRateVisitor(t.ID, &limit.RateVisitor{ - Key: v.key, + s.topicService.SetRateVisitor(t.ID, &tp.RateVisitor{ UserID: v.MaybeUserID(), - SeenAt: time.Now(), + IP: v.maskedIP(), }) + if u := v.User(); u != nil { + s.rateVisitorUsers.Put(u) // Its tier applies to the pushes right away, not from the next refresh + } t.Keepalive() } return nil diff --git a/server/server_manager.go b/server/server_manager.go index 8fa25c49..2a3b00f3 100644 --- a/server/server_manager.go +++ b/server/server_manager.go @@ -14,7 +14,7 @@ func (s *Server) execManager() { // database (and the web push job, which also sends expiry warnings) run on the cluster // leader only. Without clustering, IsLeader is always true. s.pruneVisitors() - s.rateVisitorUsers.Prune() + s.refreshRateVisitorUsers() if s.cluster.IsLeader() { s.pruneTokens() s.pruneAttachments() @@ -50,7 +50,7 @@ func (s *Server) execManager() { } t.ExpireDelivered() if subs > 0 { - s.touchRateVisitors([]*topic{t}) // An open stream keeps the rate visitor live on the other nodes too + s.keepaliveTopics([]*topic{t}) // An open stream keeps the rate visitor live on the other nodes too } subscribers += subs } diff --git a/server/server_matrix.go b/server/server_matrix.go index 79b168a1..028b5a29 100644 --- a/server/server_matrix.go +++ b/server/server_matrix.go @@ -8,7 +8,6 @@ import ( "io" "net/http" "strings" - "time" ) // Matrix Push Gateway / UnifiedPush / ntfy integration: @@ -72,14 +71,6 @@ type matrixResponse struct { Rejected []string `json:"rejected"` } -const ( - // matrixRejectPushKeyForUnifiedPushTopicWithoutRateVisitorAfter is the time after which a Matrix response - // will return an HTTP 200 with the push key (i.e. "rejected":[""]}), if no rate visitor has been set on - // the topic. Rejecting the push key will instruct the Matrix server to invalidate the pushkey and stop sending - // messages to it. This must be longer than topicExpungeAfter. See https://spec.matrix.org/v1.6/push-gateway-api/ - matrixRejectPushKeyForUnifiedPushTopicWithoutRateVisitorAfter = 12 * time.Hour -) - // errMatrixPushkeyRejected represents an error when handing Matrix gateway messages // // If the push key is set, the app server will remove it and will never send messages using the same diff --git a/server/server_middleware.go b/server/server_middleware.go index 21cbe907..f21f231a 100644 --- a/server/server_middleware.go +++ b/server/server_middleware.go @@ -36,10 +36,7 @@ func (s *Server) limitRequestsWithTopic(next handleFunc) handleFunc { return err } vrate := v - rateVisitor, err := s.rateVisitor(t) // Once per request, so the publish handler sees the same rate visitor - if err != nil { - return err - } else if rateVisitor != nil { + if rateVisitor := s.rateVisitor(t); rateVisitor != nil { // Once per request, so the publish handler sees the same rate visitor vrate = rateVisitor } r = withContext(r, map[contextKey]any{ diff --git a/server/server_test.go b/server/server_test.go index 8330f284..a7770e16 100644 --- a/server/server_test.go +++ b/server/server_test.go @@ -1966,6 +1966,8 @@ func TestServer_MatrixGateway_Push_Failure_NoSubscriber(t *testing.T) { } func TestServer_MatrixGateway_Push_Failure_NoSubscriber_After13Hours(t *testing.T) { + // A topic nobody subscribed to is a 507 for however long this node has known it: the pushkey is + // only rejected after unifiedPushAbandonedAfter, see TestServer_MatrixGateway_AbandonedTopicRejectsPushKey forEachBackend(t, func(t *testing.T, databaseURL string) { c := newTestConfig(t, databaseURL) c.VisitorSubscriberRateLimiting = true @@ -1976,15 +1978,12 @@ func TestServer_MatrixGateway_Push_Failure_NoSubscriber_After13Hours(t *testing. response := request(t, s, "POST", "/_matrix/push/v1/notify", notification, nil) require.Equal(t, 507, response.Code) require.Equal(t, 50701, toHTTPError(t, response.Body.String()).Code) - require.Nil(t, s.limiter.RateVisitor("mytopic")) + require.Nil(t, s.topicService.Get("mytopic")) // Fake: This topic has been around for 13 hours without a rate visitor s.topics["mytopic"].lastAccess = time.Now().Add(-13 * time.Hour) - - // Same request should now return HTTP 200 with a rejected pushkey response = request(t, s, "POST", "/_matrix/push/v1/notify", notification, nil) - require.Equal(t, 200, response.Code) - require.Equal(t, `{"rejected":["http://127.0.0.1:12345/mytopic?up=1"]}`, strings.TrimSpace(response.Body.String())) + require.Equal(t, 507, response.Code) // Slightly unrelated: Test that topic is pruned after 16 hours s.topics["mytopic"].lastAccess = time.Now().Add(-17 * time.Hour) @@ -3446,7 +3445,7 @@ func TestServer_SubscriberRateLimiting_VisitorExpiration(t *testing.T) { require.Equal(t, int64(1), s.visitor(netip.MustParseAddr("1.2.3.4"), nil).Stats().Messages) // Expire visitor: its subscriber has been gone for a day, and its key is pruned after a day without use - ageRateVisitor(s, "upAAAAAAAAAAAA", 25*time.Hour) + ageTopic(s, "upAAAAAAAAAAAA", 25*time.Hour) s.limiter.Prune(-time.Hour) // Everything is idle for longer than -1h // Publish message again, observe that the rate visitor is not used anymore @@ -3483,7 +3482,7 @@ func TestServer_SubscriberRateLimiting_PollKeepsStreamAsRateVisitor(t *testing.T require.Equal(t, 200, rr.Code) // A day later, the stream is still there and keeps the topic's rate visitor - ageRateVisitor(s, "upAAAAAAAAAAAA", 25*time.Hour) + ageTopic(s, "upAAAAAAAAAAAA", 25*time.Hour) response := request(t, s, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil) require.Equal(t, 200, response.Code) } @@ -3500,10 +3499,10 @@ func TestServer_SubscriberRateLimiting_PollKeepsPollingSubscriberAlive(t *testin require.Equal(t, 200, rr.Code) // Polled 23 hours ago; polling again keeps it the rate visitor for another day - ageRateVisitor(s, "upAAAAAAAAAAAA", 23*time.Hour) + ageTopic(s, "upAAAAAAAAAAAA", 23*time.Hour) rr = request(t, s, "GET", "/upAAAAAAAAAAAA/json?poll=1", "", nil, subscriberFn) require.Equal(t, 200, rr.Code) - require.WithinDuration(t, time.Now(), s.limiter.RateVisitor("upAAAAAAAAAAAA").SeenAt, time.Minute) + require.WithinDuration(t, time.Now(), s.topicService.Get("upAAAAAAAAAAAA").SubscribedAt, time.Minute) } func TestServer_SubscriberRateLimiting_OverlappingStreamsKeepTopicAlive(t *testing.T) { @@ -3543,7 +3542,7 @@ func TestServer_SubscriberRateLimiting_OverlappingStreamsKeepTopicAlive(t *testi <-doneB // A day later, stream A is still there - ageRateVisitor(s, "upAAAAAAAAAAAA", 25*time.Hour) + ageTopic(s, "upAAAAAAAAAAAA", 25*time.Hour) response := request(t, s, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil) require.Equal(t, 200, response.Code) } @@ -3572,7 +3571,7 @@ func TestServer_SubscriberRateLimiting_OpenStreamKeepsPollerAsRateVisitor(t *tes require.Equal(t, 200, rr.Code) // A day later, only the stream is still there; the publish is billed to the poller's key - ageRateVisitor(s, "upAAAAAAAAAAAA", 25*time.Hour) + ageTopic(s, "upAAAAAAAAAAAA", 25*time.Hour) response := request(t, s, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil) require.Equal(t, 200, response.Code) require.Equal(t, int64(1), s.visitor(netip.MustParseAddr("1.2.3.4"), nil).Stats().Messages) @@ -3611,9 +3610,7 @@ func TestServer_SubscriberRateLimiting_StreamDoesNotReclaimTopicAfterWriteRevoke to := lockedTopic(s, "upAAAAAAAAAAAA") return to != nil && to.SubscribersCount() == 1 }) - rateVisitor, err := s.rateVisitor(lockedTopic(s, "upAAAAAAAAAAAA")) - require.Nil(t, err) - require.True(t, rateVisitor.MessageAllowed()) + require.True(t, s.rateVisitor(lockedTopic(s, "upAAAAAAAAAAAA")).MessageAllowed()) require.Nil(t, s.userManager.AllowAccess("alice", "upAAAAAAAAAAAA", user.PermissionRead)) // Bob takes the topic with a poll, and is gone a day later @@ -3621,7 +3618,7 @@ func TestServer_SubscriberRateLimiting_StreamDoesNotReclaimTopicAfterWriteRevoke r.RemoteAddr = "2.2.2.2:1234" }) require.Equal(t, 200, response.Code) - ageRateVisitor(s, "upAAAAAAAAAAAA", 25*time.Hour) + ageTopic(s, "upAAAAAAAAAAAA", 25*time.Hour) // Alice's open stream keeps the topic alive, but it is still billed to Bob response = request(t, s, "PUT", "/upAAAAAAAAAAAA", "some message", bobHeaders, func(r *http.Request) { @@ -3649,7 +3646,7 @@ func TestServer_SubscriberRateLimiting_OpenStreamKeepsTopicAliveWithoutKeepalive to := lockedTopic(s, "upAAAAAAAAAAAA") return to != nil && to.SubscribersCount() == 1 }) - ageRateVisitor(s, "upAAAAAAAAAAAA", 25*time.Hour) + ageTopic(s, "upAAAAAAAAAAAA", 25*time.Hour) response := request(t, s, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil) require.Equal(t, 200, response.Code) require.Equal(t, int64(1), s.visitor(netip.MustParseAddr("1.2.3.4"), nil).Stats().Messages) @@ -3675,14 +3672,14 @@ func TestServer_SubscriberRateLimiting_ClosedStreamIsLastSeen(t *testing.T) { to := lockedTopic(s, "upAAAAAAAAAAAA") return to != nil && to.SubscribersCount() == 1 }) - ageRateVisitor(s, "upAAAAAAAAAAAA", 25*time.Hour) // Subscribed more than a day ago + ageTopic(s, "upAAAAAAAAAAAA", 25*time.Hour) // Subscribed more than a day ago cancel() <-done response := request(t, s, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil) require.Equal(t, 200, response.Code) // A day after the stream closed, the subscriber is gone - ageRateVisitor(s, "upAAAAAAAAAAAA", 25*time.Hour) + ageTopic(s, "upAAAAAAAAAAAA", 25*time.Hour) response = request(t, s, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil) require.Equal(t, 507, response.Code) } diff --git a/server/topic.go b/server/topic.go index a81d4581..b46399ef 100644 --- a/server/topic.go +++ b/server/topic.go @@ -12,9 +12,9 @@ import ( ) const ( - // topicExpungeAfter defines how long a topic is active before it is removed from memory. - // This must be larger than matrixRejectPushKeyForUnifiedPushTopicWithoutRateVisitorAfter to give - // time for more requests to come in, so that we can send a {"rejected":[""]} response back. + // topicExpungeAfter defines how long a topic is active before it is removed from memory. What + // has to outlive it (a UnifiedPush topic's subscriber activity and rate visitor) is kept by + // the topic service, not here. topicExpungeAfter = 16 * time.Hour // topicDeliveredRetain is how long a topic at least remembers that a message went out to its diff --git a/server/visitor.go b/server/visitor.go index d7243053..fd626ebb 100644 --- a/server/visitor.go +++ b/server/visitor.go @@ -17,9 +17,8 @@ const ( oneDay = 24 * time.Hour // visitorExpungeAfter is how long a limiter key may go unused before its buckets and counts are - // forgotten (see Server.pruneVisitors), and how long a topic keeps a rate visitor whose subscriber is - // gone (see Server.rateVisitor). 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). + // forgotten (see Server.pruneVisitors). 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 = oneDay // promoteRate is the divisor of the daily message cap (cap / promoteRate, i.e. 10%) from which @@ -93,35 +92,21 @@ func newVisitorEnv(conf *Config, limiter *limit.Service, messageCache *message.C } func newVisitor(env *visitorEnv, ip netip.Addr, u *user.User) *visitor { - v := newVisitorWithKey(env, vis.Key(ip, u, env.limitsConfig), u) - v.ip = ip - env.limiter.Touch(v.key) // The key is seen now, as the visitor's request is - return v -} - -// newRateVisitor returns the visitor of a UnifiedPush topic's rate visitor (see Server.rateVisitor): -// the subscriber's key, with its user's current limits. A user with a tier now is keyed by the -// user, as vis.Key would, even if it had none when it subscribed. Its IP is unknown. -func newRateVisitor(env *visitorEnv, key string, u *user.User) *visitor { - if u != nil && u.Tier != nil { - key = vis.Key(netip.Addr{}, u, env.limitsConfig) - } - return newVisitorWithKey(env, key, u) -} - -// newVisitorWithKey returns a visitor with the given key, and the limits of the user's tier, or -// the config's without one -func newVisitorWithKey(env *visitorEnv, key string, u *user.User) *visitor { limits := env.ipLimits + var tier *user.Tier if u != nil && u.Tier != nil { - limits = vis.LimitsFor(env.limitsConfig, u.Tier) // Resolved per request: the tier's fields may change at any time + tier = u.Tier + limits = vis.LimitsFor(env.limitsConfig, tier) // Resolved per request: the tier's fields may change at any time } - return &visitor{ + v := &visitor{ env: env, - key: key, + key: vis.Key(ip, u, env.limitsConfig), + ip: ip, user: u, limits: limits, } + env.limiter.Touch(v.key) // The key is seen now, as the visitor's request is + return v } func (v *visitor) Context() log.Context { diff --git a/visitor/key.go b/visitor/key.go index 322b8af3..107263f2 100644 --- a/visitor/key.go +++ b/visitor/key.go @@ -7,16 +7,22 @@ import ( ) // Key returns the identifier a visitor's limits are keyed by: "user:" for a user with a tier, -// or "ip:" otherwise, with the IP masked to the configured prefix bits. A user without a -// tier shares the IP key with everyone behind the same prefix. +// or "ip:" otherwise, with the IP masked to the configured prefix bits (see Mask). A user +// without a tier shares the IP key with everyone behind the same prefix. func Key(ip netip.Addr, u *user.User, conf *Config) string { if u != nil && u.Tier != nil { return "user:" + u.ID } + return "ip:" + Mask(ip, conf).String() +} + +// Mask returns the IP masked to the configured prefix bits: what the visitors behind one prefix +// share, and what is kept of a subscriber's IP (see package topic). Masking twice is harmless. +func Mask(ip netip.Addr, conf *Config) netip.Addr { if ip.Is4() { - ip = netip.PrefixFrom(ip, conf.PrefixBitsIPv4).Masked().Addr() + return netip.PrefixFrom(ip, conf.PrefixBitsIPv4).Masked().Addr() } else if ip.Is6() { - ip = netip.PrefixFrom(ip, conf.PrefixBitsIPv6).Masked().Addr() + return netip.PrefixFrom(ip, conf.PrefixBitsIPv6).Masked().Addr() } - return "ip:" + ip.String() + return ip } diff --git a/visitor/key_test.go b/visitor/key_test.go index 599b51ab..11ce1738 100644 --- a/visitor/key_test.go +++ b/visitor/key_test.go @@ -37,3 +37,13 @@ func TestKey(t *testing.T) { require.Equal(t, "ip:1.2.0.0", Key(netip.MustParseAddr("1.2.3.4"), nil, confWithShortenedPrefixes)) require.Equal(t, "ip:2a01:599:b26:2300::", Key(netip.MustParseAddr("2a01:599:b26:2397:dbe7:5aa2:95ce:1e83"), nil, confWithShortenedPrefixes)) } + +func TestMask(t *testing.T) { + conf := &Config{ + PrefixBitsIPv4: 16, + PrefixBitsIPv6: 56, + } + require.Equal(t, "1.2.0.0", Mask(netip.MustParseAddr("1.2.3.4"), conf).String()) + require.Equal(t, "2a01:599:b26:2300::", Mask(netip.MustParseAddr("2a01:599:b26:2397:dbe7:5aa2:95ce:1e83"), conf).String()) + require.Equal(t, "1.2.0.0", Mask(netip.MustParseAddr("1.2.0.0"), conf).String()) // Masking again changes nothing +}