smallwebwaf/internal/reputation/abuseipdb_test.go

package reputation_test

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"net/netip"
	"reflect"
	"slices"
	"strings"
	"sync"
	"testing"
	"testing/synctest"
	"time"

	"sneak.berlin/go/smallwebwaf/internal/alerts"
	"sneak.berlin/go/smallwebwaf/internal/metrics"
	"sneak.berlin/go/smallwebwaf/internal/reputation"
)

// The tests of AbuseIPDB run in synctest bubbles, as those of the lists
// do, and AbuseIPDB is a stand-in reached without the network, for the
// same reason. A bubble's clock starts at midnight UTC, as a day the
// checks are counted in starts.

const (
	// key is the account's key the tests give, the only one the stand-in
	// takes.
	key = "abuseipdb-key-0123456789abcdef"
	// suspect and other are clients that have committed an offence.
	suspect = "203.0.113.9"
	other   = "2001:db8::9"
)

func TestOnlyAnOffenderWithoutAScoreIsChecked(t *testing.T) {
	t.Parallel()

	synctest.Test(t, func(t *testing.T) {
		abuseIPDB := &abuseIPDBStandIn{scores: map[string]int64{suspect: 100}}
		checker := newAbuseIPDB(abuseIPDB, abuseIPDBParams())

		// A client that has committed no offence is not checked.
		wantScore(t, checker, suspect, false, 0, false)
		synctest.Wait()
		wantChecked(t, abuseIPDB)

		// An offender is, and from then on its score is used, whether or not
		// it is an offender.
		wantScore(t, checker, suspect, true, 0, false)
		synctest.Wait()
		wantScore(t, checker, suspect, true, 100, true)
		wantScore(t, checker, suspect, false, 100, true)
		synctest.Wait()
		wantChecked(t, abuseIPDB, suspect)
	})
}

func TestIPv6ClientIsCheckedOnceAndItsScoreUsedForEachOfItsAddresses(t *testing.T) {
	t.Parallel()

	synctest.Test(t, func(t *testing.T) {
		// 15 addresses of 2001:db8:1:2::/64, one client, each in a part of
		// it of its own.
		var addresses []string
		for i := 1; i < 16; i++ {
			addresses = append(addresses, fmt.Sprintf("2001:db8:1:2:%x::9", i<<12))
		}

		abuseIPDB := &abuseIPDBStandIn{scores: map[string]int64{addresses[0]: 100}}
		checker := newAbuseIPDB(abuseIPDB, abuseIPDBParams())

		// A request from each has the client checked once, by the first.
		for _, address := range addresses {
			hitFrom(t, checker, address, true)
		}

		synctest.Wait()
		wantChecked(t, abuseIPDB, addresses[0])

		// Its score is the whole client's.
		for _, address := range addresses {
			wantScore(t, checker, address, true, 100, true)
		}

		synctest.Wait()
		wantChecked(t, abuseIPDB, addresses[0])
	})
}

func TestScoreAtOrOverTheMinimumIsAHit(t *testing.T) {
	t.Parallel()

	synctest.Test(t, func(t *testing.T) {
		scores := map[string]int64{"192.0.2.74": 74, "192.0.2.75": 75, "192.0.2.100": 100}
		p := abuseIPDBParams()
		p.MinScore = 75
		checker := newAbuseIPDB(&abuseIPDBStandIn{scores: scores}, p)

		for client := range scores {
			hitFrom(t, checker, client, true)
		}

		synctest.Wait()

		for client, score := range scores {
			wantScore(t, checker, client, true, score, score >= 75)
		}
	})
}

