smallwebwaf/internal/lookup/lookup.go

// Package lookup looks up each client's country through the GeoJS web
// service, and keeps the answers in memory, for at most 100,000 clients
// and for 7 days each. The answers are written to lookups.json and read
// from it by the state package.
package lookup

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

	"github.com/hashicorp/golang-lru/v2/simplelru"
)

// URL is GeoJS's country endpoint. Asked about several addresses at once,
// comma separated in its ip parameter, it answers with a list.
const URL = "https://get.geojs.io/v1/ip/country.json"

const (
	// keepFor is how long an answer is used instead of asking GeoJS again.
	keepFor = 7 * 24 * time.Hour
	// maxAnswers is how many answers are kept. Past it, the one used
	// longest ago is dropped.
	maxAnswers = 100000
	// maxWaiting is how many clients may wait to be asked about. Past it,
	// a new client counts as not found and is not asked about until there
	// is room, so that a swarm of new addresses while GeoJS is down cannot
	// fill the memory.
	maxWaiting = 10000
	// maxPerRequest is how many addresses one request to GeoJS asks about.
	maxPerRequest = 200
	// timeout is how long a new client waits for its answer, and how long
	// a request to GeoJS may take before it is abandoned.
	timeout = time.Second
	// After a failure GeoJS is not asked again for a second, and for
	// retryDelayFactor times as long after each further failure in a row,
	// up to five minutes.
	firstRetryDelay  = time.Second
	retryDelayFactor = 2
	maxRetryDelay    = 5 * time.Minute
	// maxResponseBytes is the most of GeoJS's answer that is read.
	maxResponseBytes = 1 << 20
)

var (
	errStatus  = errors.New("GeoJS answered")
	errLeftOut = errors.New("GeoJS's answer left out")
)

// Params are what New needs.
type Params struct {
	// URL is where GeoJS is asked, normally URL.
	URL string
	// Now tells the time, normally time.Now.
	Now func() time.Time
	// ProcessLog receives GeoJS's failures.
	ProcessLog *slog.Logger
}

// GeoJS looks up clients' countries through GeoJS. At most one request
// to GeoJS is under way at a time, and it asks about every client waiting,
// up to maxPerRequest. It is safe for concurrent use.
type GeoJS struct {
	url        string
	now        func() time.Time
	processLog *slog.Logger
	// httpClient follows no redirect, so that visitors' addresses go to
	// GeoJS alone: a redirect is a failure.
	httpClient *http.Client

	mu      sync.Mutex
	answers *simplelru.LRU[netip.Prefix, *Answer]
	// waiting are the clients without an answer: those to ask GeoJS about,
	// and those it is being asked about.
	waiting map[netip.Prefix]*wait
	// asking is true while a request to GeoJS is under way.
	asking bool
	// retryDelay is how long GeoJS is left alone after its last failure,
	// zero after an answer; retryAt is when it may be asked again.
	retryDelay time.Duration
	retryAt    time.Time
}

// Answer is what GeoJS said about a client, as lookups.json holds it: its
// country, "" when GeoJS cannot place it, when GeoJS said so, and when
// the answer was last used.
type Answer struct {
	Client   netip.Prefix `json:"client"`
	Country  string       `json:"country"`
	Answered time.Time    `json:"answered"`
	Used     time.Time    `json:"used"`
}

// wait is a client waiting for its answer.
type wait struct {
	// asked is closed when the client gets its answer, and closed and
	// replaced each time GeoJS fails before then.
	asked chan struct{}
	// late is true once the client has gone without an answer, for a
	// whole timeout or because GeoJS failed: its requests no longer wait.
	late bool
}

// New returns a GeoJS with no answer kept yet.
func New(params Params) *GeoJS {
	answers, err := simplelru.NewLRU[netip.Prefix, *Answer](maxAnswers, nil)
	if err != nil {
		panic(err) // NewLRU fails only for a size below one
	}

	return &GeoJS{
		url:        params.URL,
		now:        params.Now,
		processLog: params.ProcessLog,
		httpClient: &http.Client{
			CheckRedirect: func(*http.Request, []*http.Request) error {
				return http.ErrUseLastResponse
			},
		},
		answers: answers,
		waiting: map[netip.Prefix]*wait{},
	}
}

