mirror of
https://github.com/multipleof4/ntfy.git
synced 2026-10-18 01:17:53 +00:00
Redo the banning logic, ban.Service
This commit is contained in:
+185
@@ -0,0 +1,185 @@
|
||||
// Package ban implements the abuse ban-feed: it tracks per-prefix weighted "strikes" from rejected
|
||||
// requests and appends breaching prefixes to a ban file that fail2ban tails. Keying by prefix (not
|
||||
// by visitor) makes the accounting match the unit fail2ban bans, even for shared account visitors.
|
||||
package ban
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/netip"
|
||||
"os"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"golang.org/x/time/rate"
|
||||
"heckel.io/ntfy/v2/log"
|
||||
)
|
||||
|
||||
const (
|
||||
tag = "ban"
|
||||
pruneInterval = 10 * time.Minute
|
||||
writeInterval = 3 * time.Second
|
||||
)
|
||||
|
||||
// Config is the Service's config, kept separate from server.Config to avoid an import cycle.
|
||||
type Config struct {
|
||||
File string // Ban file that fail2ban tails (must be non-empty; the caller decides whether the feature is enabled)
|
||||
Window time.Duration // Rolling window over which weighted strikes are counted
|
||||
Threshold int // Weighted strikes per Window before a prefix is banned
|
||||
Weights Weights // Code matcher -> strike weight (0 = exempt)
|
||||
PrefixBitsIPv4 int // Mask width for the ban unit, e.g. 32 (matches rate-limiting granularity)
|
||||
PrefixBitsIPv6 int // Mask width for the ban unit, e.g. 64
|
||||
}
|
||||
|
||||
// tracker is the per-prefix strike state: a weighted breach detector plus timestamps for pruning and throttling.
|
||||
type tracker struct {
|
||||
limiter *rate.Limiter
|
||||
seen time.Time // Last strike, for pruning idle prefixes
|
||||
emitted time.Time // Last ban-line write for this prefix, throttles re-emits to once per Window
|
||||
}
|
||||
|
||||
// Service owns the ban-feed: per-prefix strike accounting, buffered file writes, and idle-prefix
|
||||
// pruning. The caller owns the enable/disable decision -- only construct a Service when the feature
|
||||
// is on (see server.New, which builds one only when a ban file is configured).
|
||||
type Service struct {
|
||||
conf *Config
|
||||
mu sync.Mutex // Guards trackers and pending
|
||||
trackers map[netip.Prefix]*tracker
|
||||
pending []string // Formatted ban lines buffered by Record, flushed to the ban file by runWriteLoop
|
||||
writeDone chan struct{} // Closed when runWriteLoop exits after its final flush
|
||||
closeChan chan struct{}
|
||||
closeOnce sync.Once
|
||||
}
|
||||
|
||||
// NewService builds a Service and starts its background loops. The caller must only call it when the
|
||||
// feature is enabled (conf non-nil, File non-empty); the Service does not model a disabled state.
|
||||
func NewService(conf *Config) *Service {
|
||||
s := &Service{
|
||||
conf: conf,
|
||||
trackers: make(map[netip.Prefix]*tracker),
|
||||
closeChan: make(chan struct{}),
|
||||
writeDone: make(chan struct{}),
|
||||
}
|
||||
go s.runPruneLoop()
|
||||
go s.runWriteLoop()
|
||||
return s
|
||||
}
|
||||
|
||||
// Record counts one rejection against the IP's prefix bucket and, on breach, buffers a ban line
|
||||
// (throttled to once per Window per prefix). No-ops for a non-4xx/5xx status or a zero-weight code.
|
||||
func (s *Service) Record(ip netip.Addr, httpCode, errorCode int) {
|
||||
if httpCode < 400 {
|
||||
return
|
||||
}
|
||||
weight := s.conf.Weights.WeightFor(errorCode)
|
||||
if weight == 0 {
|
||||
return // Weight 0: exempt, no strike
|
||||
}
|
||||
prefix := s.prefix(ip)
|
||||
now := time.Now()
|
||||
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
t := s.trackers[prefix]
|
||||
if t == nil {
|
||||
t = &tracker{limiter: rate.NewLimiter(rate.Limit(float64(s.conf.Threshold)/s.conf.Window.Seconds()), s.conf.Threshold)}
|
||||
s.trackers[prefix] = t
|
||||
}
|
||||
t.seen = now
|
||||
if t.limiter.AllowN(now, weight) {
|
||||
return // Within the strike budget, no breach
|
||||
}
|
||||
if !t.emitted.IsZero() && now.Sub(t.emitted) < s.conf.Window {
|
||||
return // Already emitted this prefix within the window (one ban line per prefix per window)
|
||||
}
|
||||
t.emitted = now
|
||||
s.pending = append(s.pending, formatBanLine(now, ip, prefix, httpCode, errorCode))
|
||||
}
|
||||
|
||||
// prefix masks ip to the ban unit (PrefixBitsIPv4/IPv6) -- what fail2ban bans, e.g. a whole /64.
|
||||
func (s *Service) prefix(ip netip.Addr) netip.Prefix {
|
||||
if ip.Is4() {
|
||||
return netip.PrefixFrom(ip, s.conf.PrefixBitsIPv4).Masked()
|
||||
}
|
||||
return netip.PrefixFrom(ip, s.conf.PrefixBitsIPv6).Masked()
|
||||
}
|
||||
|
||||
// formatBanLine builds the "<RFC3339-UTC> <ip> <prefix> <http> <ntfy>" line the fail2ban filter
|
||||
// parses. The timestamp is captured at breach time, not flush time.
|
||||
func formatBanLine(t time.Time, ip netip.Addr, prefix netip.Prefix, httpCode, errorCode int) string {
|
||||
return fmt.Sprintf("%s %s %s %d %d\n", t.UTC().Format(time.RFC3339), ip.String(), prefix.String(), httpCode, errorCode)
|
||||
}
|
||||
|
||||
// runPruneLoop prunes idle prefixes until Close.
|
||||
func (s *Service) runPruneLoop() {
|
||||
ticker := time.NewTicker(pruneInterval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ticker.C:
|
||||
s.prune()
|
||||
case <-s.closeChan:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// runWriteLoop flushes buffered ban lines every writeInterval, plus a final flush on Close.
|
||||
func (s *Service) runWriteLoop() {
|
||||
defer close(s.writeDone)
|
||||
ticker := time.NewTicker(writeInterval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ticker.C:
|
||||
s.flush()
|
||||
case <-s.closeChan:
|
||||
s.flush()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// flush appends the buffered lines to the file in one open/write. Best-effort: a batch is dropped on
|
||||
// error. Concurrent calls are safe -- pending is drained under mu, so only one flush writes a batch.
|
||||
func (s *Service) flush() {
|
||||
s.mu.Lock()
|
||||
lines := s.pending
|
||||
s.pending = nil
|
||||
s.mu.Unlock()
|
||||
if len(lines) == 0 {
|
||||
return
|
||||
}
|
||||
f, err := os.OpenFile(s.conf.File, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
|
||||
if err != nil {
|
||||
log.Tag(tag).Err(err).Warn("Cannot open ban file %s, dropped %d ban(s)", s.conf.File, len(lines))
|
||||
return
|
||||
}
|
||||
defer f.Close()
|
||||
for i, line := range lines {
|
||||
if _, err := f.WriteString(line); err != nil {
|
||||
log.Tag(tag).Err(err).Warn("Cannot write to ban file %s, dropped %d ban(s)", s.conf.File, len(lines)-i)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// prune drops prefixes idle for a full Window -- their bucket has refilled, so forgetting them is a
|
||||
// no-op that bounds memory under a flood of distinct IPs.
|
||||
func (s *Service) prune() {
|
||||
now := time.Now()
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
for prefix, t := range s.trackers {
|
||||
if now.Sub(t.seen) >= s.conf.Window {
|
||||
delete(s.trackers, prefix)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Close stops the loops and blocks until the final flush completes. Idempotent.
|
||||
func (s *Service) Close() {
|
||||
s.closeOnce.Do(func() {
|
||||
close(s.closeChan)
|
||||
<-s.writeDone
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,253 @@
|
||||
package ban
|
||||
|
||||
import (
|
||||
"net/netip"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
var testIP = netip.MustParseAddr("1.2.3.4")
|
||||
|
||||
// newTestService creates a Service wired for testing, with the given ban file, weighted-bucket
|
||||
// threshold, and weights, plus a 1-minute window (so the emit throttle only fires once per test).
|
||||
func newTestService(t *testing.T, banFile string, threshold int, weights map[string]int) *Service {
|
||||
t.Helper()
|
||||
s := NewService(&Config{
|
||||
File: banFile,
|
||||
Window: time.Minute,
|
||||
Threshold: threshold,
|
||||
Weights: weights,
|
||||
PrefixBitsIPv4: 32,
|
||||
PrefixBitsIPv6: 64,
|
||||
})
|
||||
t.Cleanup(s.Close)
|
||||
return s
|
||||
}
|
||||
|
||||
// flushAndRead forces a synchronous flush of the buffered bans (writes are otherwise async, on the
|
||||
// runWriteLoop ticker) and returns the ban file's lines.
|
||||
func flushAndRead(t *testing.T, s *Service, path string) []string {
|
||||
t.Helper()
|
||||
s.flush()
|
||||
data, err := os.ReadFile(path)
|
||||
require.NoError(t, err)
|
||||
return strings.Split(strings.TrimRight(string(data), "\n"), "\n")
|
||||
}
|
||||
|
||||
func TestService_Record_Weight2BansAtHalfThreshold(t *testing.T) {
|
||||
banFile := filepath.Join(t.TempDir(), "ban.log")
|
||||
// Threshold 10, code weight 2 -> the budget covers exactly 5 hits, so the 6th breaches.
|
||||
s := newTestService(t, banFile, 10, map[string]int{"*": 2})
|
||||
for i := 0; i < 5; i++ {
|
||||
s.Record(testIP, 400, 40001)
|
||||
}
|
||||
s.flush()
|
||||
require.NoFileExists(t, banFile) // 5 hits * weight 2 = 10 == budget, exactly at the limit, not over
|
||||
s.Record(testIP, 400, 40001) // 6th hit cannot be covered -> breach
|
||||
lines := flushAndRead(t, s, banFile)
|
||||
require.Len(t, lines, 1)
|
||||
require.True(t, strings.HasSuffix(lines[0], " 1.2.3.4 1.2.3.4/32 400 40001")) // <ip> <prefix> <http> <ntfy-code>
|
||||
}
|
||||
|
||||
func TestService_Record_Weight10BansFast(t *testing.T) {
|
||||
banFile := filepath.Join(t.TempDir(), "ban.log")
|
||||
// Threshold 10, code weight 10 -> a single hit drains the whole budget, so the 2nd breaches.
|
||||
s := newTestService(t, banFile, 10, map[string]int{"42909": 10, "*": 1})
|
||||
s.Record(testIP, 429, 42909)
|
||||
s.flush()
|
||||
require.NoFileExists(t, banFile)
|
||||
s.Record(testIP, 429, 42909) // 2nd hit cannot be covered -> breach
|
||||
lines := flushAndRead(t, s, banFile)
|
||||
require.Len(t, lines, 1)
|
||||
require.True(t, strings.HasSuffix(lines[0], " 1.2.3.4 1.2.3.4/32 429 42909"))
|
||||
}
|
||||
|
||||
func TestService_Record_Weight0NeverBans(t *testing.T) {
|
||||
banFile := filepath.Join(t.TempDir(), "ban.log")
|
||||
// The legit-quota code is exempt (weight 0), so no number of hits ever bans.
|
||||
s := newTestService(t, banFile, 10, map[string]int{"42908": 0, "*": 1})
|
||||
for i := 0; i < 100; i++ {
|
||||
s.Record(testIP, 429, 42908)
|
||||
}
|
||||
s.flush()
|
||||
require.NoFileExists(t, banFile)
|
||||
}
|
||||
|
||||
func TestService_Record_SingleBucketNoRelaxation(t *testing.T) {
|
||||
banFile := filepath.Join(t.TempDir(), "ban.log")
|
||||
// One shared bucket per prefix: different codes draw down the SAME budget, so mixing them creates
|
||||
// no extra headroom (unlike per-code buckets, which would relax the effective limit for a mixed
|
||||
// offender).
|
||||
s := newTestService(t, banFile, 10, map[string]int{"403*": 2, "*": 1})
|
||||
s.Record(testIP, 403, 40301) // weight 2 -> 8 left
|
||||
s.Record(testIP, 403, 40301) // weight 2 -> 6 left
|
||||
s.Record(testIP, 403, 40301) // weight 2 -> 4 left
|
||||
s.flush()
|
||||
require.NoFileExists(t, banFile)
|
||||
for i := 0; i < 4; i++ {
|
||||
s.Record(testIP, 400, 40001) // weight 1 each -> drains the remaining 4 -> 0 left
|
||||
}
|
||||
s.flush()
|
||||
require.NoFileExists(t, banFile) // 6 + 4 = 10 == budget exactly, still not over
|
||||
s.Record(testIP, 400, 40001) // one more cannot be covered -> breach
|
||||
lines := flushAndRead(t, s, banFile)
|
||||
require.Len(t, lines, 1)
|
||||
}
|
||||
|
||||
func TestService_Record_ExactLineFormat(t *testing.T) {
|
||||
banFile := filepath.Join(t.TempDir(), "ban.log")
|
||||
s := newTestService(t, banFile, 1, map[string]int{"*": 1})
|
||||
before := time.Now().UTC().Truncate(time.Second)
|
||||
for i := 0; i < 3; i++ {
|
||||
s.Record(testIP, 429, 42901)
|
||||
}
|
||||
after := time.Now().UTC()
|
||||
lines := flushAndRead(t, s, banFile)
|
||||
require.Len(t, lines, 1)
|
||||
parts := strings.Split(lines[0], " ")
|
||||
require.Len(t, parts, 5) // "<timestamp> <ip> <prefix> <http-code> <ntfy-code>"
|
||||
ts, err := time.Parse(time.RFC3339, parts[0])
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, time.UTC, ts.Location())
|
||||
require.False(t, ts.Before(before))
|
||||
require.False(t, ts.After(after.Add(time.Second)))
|
||||
require.Equal(t, "1.2.3.4", parts[1]) // full IP
|
||||
require.Equal(t, "1.2.3.4/32", parts[2]) // masked to the default IPv4 prefix (/32)
|
||||
require.Equal(t, "429", parts[3]) // HTTP status
|
||||
require.Equal(t, "42901", parts[4]) // ntfy code
|
||||
}
|
||||
|
||||
func TestService_Record_IPv6MaskedToPrefix(t *testing.T) {
|
||||
banFile := filepath.Join(t.TempDir(), "ban.log")
|
||||
s := newTestService(t, banFile, 1, map[string]int{"*": 1})
|
||||
ip := netip.MustParseAddr("2001:db8::abcd")
|
||||
for i := 0; i < 3; i++ {
|
||||
s.Record(ip, 429, 42901)
|
||||
}
|
||||
parts := strings.Split(flushAndRead(t, s, banFile)[0], " ")
|
||||
require.Len(t, parts, 5)
|
||||
require.Equal(t, "2001:db8::abcd", parts[1]) // full IPv6 address
|
||||
require.Equal(t, "2001:db8::/64", parts[2]) // masked to the default IPv6 prefix (/64)
|
||||
}
|
||||
|
||||
func TestService_Record_PerPrefixIsolation(t *testing.T) {
|
||||
// Each source prefix gets its own bucket: one IP hammering to a breach must not push a different,
|
||||
// quiet IP over the edge. This is the whole point of keying by prefix instead of by visitor.
|
||||
banFile := filepath.Join(t.TempDir(), "ban.log")
|
||||
s := newTestService(t, banFile, 2, map[string]int{"*": 1})
|
||||
noisy := netip.MustParseAddr("1.1.1.1")
|
||||
quiet := netip.MustParseAddr("2.2.2.2")
|
||||
for i := 0; i < 5; i++ {
|
||||
s.Record(noisy, 429, 42901) // breaches its own bucket
|
||||
}
|
||||
s.Record(quiet, 429, 42901) // single hit, well under threshold
|
||||
lines := flushAndRead(t, s, banFile)
|
||||
require.Len(t, lines, 1) // only the noisy prefix is written
|
||||
require.True(t, strings.HasSuffix(lines[0], " 1.1.1.1 1.1.1.1/32 429 42901"))
|
||||
}
|
||||
|
||||
func TestService_Record_OncePerWindowThrottle(t *testing.T) {
|
||||
// Once a prefix has been written, further breaches within the window must not re-append it, so a
|
||||
// persistent offender produces exactly one line per window (mirrors the old per-visitor banEmit).
|
||||
banFile := filepath.Join(t.TempDir(), "ban.log")
|
||||
s := newTestService(t, banFile, 1, map[string]int{"*": 1})
|
||||
for i := 0; i < 50; i++ {
|
||||
s.Record(testIP, 429, 42901)
|
||||
}
|
||||
lines := flushAndRead(t, s, banFile)
|
||||
require.Len(t, lines, 1)
|
||||
}
|
||||
|
||||
func TestService_Record_BansPassedIP(t *testing.T) {
|
||||
// The Service bans the exact IP passed to Record -- the caller passes the offending request's IP,
|
||||
// which for an account-keyed visitor is not the visitor's stored IP.
|
||||
banFile := filepath.Join(t.TempDir(), "ban.log")
|
||||
s := newTestService(t, banFile, 1, map[string]int{"*": 1})
|
||||
offender := netip.MustParseAddr("5.6.7.8")
|
||||
for i := 0; i < 3; i++ {
|
||||
s.Record(offender, 429, 42901)
|
||||
}
|
||||
parts := strings.Split(flushAndRead(t, s, banFile)[0], " ")
|
||||
require.Equal(t, "5.6.7.8", parts[1]) // the IP passed to Record
|
||||
require.Equal(t, "5.6.7.8/32", parts[2]) // its prefix
|
||||
}
|
||||
|
||||
func TestService_Record_Ignores2xx3xx(t *testing.T) {
|
||||
banFile := filepath.Join(t.TempDir(), "ban.log")
|
||||
s := newTestService(t, banFile, 3, map[string]int{"*": 1})
|
||||
// Success and redirects must never count toward a ban, even over the threshold -- otherwise a
|
||||
// legit high-volume publisher (lots of 200s) would get banned.
|
||||
for i := 0; i < 20; i++ {
|
||||
s.Record(testIP, 200, 20000)
|
||||
s.Record(testIP, 302, 30000)
|
||||
}
|
||||
s.flush()
|
||||
require.NoFileExists(t, banFile)
|
||||
// A 4xx over the same budget still gets written.
|
||||
for i := 0; i < 5; i++ {
|
||||
s.Record(testIP, 400, 40001)
|
||||
}
|
||||
lines := flushAndRead(t, s, banFile)
|
||||
require.Len(t, lines, 1)
|
||||
require.True(t, strings.HasSuffix(lines[0], " 1.2.3.4 1.2.3.4/32 400 40001"))
|
||||
}
|
||||
|
||||
func TestService_Record_BuffersUntilFlush(t *testing.T) {
|
||||
// Writes are async: a breach buffers the ban line rather than writing it synchronously on the
|
||||
// request path. The line only reaches the file when runWriteLoop (or an explicit flush) runs.
|
||||
banFile := filepath.Join(t.TempDir(), "ban.log")
|
||||
s := newTestService(t, banFile, 1, map[string]int{"*": 1})
|
||||
for i := 0; i < 3; i++ {
|
||||
s.Record(testIP, 429, 42901)
|
||||
}
|
||||
require.NoFileExists(t, banFile) // not written synchronously
|
||||
s.mu.Lock()
|
||||
require.Len(t, s.pending, 1) // one line buffered (throttled to once per window)
|
||||
s.mu.Unlock()
|
||||
lines := flushAndRead(t, s, banFile)
|
||||
require.Len(t, lines, 1)
|
||||
require.True(t, strings.HasSuffix(lines[0], " 1.2.3.4 1.2.3.4/32 429 42901"))
|
||||
}
|
||||
|
||||
func TestService_Close_FlushesPending(t *testing.T) {
|
||||
// Close must flush buffered bans so nothing is lost on shutdown, and must block until it has.
|
||||
banFile := filepath.Join(t.TempDir(), "ban.log")
|
||||
s := NewService(&Config{File: banFile, Window: time.Minute, Threshold: 1, Weights: Weights{"*": 1}, PrefixBitsIPv4: 32, PrefixBitsIPv6: 64})
|
||||
for i := 0; i < 3; i++ {
|
||||
s.Record(testIP, 429, 42901)
|
||||
}
|
||||
require.NoFileExists(t, banFile) // still buffered
|
||||
s.Close() // blocks until the final flush completes
|
||||
data, err := os.ReadFile(banFile)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, strings.Split(strings.TrimRight(string(data), "\n"), "\n"), 1)
|
||||
}
|
||||
|
||||
func TestService_Prune_DropsIdlePrefixes(t *testing.T) {
|
||||
// A prefix idle for a full window has a refilled bucket, so prune drops it to bound memory. An
|
||||
// active prefix (seen within the window) is kept.
|
||||
banFile := filepath.Join(t.TempDir(), "ban.log")
|
||||
s := newTestService(t, banFile, 10, map[string]int{"*": 1})
|
||||
idle := netip.MustParseAddr("9.9.9.9")
|
||||
s.Record(idle, 400, 40001)
|
||||
s.Record(testIP, 400, 40001)
|
||||
require.Len(t, s.trackers, 2)
|
||||
|
||||
// Backdate the idle prefix past the window, then prune.
|
||||
idlePrefix := s.prefix(idle)
|
||||
s.mu.Lock()
|
||||
s.trackers[idlePrefix].seen = time.Now().Add(-2 * time.Minute)
|
||||
s.mu.Unlock()
|
||||
s.prune()
|
||||
|
||||
require.Len(t, s.trackers, 1)
|
||||
_, ok := s.trackers[idlePrefix]
|
||||
require.False(t, ok) // idle prefix dropped
|
||||
_, ok = s.trackers[s.prefix(testIP)]
|
||||
require.True(t, ok) // active prefix kept
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
package ban
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Weights maps a matcher key to a strike weight for the abuse ban-feed (see ParseWeights, WeightFor).
|
||||
type Weights map[string]int
|
||||
|
||||
// ParseWeights normalizes a list like ["42909:10","403:2","*:1"] into a Weights map. A key is an
|
||||
// exact ntfy code, a family ("429*"), a bare HTTP status ("403" -> "403*"), or "*"; weights are ints
|
||||
// >= 0 (0 = exempt). Malformed entries are rejected so misconfiguration fails at startup.
|
||||
func ParseWeights(entries []string) (Weights, error) {
|
||||
out := make(Weights, len(entries))
|
||||
for _, entry := range entries {
|
||||
key, weightStr, ok := strings.Cut(entry, ":")
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid ban-weight %q, want KEY:WEIGHT", entry)
|
||||
}
|
||||
weight, err := strconv.Atoi(strings.TrimSpace(weightStr))
|
||||
if err != nil || weight < 0 {
|
||||
return nil, fmt.Errorf("invalid ban-weight value in %q, want a non-negative integer", entry)
|
||||
}
|
||||
key = strings.TrimSpace(key)
|
||||
if !validWeightKey(key) {
|
||||
return nil, fmt.Errorf("invalid ban-weight key in %q, want %q, an ntfy code, an HTTP status, or a PREFIX*", entry, "*")
|
||||
}
|
||||
// A bare 3-digit HTTP status is shorthand for the whole family (e.g. "403" -> "403*").
|
||||
if len(key) == 3 && isAllDigits(key) {
|
||||
key += "*"
|
||||
}
|
||||
out[key] = weight
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// WeightFor returns the strike weight for an ntfy error code, longest-match-wins (exact > family > "*").
|
||||
// If nothing matches (no "*" catch-all) it returns the implied default 1, so a forgotten "*" still
|
||||
// bans; use "*:0" to exempt everything not explicitly weighted.
|
||||
func (w Weights) WeightFor(errorCode int) int {
|
||||
code := strconv.Itoa(errorCode)
|
||||
weight, bestLen, matched := 0, -1, false
|
||||
for key, wt := range w {
|
||||
matchLen := -1
|
||||
switch {
|
||||
case key == "*":
|
||||
matchLen = 0
|
||||
case strings.HasSuffix(key, "*"):
|
||||
if prefix := strings.TrimSuffix(key, "*"); strings.HasPrefix(code, prefix) {
|
||||
matchLen = len(prefix)
|
||||
}
|
||||
case key == code:
|
||||
matchLen = len(code)
|
||||
}
|
||||
if matchLen > bestLen {
|
||||
weight, bestLen, matched = wt, matchLen, true
|
||||
}
|
||||
}
|
||||
if !matched {
|
||||
return 1
|
||||
}
|
||||
return weight
|
||||
}
|
||||
|
||||
// validWeightKey reports whether key is a legal matcher: "*", an all-digits code, or DIGITS*.
|
||||
func validWeightKey(key string) bool {
|
||||
if key == "*" {
|
||||
return true
|
||||
}
|
||||
digits := strings.TrimSuffix(key, "*")
|
||||
return digits != "" && isAllDigits(digits)
|
||||
}
|
||||
|
||||
func isAllDigits(s string) bool {
|
||||
for _, r := range s {
|
||||
if r < '0' || r > '9' {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return s != ""
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
package ban
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestParseWeights(t *testing.T) {
|
||||
// Exact codes, a bare 3-digit HTTP status (normalized to a family), an exempt code, and "*".
|
||||
weights, err := ParseWeights([]string{"42909:10", "403:2", "42908:0", "*:1"})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, Weights{"42909": 10, "403*": 2, "42908": 0, "*": 1}, weights)
|
||||
|
||||
// A bare 3-digit HTTP status normalizes to its family.
|
||||
weights, err = ParseWeights([]string{"429:5"})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, Weights{"429*": 5}, weights)
|
||||
|
||||
// An explicit family key stays as-is.
|
||||
weights, err = ParseWeights([]string{"429*:5"})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, Weights{"429*": 5}, weights)
|
||||
|
||||
// Weight 0 is valid and means exempt.
|
||||
weights, err = ParseWeights([]string{"42908:0"})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, Weights{"42908": 0}, weights)
|
||||
|
||||
_, err = ParseWeights([]string{"401"}) // Missing weight
|
||||
require.Error(t, err)
|
||||
_, err = ParseWeights([]string{"401:-1"}) // Negative weight
|
||||
require.Error(t, err)
|
||||
_, err = ParseWeights([]string{"401:abc"}) // Non-integer weight
|
||||
require.Error(t, err)
|
||||
_, err = ParseWeights([]string{"abc:10"}) // Non-numeric key
|
||||
require.Error(t, err)
|
||||
_, err = ParseWeights([]string{"4*3:10"}) // Star not at the end
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestWeights_WeightFor(t *testing.T) {
|
||||
weights, err := ParseWeights([]string{"42908:0", "42903:0", "42905:0", "42910:0", "42909:10", "429*:1", "403*:2", "4*:1", "5*:1"})
|
||||
require.NoError(t, err)
|
||||
// Longest-match-wins: exact 5-digit beats "429*" beats "4*" beats "*".
|
||||
require.Equal(t, 0, weights.WeightFor(42908))
|
||||
require.Equal(t, 10, weights.WeightFor(42909))
|
||||
require.Equal(t, 1, weights.WeightFor(42901))
|
||||
require.Equal(t, 2, weights.WeightFor(40311))
|
||||
require.Equal(t, 1, weights.WeightFor(40011))
|
||||
require.Equal(t, 1, weights.WeightFor(50312))
|
||||
// No rule matches (this config has 4*/5* but no "*"), so the implied default weight 1 applies.
|
||||
require.Equal(t, 1, weights.WeightFor(30012))
|
||||
}
|
||||
|
||||
func TestWeights_WeightFor_NoStarRuleImpliesWeight1(t *testing.T) {
|
||||
// With no "*" rule, a code that matches nothing defaults to weight 1 (can be banned), so the
|
||||
// feature can't be silently turned into a no-op by forgetting "*". Explicit codes still win.
|
||||
weights, err := ParseWeights([]string{"42908:0", "42909:10"})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 0, weights.WeightFor(42908)) // explicitly exempt
|
||||
require.Equal(t, 10, weights.WeightFor(42909)) // explicit
|
||||
require.Equal(t, 1, weights.WeightFor(42901)) // unmatched -> implied 1
|
||||
require.Equal(t, 1, weights.WeightFor(40001)) // unmatched -> implied 1
|
||||
}
|
||||
|
||||
func TestWeights_WeightFor_ExplicitStarZeroExemptsAll(t *testing.T) {
|
||||
// An explicit "*:0" is the opt-out: exempt everything not otherwise weighted.
|
||||
weights, err := ParseWeights([]string{"42909:10", "*:0"})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 10, weights.WeightFor(42909)) // explicit
|
||||
require.Equal(t, 0, weights.WeightFor(42901)) // *:0 -> exempt everything else
|
||||
}
|
||||
Reference in New Issue
Block a user