func TestScoreUsedUntilTheCacheTTLHasPassedSinceItWasFetched(t *testing.T) {
	t.Parallel()

	synctest.Test(t, func(t *testing.T) {
		abuseIPDB := &abuseIPDBStandIn{scores: map[string]int64{suspect: 100}}
		checker := newAbuseIPDB(abuseIPDB, abuseIPDBParams())

		hitFrom(t, checker, suspect, true)
		synctest.Wait()

		// AbuseIPDB gives another score from now on, but the one kept is
		// used, and the client is not checked again, until the TTL has
		// passed.
		abuseIPDB.setScore(suspect, 80)
		time.Sleep(cacheTTL - time.Nanosecond)
		wantScore(t, checker, suspect, true, 100, true)
		synctest.Wait()
		wantChecked(t, abuseIPDB, suspect)

		// Then it is not used, and the client is checked again.
		time.Sleep(time.Nanosecond)
		wantScore(t, checker, suspect, true, 0, false)
		synctest.Wait()
		wantScore(t, checker, suspect, true, 80, true)
		wantChecked(t, abuseIPDB, suspect, suspect)
	})
}

func TestDailyBudgetKeptAcrossARestartAndWholeAgainAsTheDayEnds(t *testing.T) {
	t.Parallel()

	synctest.Test(t, func(t *testing.T) {
		var log bytes.Buffer

		queue := newQueue()
		p := abuseIPDBParams()
		p.DailyBudget = 3
		p.Alerts = queue
		p.ProcessLog = slog.New(slog.NewJSONHandler(&log, nil))
		abuseIPDB := &abuseIPDBStandIn{scores: map[string]int64{suspect: 100}}
		checker := newAbuseIPDB(abuseIPDB, p)

		// At noon, the first three offenders spend the budget, and the
		// fourth, unchecked, is not.
		time.Sleep(12 * time.Hour)

		const unchecked = "192.0.2.4"

		clients := []string{suspect, "192.0.2.2", "192.0.2.3", unchecked}
		for _, client := range clients {
			hitFrom(t, checker, client, true)
		}

		synctest.Wait()
		wantChecked(t, abuseIPDB, clients[:3]...)
		wantBudgetLeft(t, checker, 0)

		// The check that used the budget up raised the alert, and logged it.
		const usedUp = "the daily budget of AbuseIPDB checks is used up"

		wantFailureAlert(t, queue, time.Now(), usedUp,
			"3 checks spent; none is made until the day ends at 00:00 UTC", 0)

		if !strings.Contains(log.String(), `"msg":"`+usedUp+`"`) {
			t.Errorf("logged\n%s\nwant the budget used up", log.String())
		}

		// Restarted with what reputation.json keeps, it uses the scores, and
		// checks no client until the day ends.
		restarted := &abuseIPDBStandIn{}
		again := newAbuseIPDB(restarted, p)
		again.Load(checker.Snapshot())

		wantScore(t, again, suspect, true, 100, true)
		wantBudgetLeft(t, again, 0)
		time.Sleep(12*time.Hour - time.Nanosecond)
		wantScore(t, again, unchecked, true, 0, false)
		synctest.Wait()
		wantChecked(t, restarted)

		// At midnight the budget is whole again.
		time.Sleep(time.Nanosecond)
		wantBudgetLeft(t, again, 3)
		wantScore(t, again, unchecked, true, 0, false)
		synctest.Wait()
		wantChecked(t, restarted, unchecked)
		wantBudgetLeft(t, again, 2)
	})
}

