mtfyun/ntfy
1
0 Fork 0
mirror of https://github.com/binwiederhier/ntfy.git synced 2026-11-10 14:44:04 +08:00
Code Issues Packages Projects Releases Wiki Activity

Keep the UnifiedPush rate visitors in the topic service, derive the limiter key at push time

Browse Source 
The server uses the topic package instead of the limit service for the rate
visitors: a subscribe sets the rate visitor (user ID and masked IP), a poll,
a stream close and the manager's keepalive record subscriber activity, a
refused push notes an unknown topic (insert-only). The limiter key is derived
on every push from the rate visitor's current tier (refreshed in the
background every minute, never on the publish path), so a downgrade or a
deleted account charges the IP. Before a 404 for an abandoned topic the row
is read from the database. The Matrix gateway rejects the pushkey under the
same 7-day rule; its 12 h per-node rule is gone.
This commit is contained in:
binwiederhier committed 2026-10-31 14:05:37 +00:00
1 parent 14f631da8f
commit ac9c6b2bac
19 files changed
+500 -882
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

limit/rate_visitor.go
-169
View File
@@ -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,
}
}
limit/rate_visitor_test.go
-183
View File
@@ -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,
}
}
limit/service.go
+6 -37
View File
@@ -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.
@@ -67,7 +64,6 @@ type Service struct {
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
@@ -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)
}
}
}
limit/store.go
+1 -59
View File
@@ -32,15 +32,6 @@ 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.
@@ -48,7 +39,6 @@ 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
}
// queries are the dialect's SQL statements
@@ -56,9 +46,6 @@ type queries struct {
selectSince string
upsert string
prune string
selectRateVisitorsSince string
upsertRateVisitor string
pruneRateVisitors 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)
}
limit/store_postgres.go
-32
View File
@@ -28,23 +28,6 @@ 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 {
@@ -54,13 +37,9 @@ func NewPostgresStore(d *db.DB) *Store {
selectSince: postgresSelectSinceQuery,
upsert: postgresUpsertQuery,
prune: postgresPruneQuery,
selectRateVisitorsSince: postgresSelectRateVisitorsSinceQuery,
upsertRateVisitor: postgresUpsertRateVisitorQuery,
pruneRateVisitors: postgresPruneRateVisitorsQuery,
},
}
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.
limit/store_sqlite.go
-36
View File
@@ -24,21 +24,6 @@ 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 {
@@ -48,13 +33,9 @@ func NewSQLiteStore(d *db.DB) *Store {
selectSince: sqliteSelectSinceQuery,
upsert: sqliteUpsertQuery,
prune: sqlitePruneQuery,
selectRateVisitorsSince: sqliteSelectRateVisitorsSinceQuery,
upsertRateVisitor: sqliteUpsertRateVisitorQuery,
pruneRateVisitors: sqlitePruneRateVisitorsQuery,
},
}
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
})
}
limit/store_test.go
-81
View File
@@ -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:<id>") 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,
}
}
limit/types.go
-31
View File
@@ -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 {
server/rate_visitor.go
+125 -73
View File
@@ -1,137 +1,189 @@
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),
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
// UnifiedPush length
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()
}
server/rate_visitor_test.go
+279 -92
View File
@@ -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,7 +112,8 @@ 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, "")
forEachBackend(t, func(t *testing.T, databaseURL string) {
c := newTestConfigWithAuthFile(t, databaseURL)
c.VisitorSubscriberRateLimiting = true
s := newTestServer(t, c)
push := func() *errHTTP {
@@ -88,9 +126,9 @@ func TestServer_UnifiedPush_AbandonedTopicAnswers404(t *testing.T) {
// Never subscribed: 507 for a week from the first push, then 404
require.Equal(t, 50701, push().Code)
ageRateVisitor(s, "upAAAAAAAAAAAA", unifiedPushAbandonedAfter-time.Hour)
ageTopic(s, "upAAAAAAAAAAAA", unifiedPushAbandonedAfter-time.Hour)
require.Equal(t, 50701, push().Code)
ageRateVisitor(s, "upAAAAAAAAAAAA", unifiedPushAbandonedAfter+time.Hour)
ageTopic(s, "upAAAAAAAAAAAA", 2*time.Hour)
require.Equal(t, 40402, push().Code)
require.Equal(t, 404, push().HTTPCode)
@@ -100,9 +138,9 @@ func TestServer_UnifiedPush_AbandonedTopicAnswers404(t *testing.T) {
})
require.Equal(t, 200, rr.Code)
require.Nil(t, push())
ageRateVisitor(s, "upAAAAAAAAAAAA", 2*oneDay)
ageTopic(s, "upAAAAAAAAAAAA", 2*oneDay)
require.Equal(t, 50701, push().Code)
ageRateVisitor(s, "upAAAAAAAAAAAA", unifiedPushAbandonedAfter+time.Hour)
ageTopic(s, "upAAAAAAAAAAAA", unifiedPushAbandonedAfter)
require.Equal(t, 40402, push().Code)
// An open stream keeps the topic, however long ago it was seen otherwise
@@ -117,8 +155,49 @@ func TestServer_UnifiedPush_AbandonedTopicAnswers404(t *testing.T) {
waitFor(t, func() bool {
return lockedTopic(s, "upAAAAAAAAAAAA").SubscribersCount() == 1
})
ageRateVisitor(s, "upAAAAAAAAAAAA", unifiedPushAbandonedAfter+time.Hour)
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)
})
require.Equal(t, 507, response.Code)
response = request(t, s, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil)
require.Equal(t, 404, response.Code)
})
}
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
ageTopic(s, "upAAAAAAAAAAAA", 2*rateVisitorLiveFor) // 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)
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, 200, push())
})
}
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, 0, lookups)
s.execManager()
require.Equal(t, 1, lookups)
require.Equal(t, 200, push())
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
}
server/server.go
+37 -23
View File
@@ -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()
@@ -647,6 +667,7 @@ func (s *Server) stop() {
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.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
server/server_manager.go
+2 -2
View File
@@ -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
}
server/server_matrix.go
-9
View File
@@ -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":["<pushkey>"]}), 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
server/server_middleware.go
+1 -4
View File
@@ -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{
server/server_test.go
+15 -18
View File
@@ -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)
}
server/topic.go
+3 -3
View File
@@ -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":["<pushkey>"]} 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
server/visitor.go
+10 -25
View File
@@ -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 {
visitor/key.go
+11 -5
View File
@@ -7,16 +7,22 @@ import (
)
// Key returns the identifier a visitor's limits are keyed by: "user:<id>" for a user with a tier,
// or "ip:<prefix>" 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:<prefix>" 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
}
visitor/key_test.go
+10
View File
@@ -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
}