routewatch/internal/geo/locator.go

// Package geo finds where IP addresses are. It asks geojs.io, and keeps those
// answers in memory and in cache.db. When geojs.io gives no location for an
// address, it reads a copy of DB-IP's City Lite database, which it downloads
// once a month.
package geo

import (
	"context"
	"database/sql"
	"fmt"
	"log/slog"
	"net/http"
	"net/netip"
	"sync"
	"time"

	lru "github.com/hashicorp/golang-lru/v2"
	"github.com/oschwald/maxminddb-golang/v2"
)

// The values of Location.Source.
const (
	sourceGeoJS = "geojs"
	sourceDBIP  = "dbip"
)

const (
	// maxEntries limits how many geojs answers are held in memory.
	maxEntries = 250_000

	// userAgentHeader and userAgent name routewatch to geojs and DB-IP.
	userAgentHeader = "User-Agent"
	userAgent       = "routewatch (+https://git.eeqj.de/sneak/routewatch)"
)

// Params is what New needs.
type Params struct {
	// DB is the open cache.db.
	DB *sql.DB

	// Dir is the directory that holds the copy of the database.
	Dir string

	// GeoJSURL is the URL of the geojs geo endpoint, GeoJSEndpoint in use.
	GeoJSURL string

	// DatabaseURLFormat is the URL of the database with %s for the year and
	// month, DatabaseDownloadFormat in use.
	DatabaseURLFormat string

	Logger *slog.Logger
}

// Location is where an address is. A part the source does not know is nil.
//
//nolint:tagliatelle // snake_case keys, as the Go style guide asks
type Location struct {
	CountryCode      *string  `json:"country_code"`
	Country          *string  `json:"country"`
	ContinentCode    *string  `json:"continent_code"`
	Region           *string  `json:"region"`
	City             *string  `json:"city"`
	Latitude         *float64 `json:"latitude"`
	Longitude        *float64 `json:"longitude"`
	AccuracyRadiusKm *int     `json:"accuracy_radius_km"`
	TimeZone         *string  `json:"time_zone"`

	// Source is geojs or dbip.
	Source string `json:"source"`

	// Updated is when the value was fetched: the geojs answer, or the copy of
	// the database.
	Updated time.Time `json:"updated"`
}

// Locator answers where addresses are. Start loads its cache and starts its
// goroutines.
type Locator struct {
	db                *sql.DB
	dir               string
	geojsURL          string
	databaseURLFormat string
	timeLimit         time.Duration
	maxEntries        int
	maxDatabaseBytes  int64
	logger            *slog.Logger
	client            *http.Client
	now               func() time.Time

	// entries holds geojs answers by address, and drops the least recently
	// used first.
	entries *lru.Cache[string, *Location]

	// pending holds the answers that cache.db does not have yet.
	pendingMu sync.Mutex
	pending   map[string]*Location

	// database is the open copy, nil until one has opened, and databaseFetched
	// is when that copy was downloaded. A lookup holds the read lock, so that
	// a new copy never closes the old one under it.
	databaseMu      sync.RWMutex
	database        *maxminddb.Reader
	databaseFetched time.Time

	// stop ends the goroutines Start runs, and done is closed when they have
	// returned.
	stop context.CancelFunc
	done chan struct{}
}

// New returns a Locator with an empty cache. Start loads it.
func New(params Params) (*Locator, error) {
	entries, err := lru.New[string, *Location](maxEntries)
	if err != nil {
		return nil, fmt.Errorf("creating the geolocation cache: %w", err)
	}

	return &Locator{
		db:                params.DB,
		dir:               params.Dir,
		geojsURL:          params.GeoJSURL,
		databaseURLFormat: params.DatabaseURLFormat,
		timeLimit:         geojsTimeLimit,
		maxEntries:        maxEntries,
		maxDatabaseBytes:  maxDatabaseSize,
		logger:            params.Logger.With("component", "geo"),
		client:            &http.Client{},
		now:               time.Now,
		entries:           entries,
		pending:           make(map[string]*Location),
		done:              make(chan struct{}),
	}, nil
}

// Lookup returns where each of addrs is, in the same order. The entry for an
// address neither source answers for is nil. Addresses with no current cached
// answer go to geojs in one request. The returned Locations are shared with
// the cache and must not be changed.
func (l *Locator) Lookup(ctx context.Context, addrs []netip.Addr) []*Location {
	locations := make([]*Location, len(addrs))

	var unanswered []string

	for i, addr := range addrs {
		address := plainAddress(addr)

		locations[i] = l.cached(ctx, address)
		if locations[i] == nil {
			unanswered = append(unanswered, address)
		}
	}

	if len(unanswered) == 0 {
		return locations
	}

	answers, err := l.askGeoJS(ctx, unanswered)
	if err != nil {
		l.logger.Warn("geojs lookup failed; using the database", "error", err)
	}

	for i, addr := range addrs {
		if locations[i] != nil {
			continue
		}

		address := plainAddress(addr)

		answer, ok := answers[address]
		if ok {
			l.store(address, answer)
			locations[i] = answer

			continue
		}

		locations[i] = l.fromDatabase(addr)
	}

	return locations
}

// plainAddress returns addr as the requests to geojs and the cache name it: an
// IPv4 address mapped into IPv6 as plain IPv4, and with no zone. A client
// chooses the zone, so it must not reach the request.
func plainAddress(addr netip.Addr) string {
	return addr.Unmap().WithZone("").String()
}

// Start loads the cached answers that have not expired and opens the copy of
// the database, if there is one. It then starts the goroutines that write new
// answers, delete expired ones and download the database. They run until Stop:
// the end of ctx does not stop them.
func (l *Locator) Start(ctx context.Context) error {
	err := l.load(ctx)
	if err != nil {
		return err
	}

	l.logger.Info("Geolocation cache loaded", "entries", l.entries.Len())

	l.openDatabaseCopy()

	ctx, l.stop = context.WithCancel(context.WithoutCancel(ctx))

	go func() {
		defer close(l.done)
		defer l.closeDatabase()

		var workers sync.WaitGroup

		workers.Go(func() { l.keepCache(ctx) })
		workers.Go(func() { l.keepDatabaseCurrent(ctx) })
		workers.Wait()
	}()

	return nil
}

// Stop cancels a download in progress, writes the answers cache.db does not
// have yet, and waits until the goroutines have returned or ctx ends.
func (l *Locator) Stop(ctx context.Context) {
	l.stop()

	select {
	case <-l.done:
		l.logger.Info("Geolocation stopped")
	case <-ctx.Done():
		l.logger.Warn("Stop deadline passed before the geolocation cache " +
			"was written")
	}
}