func TestFailedCheckGivesNoScoreAndNoClientIsCheckedForAMinute(t *testing.T) {
	t.Parallel()

	for _, tc := range []struct {
		name string
		// key is the key sent, status and body what AbuseIPDB answers with,
		// and error the failure.
		key, body string
		status    int
		error     string
	}{
		{
			"a refusal, past AbuseIPDB's own limit", key,
			`{"errors":[{"detail":"Daily rate limit of 1000 requests exceeded"}]}`,
			http.StatusTooManyRequests, "the server answered 429 Too Many Requests",
		},
		{
			"a refusal of a wrong key", "wrong-key-0123456789abcdef", "", 0,
			"the server answered 401 Unauthorized",
		},
		{
			"a server failure", key, "", http.StatusInternalServerError,
			"the server answered 500 Internal Server Error",
		},
		{
			"an answer without a score", key, `{"data":{"ipAddress":"` + suspect + `"}}`,
			http.StatusOK, "the answer gives no abuseConfidenceScore",
		},
		{
			"an answer that is not JSON", key, "<html>", http.StatusOK,
			"read the answer: invalid character '<' looking for beginning of value",
		},
	} {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			synctest.Test(t, func(t *testing.T) {
				var log bytes.Buffer

				queue := newQueue()
				p := abuseIPDBParams()
				p.Key = tc.key
				p.Alerts = queue
				p.ProcessLog = slog.New(slog.NewJSONHandler(&log, nil))
				abuseIPDB := &abuseIPDBStandIn{status: tc.status, body: tc.body}
				checker := newAbuseIPDB(abuseIPDB, p)

				// The failure gives no score, and no client is checked within a
				// minute of it.
				wantScore(t, checker, suspect, true, 0, false)
				synctest.Wait()
				time.Sleep(time.Minute - time.Nanosecond)
				wantScore(t, checker, other, true, 0, false)
				synctest.Wait()
				wantChecked(t, abuseIPDB, suspect)
				wantFailures(t, checker, 1)

				time.Sleep(time.Nanosecond)
				wantScore(t, checker, other, true, 0, false)
				synctest.Wait()
				wantChecked(t, abuseIPDB, suspect, other)
				wantFailures(t, checker, 2)

				if scores := checker.Snapshot().Scores; len(scores) != 0 {
					t.Errorf("scores %+v, want none", scores)
				}

				// One alert for the first failure; the cooldown holds back the
				// second.
				wantFailureAlert(t, queue, time.Now().Add(-time.Minute),
					"checking a client with AbuseIPDB failed", tc.error, 1)

				if !strings.Contains(log.String(), `"msg":"checking a client with `+
					`AbuseIPDB failed","source":"abuseipdb","error":"`+tc.error) {
					t.Errorf("logged\n%s\nwant the failures", log.String())
				}
			})
		})
	}
}

func TestCheckNotAnsweredWithinTheTimeoutFailsAndHitNeverWaits(t *testing.T) {
	t.Parallel()

	synctest.Test(t, func(t *testing.T) {
		queue := newQueue()
		p := abuseIPDBParams()
		p.Alerts = queue
		abuseIPDB := &abuseIPDBStandIn{hanging: true}
		checker := newAbuseIPDB(abuseIPDB, p)
		began := time.Now()

		// The second, while the first's check is under way, starts none.
		wantScore(t, checker, suspect, true, 0, false)
		wantScore(t, checker, suspect, true, 0, false)

		if waited := time.Since(began); waited != 0 {
			t.Errorf("waited %s for the check, want no wait", waited)
		}

		time.Sleep(timeout - time.Nanosecond)
		synctest.Wait()
		wantChecked(t, abuseIPDB, suspect)
		wantFailures(t, checker, 0)

		time.Sleep(time.Nanosecond)
		synctest.Wait()
		wantFailures(t, checker, 1)
		wantFailureAlert(t, queue, time.Now(), "checking a client with AbuseIPDB failed",
			"check the client: context deadline exceeded", 0)
	})
}