// Country returns the country GeoJS places client in, as a two-letter
// code in capitals, or "" when the country cannot be found: GeoJS cannot
// place the client, or has not answered in time. An answer is kept for 7
// days. Without one, a client waits up to timeout for it, unless it has
// gone without one before; until GeoJS answers, the client is asked about
// again in the background. ctx is the context of the client's request,
// and ends the wait when it ends.
//
// GeoJS is asked about the client's first address, which is the client's
// own address for IPv4, and an address in the same place for an IPv6 /64.
func (g *GeoJS) Country(ctx context.Context, client netip.Prefix) string {
	country, asked := g.answerOrWait(ctx, client)
	if asked == nil {
		return country
	}

	timer := time.NewTimer(timeout)
	defer timer.Stop()

	select {
	case <-asked:
	case <-timer.C:
	case <-ctx.Done():
	}

	g.mu.Lock()
	defer g.mu.Unlock()

	country, found := g.kept(client)

	w, waiting := g.waiting[client]
	if !found && waiting {
		w.late = true
	}

	return country
}

// Snapshot returns every answer kept, sorted by client, as lookups.json
// lists them.
func (g *GeoJS) Snapshot() []Answer {
	g.mu.Lock()

	answers := make([]Answer, 0, g.answers.Len())
	for _, kept := range g.answers.Values() {
		answers = append(answers, *kept)
	}

	g.mu.Unlock()

	slices.SortFunc(answers, func(a, b Answer) int {
		return a.Client.Compare(b.Client)
	})

	return answers
}

// Load keeps answers read from lookups.json, in a GeoJS that keeps none
// yet, in the order they were last used, so that the one used longest
// ago is dropped first. Answers GeoJS gave keepFor ago or more are
// dropped.
func (g *GeoJS) Load(answers []Answer) {
	g.mu.Lock()
	defer g.mu.Unlock()

	answers = slices.Clone(answers)
	slices.SortStableFunc(answers, func(a, b Answer) int {
		return a.Used.Compare(b.Used)
	})

	now := g.now()

	for _, answer := range answers {
		if now.Sub(answer.Answered) < keepFor {
			g.answers.Add(answer.Client, &answer)
		}
	}
}

// answerOrWait returns client's kept answer if it has one. Otherwise it
// puts the client among those waiting if there is room, has GeoJS asked
// about them if it can be, and returns what to wait on for the answer, or
// nil when there is nothing to wait for.
func (g *GeoJS) answerOrWait(
	ctx context.Context, client netip.Prefix,
) (string, <-chan struct{}) {
	g.mu.Lock()
	defer g.mu.Unlock()

	country, found := g.kept(client)
	if found {
		return country, nil
	}

	w, waiting := g.waiting[client]
	if !waiting && len(g.waiting) < maxWaiting {
		w = &wait{asked: make(chan struct{})}
		g.waiting[client] = w
	}

	g.ask(ctx)

	if w == nil {
		return "", nil // too many clients wait already
	}

	if !g.asking {
		// GeoJS is left alone after a failure, so no answer can come.
		w.late = true
	}

	if w.late {
		return "", nil
	}

	return "", w.asked
}

// kept returns client's answer, if GeoJS gave it less than keepFor ago,
// and notes that it was used.
func (g *GeoJS) kept(client netip.Prefix) (string, bool) {
	now := g.now()

	kept, found := g.answers.Get(client)
	if !found || now.Sub(kept.Answered) >= keepFor {
		return "", false
	}

	kept.Used = now

	return kept.Country, true
}

// ask starts asking GeoJS about the waiting clients, unless a request to
// it is under way or it is left alone after a failure. The requests to
// GeoJS are for every client waiting, so they go on when the client's
// request whose ctx is given ends.
func (g *GeoJS) ask(ctx context.Context) {
	if g.asking || g.now().Before(g.retryAt) {
		return
	}

	g.asking = true

	go g.askAboutWaiting(context.WithoutCancel(ctx))
}

