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.
package lookup

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/netip"
	"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")

// 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

	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: its country, "" when GeoJS
// cannot place it, and when GeoJS said so.
type answer struct {
	country  string
	received time.Time
}

// 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,
		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
}

// 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 one was received less than keepFor
// ago.
func (g *GeoJS) kept(client netip.Prefix) (string, bool) {
	kept, found := g.answers.Get(client)
	if !found || g.now().Sub(kept.received) >= keepFor {
		return "", false
	}

	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. Each client asked about gets its answer, with
// no country for one GeoJS gave none for. After a failure GeoJS is left
// alone for a while, and every waiting client 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()

	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

	for _, client := range clients {
		g.answers.Add(client, answer{country: countries[client.Addr()], received: now})
		close(g.waiting[client].asked)
		delete(g.waiting, client)
	}

	return true
}

// request asks GeoJS about clients in one request, and returns the
// country it gave for each address it answered for.
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+"?ip="+strings.Join(addrs, ","), http.NoBody)
	if err != nil {
		return nil, fmt.Errorf("make the request to GeoJS: %w", err)
	}

	res, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("ask GeoJS: %w", 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] = item.Country
		}
	}

	return countries, nil
}