func TestKeyIsSentInTheKeyHeaderAndNeverShown(t *testing.T) {
	t.Parallel()

	synctest.Test(t, func(t *testing.T) {
		var log bytes.Buffer

		queue := newQueue()
		p := abuseIPDBParams()
		p.Alerts = queue
		p.ProcessLog = slog.New(slog.NewJSONHandler(&log, nil))
		abuseIPDB := &abuseIPDBStandIn{scores: map[string]int64{suspect: 100}}
		checker := newAbuseIPDB(abuseIPDB, p)
		m := metrics.New(1, "app")
		m.AddReputation(reputation.New(params()), reputation.NewDNSBL(dnsblParams()))
		m.AddAbuseIPDB(checker)

		// One check that AbuseIPDB answers, and one it refuses with an answer
		// that names the key.
		wantScore(t, checker, suspect, true, 0, false)
		synctest.Wait()
		wantScore(t, checker, suspect, true, 100, true)

		abuseIPDB.answerWith(http.StatusUnauthorized, `{"errors":[{"detail":"`+key+`"}]}`)
		wantScore(t, checker, other, true, 0, false)
		synctest.Wait()
		wantFailures(t, checker, 1)

		abuseIPDB.mu.Lock()
		sent := slices.Clone(abuseIPDB.keys)
		abuseIPDB.mu.Unlock()

		if !slices.Equal(sent, []string{key, key}) {
			t.Errorf("checks sent the keys %v, want %s twice", sent, key)
		}

		alerted, err := json.Marshal(waiting(queue))
		if err != nil {
			t.Fatalf("encode the alerts: %v", err)
		}

		kept, err := json.Marshal(checker.Snapshot())
		if err != nil {
			t.Fatalf("encode the checks: %v", err)
		}

		for name, shown := range map[string]string{
			"the log": log.String(), "the alerts": string(alerted),
			"the metrics": scrapeMetrics(t, m), "reputation.json": string(kept),
		} {
			if strings.Contains(shown, key) {
				t.Errorf("%s shows the key:\n%s", name, shown)
			}
		}
	})
}

func TestMetricsCountTheChecksTheFailuresAndTheBudgetLeft(t *testing.T) {
	t.Parallel()

	synctest.Test(t, func(t *testing.T) {
		abuseIPDB := &abuseIPDBStandIn{}
		p := abuseIPDBParams()
		p.DailyBudget = 5
		checker := newAbuseIPDB(abuseIPDB, p)
		m := metrics.New(1, "app")
		m.AddReputation(reputation.New(params()), reputation.NewDNSBL(dnsblParams()))
		m.AddAbuseIPDB(checker)

		// One check that AbuseIPDB answers, and one that fails.
		hitFrom(t, checker, suspect, true)
		synctest.Wait()
		abuseIPDB.answerWith(http.StatusInternalServerError, "")
		hitFrom(t, checker, other, true)
		synctest.Wait()

		scraped := scrapeMetrics(t, m)

		for series, want := range map[string]string{
			"queries_total":          "2",
			"failures_total":         "1",
			"daily_budget_remaining": "3",
		} {
			line := "\nsmallwebwaf_reputation_" + series +
				`{instance="app",source="abuseipdb"} ` + want + "\n"
			if !strings.Contains(scraped, line) {
				t.Errorf("metrics\n%s\nwant%s", scraped, line)
			}
		}
	})
}

func TestScoreFetchedATTLAgoIsNeitherUsedNorKept(t *testing.T) {
	t.Parallel()

	now := time.Date(2026, 10, 7, 0, 0, 0, 0, time.UTC)
	p := abuseIPDBParams()
	p.Now = func() time.Time { return now }
	checker := reputation.NewAbuseIPDB(p)

	// The last score still in use, and one, of other's /64, fetched a TTL
	// ago.
	inUse := reputation.Score{
		Client: netip.MustParsePrefix(suspect + "/32"), Score: 100,
		Fetched: now.Add(-cacheTTL + time.Nanosecond),
	}
	stale := reputation.Score{
		Client: netip.MustParsePrefix("2001:db8::/64"), Score: 100,
		Fetched: now.Add(-cacheTTL),
	}

	checker.Load(reputation.Checks{Scores: []reputation.Score{stale, inUse}})

	wantScore(t, checker, suspect, false, 100, true)
	wantScore(t, checker, other, false, 0, false)

	got := checker.Snapshot().Scores
	if !reflect.DeepEqual(got, []reputation.Score{inUse}) {
		t.Errorf("scores %+v, want only %+v", got, inUse)
	}
}