// askAboutWaiting asks GeoJS about the waiting clients, one request at a
// time, until none is left or GeoJS fails.
func (g *GeoJS) askAboutWaiting(ctx context.Context) {
	for {
		clients := g.nextClients()
		if len(clients) == 0 {
			return
		}

		countries, err := g.request(ctx, clients)
		if !g.keep(clients, countries, err) {
			return
		}
	}
}

// nextClients returns up to maxPerRequest of the waiting clients. When
// none is waiting, it returns none and notes that no request to GeoJS is
// under way.
func (g *GeoJS) nextClients() []netip.Prefix {
	g.mu.Lock()
	defer g.mu.Unlock()

	if len(g.waiting) == 0 {
		g.asking = false

		return nil
	}

	clients := make([]netip.Prefix, 0, min(len(g.waiting), maxPerRequest))

	for client := range g.waiting {
		if len(clients) == maxPerRequest {
			break
		}

		clients = append(clients, client)
	}

	return clients
}

// keep notes how a request to GeoJS about clients ended, and reports
// whether GeoJS answered about all of them. Each client whose address
// GeoJS's answer names gets its answer, with no country when GeoJS gave
// none. An answer that leaves an address out is a failure. After a
// failure GeoJS is left alone for a while, and every client still waiting
// stops waiting and is asked about once GeoJS is asked again.
func (g *GeoJS) keep(
	clients []netip.Prefix, countries map[netip.Addr]string, err error,
) bool {
	g.mu.Lock()
	defer g.mu.Unlock()

	now := g.now()
	leftOut := 0

	for _, client := range clients {
		country, named := countries[client.Addr()]
		if !named {
			leftOut++

			continue
		}

		g.answers.Add(client, &Answer{
			Client: client, Country: country, Answered: now, Used: now,
		})
		close(g.waiting[client].asked)
		delete(g.waiting, client)
	}

	if err == nil && leftOut > 0 {
		err = fmt.Errorf("%w %d of %d addresses", errLeftOut, leftOut, len(clients))
	}

	if err != nil {
		g.retryDelay = min(max(retryDelayFactor*g.retryDelay, firstRetryDelay),
			maxRetryDelay)
		g.retryAt = now.Add(g.retryDelay)
		g.asking = false

		for _, w := range g.waiting {
			close(w.asked)

			w.asked = make(chan struct{})
			w.late = true
		}

		g.processLog.Warn("asking GeoJS failed",
			"error", err.Error(), "asking_again_in", g.retryDelay.String())

		return false
	}

	g.retryDelay = 0

	return true
}

// request asks GeoJS about clients in one request, and returns the
// country it gave, in capitals, for each address its answer names.
func (g *GeoJS) request(
	ctx context.Context, clients []netip.Prefix,
) (map[netip.Addr]string, error) {
	addrs := make([]string, 0, len(clients))

	for _, client := range clients {
		addrs = append(addrs, client.Addr().String())
	}

	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, g.url, http.NoBody)
	if err != nil {
		return nil, fmt.Errorf("make the request to GeoJS: %w", err)
	}

	req.URL.RawQuery = "ip=" + strings.Join(addrs, ",")

	res, err := g.httpClient.Do(req)
	if err != nil {
		// Do's error names the URL, and so the visitors' addresses, which
		// are not to be logged: only what went wrong is kept.
		return nil, fmt.Errorf("ask GeoJS: %w", errors.Unwrap(err))
	}

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

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

	var answers []struct {
		IP      string `json:"ip"`
		Country string `json:"country"`
	}

	err = json.NewDecoder(io.LimitReader(res.Body, maxResponseBytes)).Decode(&answers)
	if err != nil {
		return nil, fmt.Errorf("read GeoJS's answer: %w", err)
	}

	countries := make(map[netip.Addr]string, len(answers))

	for _, item := range answers {
		addr, err := netip.ParseAddr(item.IP)
		if err == nil {
			countries[addr] = strings.ToUpper(item.Country)
		}
	}

	return countries, nil
}