routewatch/internal/geo/database.go

package geo

import (
	"compress/gzip"
	"context"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"net/http"
	"net/netip"
	"os"
	"path/filepath"
	"time"

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

const (
	// DatabaseDownloadFormat is the URL of DB-IP's City Lite database, with %s
	// for the year and month.
	DatabaseDownloadFormat = "https://download.db-ip.com/free/" +
		"dbip-city-lite-%s.mmdb.gz"

	// databaseMonthFormat is the year and month in the URL.
	databaseMonthFormat = "2006-01"

	// databaseFileName is the copy of the database in the state directory.
	databaseFileName = "dbip-city-lite.mmdb"

	// databaseDownloadSuffix ends the name of the file a download is written
	// to, until it has opened.
	databaseDownloadSuffix = ".download"

	// maxDatabaseSize is the largest database accepted once unpacked, in
	// bytes.
	maxDatabaseSize = 200 << 20

	// databaseDownloadTimeout limits one download.
	databaseDownloadTimeout = 10 * time.Minute

	// databaseCheckInterval is how often RouteWatch checks whether a new
	// month's database is due. A new month's file appears some time on the
	// first of the month, so a failed download is tried again this much later.
	databaseCheckInterval = time.Hour
)

var (
	errDatabaseStatus  = errors.New("unexpected HTTP status from DB-IP")
	errDatabaseTooLong = errors.New("database is over the size limit")
)

// dbipName is the part of a database record that holds an English name.
type dbipName struct {
	Names struct {
		English string `maxminddb:"en"`
	} `maxminddb:"names"`
}

// dbipRecord is the part of a database record RouteWatch reads.
type dbipRecord struct {
	City      dbipName `maxminddb:"city"`
	Continent struct {
		Code string `maxminddb:"code"`
	} `maxminddb:"continent"`
	Country struct {
		dbipName

		ISOCode string `maxminddb:"iso_code"`
	} `maxminddb:"country"`
	Location struct {
		Latitude  *float64 `maxminddb:"latitude"`
		Longitude *float64 `maxminddb:"longitude"`
	} `maxminddb:"location"`
	Subdivisions []dbipName `maxminddb:"subdivisions"`
}

// location returns the record as a Location fetched at the time given. The
// first subdivision is the region. The database has no accuracy radius and no
// time zone.
func (r *dbipRecord) location(fetched time.Time) *Location {
	found := &Location{
		CountryCode:   nilIfEmpty(r.Country.ISOCode),
		Country:       nilIfEmpty(r.Country.Names.English),
		ContinentCode: nilIfEmpty(r.Continent.Code),
		City:          nilIfEmpty(r.City.Names.English),
		Latitude:      r.Location.Latitude,
		Longitude:     r.Location.Longitude,
		Source:        sourceDBIP,
		Updated:       fetched,
	}

	if len(r.Subdivisions) > 0 {
		found.Region = nilIfEmpty(r.Subdivisions[0].Names.English)
	}

	return found
}

// nilIfEmpty returns a pointer to text, or nil when text is empty.
func nilIfEmpty(text string) *string {
	if text == "" {
		return nil
	}

	return &text
}

// fromDatabase returns where the copy of the database places addr. It returns
// nil when there is no copy yet or the copy has no record for addr. It also
// returns nil, and logs the error, when reading the copy fails.
func (l *Locator) fromDatabase(addr netip.Addr) *Location {
	l.databaseMu.RLock()
	defer l.databaseMu.RUnlock()

	if l.database == nil {
		return nil
	}

	result := l.database.Lookup(addr)

	err := result.Err()
	if err != nil {
		l.logger.Error("Failed to look up an address in the database",
			"address", addr, "error", err)

		return nil
	}

	if !result.Found() {
		return nil
	}

	var record dbipRecord

	err = result.Decode(&record)
	if err != nil {
		l.logger.Error("Failed to read a database record",
			"address", addr, "error", err)

		return nil
	}

	return record.location(l.databaseFetched)
}

// keepDatabaseCurrent downloads the database when it is due, at once and then
// every databaseCheckInterval, until ctx ends.
func (l *Locator) keepDatabaseCurrent(ctx context.Context) {
	l.downloadDatabaseIfDue(ctx)

	checks := time.NewTicker(databaseCheckInterval)
	defer checks.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-checks.C:
			l.downloadDatabaseIfDue(ctx)
		}
	}
}

// downloadDatabaseIfDue downloads this month's database, unless the copy in
// use was downloaded this month. A failure leaves the copy as it is.
func (l *Locator) downloadDatabaseIfDue(ctx context.Context) {
	now := l.now().UTC()

	if !l.databaseDue(now) {
		return
	}

	err := l.downloadDatabase(ctx, now)
	if err != nil && ctx.Err() == nil {
		l.logger.Error("Database download failed; keeping the previous copy",
			"error", err)
	}
}