func TestAtMost100000ScoresKeptTheOneFetchedLongestAgoDroppedFirst(t *testing.T) {
	t.Parallel()

	now := time.Date(2026, 10, 7, 0, 0, 0, 0, time.UTC)
	p := abuseIPDBParams()
	p.Now = func() time.Time { return now }
	checker := reputation.NewAbuseIPDB(p)

	// 100,001 scores, listed by client, as reputation.json lists them, each
	// fetched a millisecond before the one before it: the last is one too
	// many.
	const count = 100001

	scores := make([]reputation.Score, 0, count)

	addr := netip.MustParseAddr("198.18.0.0")
	for i := range count {
		scores = append(scores, reputation.Score{
			Client:  netip.PrefixFrom(addr, 32),
			Fetched: now.Add(-time.Duration(i) * time.Millisecond),
		})
		addr = addr.Next()
	}

	checker.Load(reputation.Checks{Scores: scores})

	got := checker.Snapshot().Scores
	if len(got) != count-1 || !slices.Contains(got, scores[0]) ||
		slices.Contains(got, scores[count-1]) {
		t.Errorf("%d scores kept, want all but the one fetched longest ago", len(got))
	}
}

// abuseIPDBStandIn is a stand-in for AbuseIPDB. It answers a check sent
// with key by the client's score, as scores gives it, 0 for a client it
// does not give; a check sent with another key with 401; and, while
// status is not 0, every check with status and body; and while hanging,
// none at all. It notes each client checked, and the key sent.
type abuseIPDBStandIn struct {
	mu      sync.Mutex
	scores  map[string]int64
	status  int
	body    string
	hanging bool
	checked []string
	keys    []string
}

// RoundTrip has the stand-in answer req, in place of the network. A check
// abandoned before the stand-in answers fails, as over the network.
func (s *abuseIPDBStandIn) RoundTrip(req *http.Request) (*http.Response, error) {
	client := req.URL.Query().Get("ipAddress")
	sent := req.Header.Get("Key")

	s.mu.Lock()
	s.checked = append(s.checked, client)
	s.keys = append(s.keys, sent)
	score := s.scores[client]
	status, body, hanging := s.status, s.body, s.hanging
	s.mu.Unlock()

	switch {
	case hanging:
		<-req.Context().Done()

		return nil, req.Context().Err()
	case sent != key:
		status = http.StatusUnauthorized
	case status == 0:
		status = http.StatusOK
		body = fmt.Sprintf(`{"data":{"ipAddress":%q,"abuseConfidenceScore":%d}}`, client,
			score)
	}

	return &http.Response{
		StatusCode: status,
		Status:     fmt.Sprintf("%d %s", status, http.StatusText(status)),
		Header:     http.Header{},
		Body:       io.NopCloser(strings.NewReader(body)),
		Request:    req,
	}, nil
}

// setScore has the stand-in give client score.
func (s *abuseIPDBStandIn) setScore(client string, score int64) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.scores[client] = score
}

// answerWith has the stand-in answer every check with status and body.
func (s *abuseIPDBStandIn) answerWith(status int, body string) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.status, s.body = status, body
}

// abuseIPDBParams returns the AbuseIPDBParams of the tests: key, a minimum
// score of 75, a daily budget of 900, and the cache TTL and timeout of the
// DNSBL tests, by the bubble's clock, with alerts to a queue that sends
// none.
func abuseIPDBParams() reputation.AbuseIPDBParams {
	return reputation.AbuseIPDBParams{
		URL:         "https://abuseipdb.example/api/v2/check",
		Key:         key,
		MinScore:    75,
		DailyBudget: 900,
		CacheTTL:    cacheTTL,
		Timeout:     timeout,
		Now:         time.Now,
		ProcessLog:  slog.New(slog.DiscardHandler),
		Alerts:      newQueue(),
	}
}

