smallwebwaf/internal/reputation/abuseipdb.go

package reputation

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/netip"
	"net/url"
	"slices"
	"sync"
	"time"

	"github.com/hashicorp/golang-lru/v2/simplelru"
	"sneak.berlin/go/smallwebwaf/internal/alerts"
)

const (
	// AbuseIPDBURL is where clients are checked: the check endpoint of
	// AbuseIPDB's API.
	AbuseIPDBURL = "https://api.abuseipdb.com/api/v2/check"
	// AbuseIPDBSource is how the request log, the alerts and the metrics
	// name AbuseIPDB.
	AbuseIPDBSource = "abuseipdb"
	// maxAnswerBytes is the most of an answer of AbuseIPDB that is read.
	maxAnswerBytes = 64 << 10
	// day is the length of the day the checks are counted in, in UTC.
	day = 24 * time.Hour
)

var (
	errNoScore      = errors.New("the answer gives no abuseConfidenceScore")
	errBudgetUsedUp = errors.New(
		"checks spent; none is made until the day ends at 00:00 UTC")
)

// Score is what AbuseIPDB said about a client, as reputation.json holds
// it: the client, an IPv4 address or an IPv6 group, its abuse confidence
// score, from 0 to 100, and when AbuseIPDB answered.
type Score struct {
	Client  netip.Prefix `json:"client"`
	Score   int64        `json:"score"`
	Fetched time.Time    `json:"fetched"`
}

// Checks are what reputation.json keeps of the checks of clients with
// AbuseIPDB: the day, in UTC, of the checks Spent counts, zero before the
// first, and the scores still in use.
type Checks struct {
	Day    time.Time `json:"day,omitzero"`
	Spent  int       `json:"spent"`
	Scores []Score   `json:"scores"`
}

// AbuseIPDBParams are what NewAbuseIPDB needs.
type AbuseIPDBParams struct {
	// URL is where clients are checked, normally AbuseIPDBURL, with Key,
	// the account's key (SWW_WAF_ABUSEIPDB_KEY).
	URL string
	Key string
	// MinScore is the least score that is a hit (SWW_WAF_ABUSEIPDB_MIN_SCORE),
	// and DailyBudget the most checks made in a day, in UTC
	// (SWW_WAF_ABUSEIPDB_DAILY_BUDGET).
	MinScore    int64
	DailyBudget int
	// CacheTTL is how long a score is used after it was fetched
	// (SWW_WAF_REPUTATION_CACHE_TTL), and Timeout how long a check may take
	// (SWW_WAF_REPUTATION_TIMEOUT).
	CacheTTL time.Duration
	Timeout  time.Duration
	// Now tells the time, normally time.Now in UTC.
	Now func() time.Time
	// ProcessLog receives each check that fails, and why, and the day's
	// budget used up.
	ProcessLog *slog.Logger
	// Alerts receive a source_failure alert for each.
	Alerts *alerts.Queue
}

// AbuseIPDB checks clients with AbuseIPDB, in the background, and keeps
// their scores. It is safe for concurrent use.
type AbuseIPDB struct {
	params     AbuseIPDBParams
	httpClient *http.Client

	mu sync.Mutex
	// scores are by client. Each is added as it is fetched and never moved
	// up, so that the one fetched longest ago is the first dropped.
	scores *simplelru.LRU[netip.Prefix, Score]
	// checking are the clients whose check is under way.
	checking map[netip.Prefix]bool
	// day is the day, in UTC, of the checks spent counts.
	day   time.Time
	spent int
	// checks and failures count the checks made and those that failed,
	// and retryAt is when a client may be checked again after the last
	// check failed.
	checks   int
	failures int
	retryAt  time.Time
}

// NewAbuseIPDB returns an AbuseIPDB with no score yet, and no check spent.
func NewAbuseIPDB(params AbuseIPDBParams) *AbuseIPDB {
	scores, err := simplelru.NewLRU[netip.Prefix, Score](maxVerdicts, nil)
	if err != nil {
		panic(err) // NewLRU fails only for a size below one
	}

	return &AbuseIPDB{
		params:     params,
		httpClient: &http.Client{},
		scores:     scores,
		checking:   map[netip.Prefix]bool{},
	}
}

// Hit returns AbuseIPDB's score of client, an IPv4 address or an IPv6
// group, and whether it is a hit: MinScore or more. A score is used until
// CacheTTL has passed since it was fetched, whichever of the client's
// addresses its request comes from. A client without one is checked in
// the background, by addr, the address its request came from, if
// offender, if it has committed an offence, unless its check is under
// way, a check failed less than failureDelay ago, or the day's checks
// have used up DailyBudget; Hit never waits for a check. The check that
// uses the budget up is logged and raised as a source_failure alert. ctx
// is the context of the client's request, and a check goes on after the
// request ends.
func (a *AbuseIPDB) Hit(
	ctx context.Context, client netip.Prefix, addr netip.Addr, offender bool,
) (int64, bool) {
	a.mu.Lock()

	now := a.params.Now()

	kept, found := a.scores.Peek(client)
	if found && now.Sub(kept.Fetched) < a.params.CacheTTL {
		a.mu.Unlock()

		return kept.Score, kept.Score >= a.params.MinScore
	}

	if today := now.Truncate(day); !a.day.Equal(today) {
		a.day, a.spent = today, 0
	}

	check := offender && !a.checking[client] && !now.Before(a.retryAt) &&
		a.spent < a.params.DailyBudget
	if check {
		a.checking[client] = true
		a.checks++
		a.spent++

		go a.check(context.WithoutCancel(ctx), client, addr)
	}

	usedUp := check && a.spent == a.params.DailyBudget

	a.mu.Unlock()

	if usedUp {
		a.alert("the daily budget of AbuseIPDB checks is used up",
			fmt.Errorf("%d %w", a.params.DailyBudget, errBudgetUsedUp))
	}

	return 0, false
}