// databaseDue says whether there is no open copy or it is from an earlier
// month than now.
func (l *Locator) databaseDue(now time.Time) bool {
	l.databaseMu.RLock()
	defer l.databaseMu.RUnlock()

	return l.database == nil ||
		l.databaseFetched.Format(databaseMonthFormat) !=
			now.Format(databaseMonthFormat)
}

// downloadDatabase fetches the database of now's month. The file replaces the
// copy, on disk and in use, only once it has opened.
func (l *Locator) downloadDatabase(ctx context.Context, now time.Time) error {
	ctx, cancel := context.WithTimeout(ctx, databaseDownloadTimeout)
	defer cancel()

	url := fmt.Sprintf(l.databaseURLFormat, now.Format(databaseMonthFormat))

	request, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return err
	}

	request.Header.Set(userAgentHeader, userAgent)

	response, err := l.client.Do(request)
	if err != nil {
		return err
	}

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

	if response.StatusCode != http.StatusOK {
		return fmt.Errorf("%w: %s", errDatabaseStatus, response.Status)
	}

	return l.replaceDatabase(response.Body, now)
}

// replaceDatabase unpacks the gzip stream body beside the copy, then makes the
// result the copy once it has opened. Whatever fails, the file it wrote is
// removed and the copy stays.
func (l *Locator) replaceDatabase(body io.Reader, now time.Time) error {
	root, err := os.OpenRoot(l.dir)
	if err != nil {
		return err
	}

	defer func() { _ = root.Close() }()

	downloadName := databaseFileName + databaseDownloadSuffix

	err = unpackDatabase(root, downloadName, body, l.maxDatabaseBytes)
	if err != nil {
		_ = root.Remove(downloadName)

		return err
	}

	opened, err := maxminddb.Open(filepath.Join(l.dir, downloadName))
	if err != nil {
		_ = root.Remove(downloadName)

		return fmt.Errorf("opening the downloaded database: %w", err)
	}

	err = root.Rename(downloadName, databaseFileName)
	if err != nil {
		_ = opened.Close()
		_ = root.Remove(downloadName)

		return err
	}

	l.useDatabase(opened, now.Truncate(time.Second))
	l.logger.Info("Downloaded a new copy of the database")

	return nil
}

// unpackDatabase writes the gzip stream body to the file name in root, up to
// limit bytes once unpacked.
func unpackDatabase(
	root *os.Root, name string, body io.Reader, limit int64,
) error {
	unpacked, err := gzip.NewReader(body)
	if err != nil {
		return fmt.Errorf("unpacking the database: %w", err)
	}

	f, err := root.Create(name)
	if err != nil {
		return err
	}

	defer func() { _ = f.Close() }()

	written, err := io.Copy(f, io.LimitReader(unpacked, limit+1))
	if err != nil {
		return fmt.Errorf("unpacking the database: %w", err)
	}

	if written > limit {
		return errDatabaseTooLong
	}

	return nil
}

// openDatabaseCopy opens the copy left by an earlier run, if there is one. A
// copy that does not open is treated as no copy, so it is downloaded again.
func (l *Locator) openDatabaseCopy() {
	path := filepath.Join(l.dir, databaseFileName)

	info, err := os.Stat(path)
	if errors.Is(err, fs.ErrNotExist) {
		return
	}

	if err != nil {
		l.logger.Error("Failed to read the database copy", "error", err)

		return
	}

	opened, err := maxminddb.Open(path)
	if err != nil {
		l.logger.Error("Failed to open the database copy", "error", err)

		return
	}

	l.useDatabase(opened, info.ModTime().UTC().Truncate(time.Second))
}

// useDatabase makes opened the copy in use and closes the one it replaces.
func (l *Locator) useDatabase(opened *maxminddb.Reader, fetched time.Time) {
	l.databaseMu.Lock()
	defer l.databaseMu.Unlock()

	if l.database != nil {
		err := l.database.Close()
		if err != nil {
			l.logger.Error("Failed to close the old database copy", "error", err)
		}
	}

	l.database = opened
	l.databaseFetched = fetched
}

// closeDatabase closes the copy in use, after which no source can answer from
// it.
func (l *Locator) closeDatabase() {
	l.databaseMu.Lock()
	defer l.databaseMu.Unlock()

	if l.database == nil {
		return
	}

	err := l.database.Close()
	if err != nil {
		l.logger.Error("Failed to close the database copy", "error", err)
	}

	l.database = nil
}