// newAbuseIPDB returns the AbuseIPDB of p, checking clients with
// abuseIPDB.
func newAbuseIPDB(
	abuseIPDB *abuseIPDBStandIn, p reputation.AbuseIPDBParams,
) *reputation.AbuseIPDB {
	checker := reputation.NewAbuseIPDB(p)
	checker.SetTransport(abuseIPDB)

	return checker
}

// wantScore checks the score checker gives client, and whether it is a
// hit, as a request from client finds them, offender or not.
func wantScore(
	t *testing.T, checker *reputation.AbuseIPDB, client string, offender bool,
	score int64, hit bool,
) {
	t.Helper()

	gotScore, gotHit := hitFrom(t, checker, client, offender)
	if gotScore != score || gotHit != hit {
		t.Errorf("%s has the score %d, a hit %t, want %d, %t", client, gotScore, gotHit,
			score, hit)
	}
}

// hitFrom is checker's Hit for a request from address, offender or not.
// Its client is address for an IPv4 address, and its /64 for an IPv6 one,
// as smallwebwaf counts clients.
func hitFrom(
	t *testing.T, checker *reputation.AbuseIPDB, address string, offender bool,
) (int64, bool) {
	t.Helper()

	addr := netip.MustParseAddr(address)

	client := netip.PrefixFrom(addr, addr.BitLen())
	if addr.Is6() {
		client = netip.PrefixFrom(addr, 64).Masked()
	}

	return checker.Hit(t.Context(), client, addr, offender)
}

// wantChecked checks the clients the stand-in was asked about, in any
// order.
func wantChecked(t *testing.T, abuseIPDB *abuseIPDBStandIn, want ...string) {
	t.Helper()

	abuseIPDB.mu.Lock()
	got := slices.Sorted(slices.Values(abuseIPDB.checked))
	abuseIPDB.mu.Unlock()

	want = slices.Sorted(slices.Values(want))

	if !slices.Equal(got, want) {
		t.Errorf("checked %v, want %v", got, want)
	}
}

// wantFailures checks how many checks failed.
func wantFailures(t *testing.T, checker *reputation.AbuseIPDB, want int) {
	t.Helper()

	if got := checker.Failures(); got != want {
		t.Errorf("%d checks failed, want %d", got, want)
	}
}

// wantFailureAlert checks that the one alert waiting in queue is a
// source_failure alert from AbuseIPDB, raised at raised, with reason and
// the error failure, and that the cooldown has held back held repeats of
// it.
func wantFailureAlert(
	t *testing.T, queue *alerts.Queue, raised time.Time, reason, failure string,
	held int64,
) {
	t.Helper()

	got := waiting(queue)
	if len(got) != 1 || !got[0].Time.Equal(raised) ||
		got[0].Event != alerts.EventSourceFailure || got[0].Reason != reason ||
		got[0].Detail["source"] != reputation.AbuseIPDBSource ||
		got[0].Detail["error"] != failure || queue.Suppressed() != held {
		t.Errorf("alerts waiting %+v, %d held back, want only AbuseIPDB's %q with %q, "+
			"and %d", got, queue.Suppressed(), reason, failure, held)
	}
}

// wantBudgetLeft checks how many checks the day's budget has left.
func wantBudgetLeft(t *testing.T, checker *reputation.AbuseIPDB, want int) {
	t.Helper()

	if got := checker.BudgetLeft(); got != want {
		t.Errorf("%d checks left, want %d", got, want)
	}
}

// scrapeMetrics returns the metrics m serves.
func scrapeMetrics(t *testing.T, m *metrics.Metrics) string {
	t.Helper()

	scraped := httptest.NewRecorder()
	m.ServeHTTP(scraped, httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/",
		http.NoBody))

	return scraped.Body.String()
}