// Checked returns how many checks were made.
func (a *AbuseIPDB) Checked() int {
	a.mu.Lock()
	defer a.mu.Unlock()

	return a.checks
}

// Failures returns how many checks failed.
func (a *AbuseIPDB) Failures() int {
	a.mu.Lock()
	defer a.mu.Unlock()

	return a.failures
}

// BudgetLeft returns how many checks the day's budget has left.
func (a *AbuseIPDB) BudgetLeft() int {
	a.mu.Lock()
	defer a.mu.Unlock()

	if !a.day.Equal(a.params.Now().Truncate(day)) {
		return a.params.DailyBudget
	}

	return max(a.params.DailyBudget-a.spent, 0)
}

// Snapshot returns the checks spent and every score still in use, sorted
// by client, as reputation.json keeps them.
func (a *AbuseIPDB) Snapshot() Checks {
	a.mu.Lock()

	now := a.params.Now()
	checks := Checks{Day: a.day, Spent: a.spent, Scores: make([]Score, 0, a.scores.Len())}

	for _, kept := range a.scores.Values() {
		if now.Sub(kept.Fetched) < a.params.CacheTTL {
			checks.Scores = append(checks.Scores, kept)
		}
	}

	a.mu.Unlock()

	slices.SortFunc(checks.Scores, func(x, y Score) int {
		return x.Client.Compare(y.Client)
	})

	return checks
}

// Load keeps checks, read from reputation.json, in place of those it
// keeps, but for the scores past maxVerdicts, those fetched longest ago.
// One fetched CacheTTL ago or more is neither used nor written, as for any
// score.
func (a *AbuseIPDB) Load(checks Checks) {
	scores := slices.Clone(checks.Scores)
	slices.SortStableFunc(scores, func(x, y Score) int {
		return x.Fetched.Compare(y.Fetched)
	})

	a.mu.Lock()
	defer a.mu.Unlock()

	a.day, a.spent = checks.Day, checks.Spent
	a.scores.Purge()

	for _, kept := range scores {
		a.scores.Add(kept.Client, kept)
	}
}

// check checks client with AbuseIPDB by addr, one of its addresses, keeps
// the score as client's, and notes the check as no longer under way. A
// check that fails gives no score: it is counted, logged and raised as a
// source_failure alert, and no client is checked for failureDelay.
func (a *AbuseIPDB) check(ctx context.Context, client netip.Prefix, addr netip.Addr) {
	score, err := a.ask(ctx, addr)
	now := a.params.Now()

	a.mu.Lock()

	delete(a.checking, client)

	if err == nil {
		a.scores.Add(client, Score{Client: client, Score: score, Fetched: now})
	} else {
		a.failures++
		a.retryAt = now.Add(failureDelay)
	}

	a.mu.Unlock()

	if err != nil {
		a.alert("checking a client with AbuseIPDB failed", err)
	}
}

// ask asks AbuseIPDB for addr's abuse confidence score, sending the key
// in the header Key. An answer other than 200, one that gives no score,
// and none within Timeout, fail.
func (a *AbuseIPDB) ask(ctx context.Context, addr netip.Addr) (int64, error) {
	ctx, cancel := context.WithTimeout(ctx, a.params.Timeout)
	defer cancel()

	query := url.Values{"ipAddress": {addr.String()}}

	req, err := http.NewRequestWithContext(ctx, http.MethodGet,
		a.params.URL+"?"+query.Encode(), http.NoBody)
	if err != nil {
		return 0, fmt.Errorf("make the request: %w", err)
	}

	req.Header.Set("Key", a.params.Key)
	req.Header.Set("Accept", "application/json")

	res, err := a.httpClient.Do(req)
	if err != nil {
		// Do's error names the URL, which holds the client's address, which
		// is not to be logged: only what went wrong is kept.
		return 0, fmt.Errorf("check the client: %w", errors.Unwrap(err))
	}

	defer func() {
		_ = res.Body.Close()
	}()

	if res.StatusCode != http.StatusOK {
		return 0, fmt.Errorf("%w %s", errStatus, res.Status)
	}

	var answer struct {
		Data struct {
			AbuseConfidenceScore *int64 `json:"abuseConfidenceScore"`
		} `json:"data"`
	}

	err = json.NewDecoder(io.LimitReader(res.Body, maxAnswerBytes)).Decode(&answer)
	if err != nil {
		return 0, fmt.Errorf("read the answer: %w", err)
	}

	if answer.Data.AbuseConfidenceScore == nil {
		return 0, errNoScore
	}

	return *answer.Data.AbuseConfidenceScore, nil
}

// alert raises a source_failure alert from AbuseIPDB with reason and err,
// and logs them.
func (a *AbuseIPDB) alert(reason string, err error) {
	// Raised before it is logged, so that the alert is there once the log
	// line is.
	raiseFailure(a.params.Alerts, reason, AbuseIPDBSource, err)
	a.params.ProcessLog.Warn(reason, "source", AbuseIPDBSource, "error", err.Error())
}