l-ui/internal/web/service/geoip_country.go

package service

import (
	_ "embed"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/oschwald/maxminddb-golang"

	"github.com/mhsanaei/3x-ui/v3/internal/config"
	"github.com/mhsanaei/3x-ui/v3/internal/database/model"
	"github.com/mhsanaei/3x-ui/v3/internal/logger"
)

// Node country flags (item 7 Q7, resolved 2026-10-24): the public
// views carry an ISO country code per node, which the themes render
// as a flag. The upstream hub resolves it by calling ipinfo.io per
// node; this build never sends a node address to a third party — the
// lookup is a local MMDB read on the master, from the node's own
// reported addresses. The agent is untouched: it already reports its
// interface addresses in its facts, and the master already observes
// the connection's source addresses (SeenIPv4/SeenIPv6).
//
// The data is the DB-IP Lite country database (CC BY 4.0 by DB-IP,
// https://db-ip.com; attribution and snapshot provenance live in
// geoipdata/NOTICE), in the MMDB conversion published by the
// sapics/ip-location-db project. Three layers, in winning order:
//
//  1. a local file in the data folder (geoip-country.mmdb beside the
//     panel database) — written by the weekly updater below, or by
//     the operator by hand;
//  2. the factory snapshot embedded in the binary (geoipdata/), so a
//     fresh install shows flags with no network at all;
//  3. no database: lookups answer "", the views render exactly the
//     pre-flag shape.
//
// Reads are mmap-backed (maxminddb.Open) or run against the embedded
// bytes in place; either way a lookup is a pure in-memory tree walk
// on the snapshot path, and answers are cached per address for the
// life of the loaded database.

//go:embed geoipdata/country.mmdb
var geoipCountryEmbedded []byte

// geoipCountryFileName is the local copy's name inside the data
// folder (config.GetDBFolderPath, where x-ui.db and probe.db live).
const geoipCountryFileName = "geoip-country.mmdb"

// geoipCountryDefaultURL is the weekly updater's fallback source,
// used when the operator has not set the geoip_country_url setting.
//
// RELEASE-TIME (main session): point this at the self-hosted Gitea
// mirror asset for geoip-country.mmdb. Empty disables auto-update;
// the embedded snapshot and any hand-placed local file still apply.
const geoipCountryDefaultURL = ""

// Settings keys (main settings KV, alongside the other probe_* keys).
const (
	geoipCountryURLSetting  = "geoip_country_url"
	geoipCountryETagSetting = "geoip_country_etag"
)

// geoipCountryMaxBytes caps one update download; the country
// database is ~8 MB today, the cap only stops a runaway response.
const geoipCountryMaxBytes = 64 << 20

// geoipCountryProbeIP / geoipCountryProbeIPv6 are the well-known
// addresses a candidate database must answer to be trusted (Google
// DNS v4, AliDNS v6 — both resolve in the DB-IP data).
const (
	geoipCountryProbeIP   = "8.8.8.8"
	geoipCountryProbeIPv6 = "2400:3200::1"
)

// GeoIPCountryLocalPath is where the local country database lives.
func GeoIPCountryLocalPath() string {
	return filepath.Join(config.GetDBFolderPath(), geoipCountryFileName)
}

// geoipCountryUpdateURL resolves the updater's source: the operator
// setting wins, else the compiled-in default (possibly empty).
func geoipCountryUpdateURL() string {
	if v, ok := probeSetting(geoipCountryURLSetting); ok {
		if v = strings.TrimSpace(v); v != "" {
			return v
		}
	}
	return geoipCountryDefaultURL
}

// ---- the loaded database ----

// countryRecord is the DB-IP Lite record shape: a flat country_code.
type countryRecord struct {
	CountryCode string `maxminddb:"country_code"`
}

// geoipDB is one loaded database plus its per-address answer cache
// ("" answers are cached too — an unlisted address stays unlisted
// until the database is swapped).
type geoipDB struct {
	reader *maxminddb.Reader
	origin string // "local" | "embedded"

	mu    sync.Mutex
	cache map[string]string
}

func (db *geoipDB) lookup(ipStr string) string {
	db.mu.Lock()
	if cc, ok := db.cache[ipStr]; ok {
		db.mu.Unlock()
		return cc
	}
	db.mu.Unlock()

	cc := ""
	var rec countryRecord
	if err := db.reader.Lookup(net.ParseIP(ipStr), &rec); err == nil {
		cc = rec.CountryCode
	}

	db.mu.Lock()
	db.cache[ipStr] = cc
	db.mu.Unlock()
	return cc
}

var (
	geoipMu   sync.RWMutex
	geoipCur  *geoipDB
	geoipOnce sync.Once
)

// currentGeoIP returns the live database, loading it on first use.
func currentGeoIP() *geoipDB {
	geoipOnce.Do(func() {
		if db := loadGeoIP(); db != nil {
			swapGeoIP(db)
		}
	})
	geoipMu.RLock()
	defer geoipMu.RUnlock()
	return geoipCur
}

// swapGeoIP makes db the live database. The replaced reader is
// deliberately NOT closed: it is mmap-backed (or backed by the
// binary's own image), a lookup racing the swap could still be
// walking it, and an mmap cannot be safely unmapped under a reader.
// The mapping is one file per week at most and is reclaimed when the
// process exits.
func swapGeoIP(db *geoipDB) {
	geoipMu.Lock()
	geoipCur = db
	geoipMu.Unlock()
}

// loadGeoIP applies the layer order: a valid local file wins, else
// the embedded snapshot. A local file that will not open or will not
// answer falls through to the snapshot with one warning — a corrupt
// download must never take the flags (or the panel) down.
func loadGeoIP() *geoipDB {
	if path := GeoIPCountryLocalPath(); path != "" {
		if r, err := openGeoIPFile(path); err == nil {
			return &geoipDB{reader: r, origin: "local", cache: map[string]string{}}
		} else {
			logger.Warningf("geoip: local country database %s unusable (%v); using the embedded snapshot", path, err)
		}
	}
	if r, err := maxminddb.FromBytes(geoipCountryEmbedded); err == nil {
		if err := probeGeoIPReader(r); err == nil {
			return &geoipDB{reader: r, origin: "embedded", cache: map[string]string{}}
		} else {
			logger.Warningf("geoip: embedded country snapshot fails validation: %v", err)
		}
	} else {
		logger.Warningf("geoip: embedded country snapshot will not open: %v", err)
	}
	return &geoipDB{origin: "none", cache: map[string]string{}}
}

// openGeoIPFile opens a database file and checks it answers.
func openGeoIPFile(path string) (*maxminddb.Reader, error) {
	r, err := maxminddb.Open(path)
	if err != nil {
		return nil, err
	}
	if err := probeGeoIPReader(r); err != nil {
		_ = r.Close()
		return nil, err
	}
	return r, nil
}

// probeGeoIPReader is the functional validation: the well-known
// probe address must resolve to a country code, and a v6 probe must
// at least answer without an error. (maxminddb's Verify is NOT used:
// it also audits metadata niceties — the DB-IP/sapics files carry an
// empty description map and fail it despite being fully readable,
// verified against the factory snapshot on 2026-10-24. A candidate
// that opens and answers both probes is a working country DB; a
// truncated or foreign file fails to open or errors a lookup.)
func probeGeoIPReader(r *maxminddb.Reader) error {
	var rec countryRecord
	if err := r.Lookup(net.ParseIP(geoipCountryProbeIP), &rec); err != nil {
		return err
	}
	if rec.CountryCode == "" {
		return fmt.Errorf("probe address %s resolves to no country", geoipCountryProbeIP)
	}
	return r.Lookup(net.ParseIP(geoipCountryProbeIPv6), &rec)
}

// verifyGeoIPFile validates a fresh download in place before it may
// replace the live database.
func verifyGeoIPFile(path string) error {
	r, err := maxminddb.Open(path)
	if err != nil {
		return err
	}
	defer func() { _ = r.Close() }()
	return probeGeoIPReader(r)
}

// ---- lookups used by the public views ----

// countryCodeForIP resolves one public IP literal to its ISO country
// code, "" for anything non-public, unknown, or unanswerable. The
// public gate runs first so private/held addresses never reach the
// database (they have no country, exactly as on the hub).
func countryCodeForIP(ipStr string) string {
	ipStr = strings.TrimSpace(ipStr)
	if !IsPublicIPLiteral(ipStr) {
		return ""
	}
	db := currentGeoIP()
	if db == nil || db.reader == nil {
		return ""
	}
	return db.lookup(ipStr)
}

// probeCountryAuto is a node's auto-detected country: the first
// public address in the hub's order — the held interface v4, then
// held v6, then the master-observed exit v4, then exit v6 — resolved
// locally. "" when the node shows no public address or the database
// has no answer for it.
func probeCountryAuto(a *model.AgentNode) string {
	for _, s := range []string{a.IPv4, a.IPv6, a.SeenIPv4, a.SeenIPv6} {
		if IsPublicIPLiteral(s) {
			return countryCodeForIP(s)
		}
	}
	return ""
}

// probeCountry resolves the pair the views publish (hub api.rs): the
// effective country is the operator's pin when one is set, else the
// auto-detected code; country_auto always carries the detected code.
func probeCountry(pin string, a *model.AgentNode) (country, auto string) {
	auto = probeCountryAuto(a)
	if pin != "" {
		return pin, auto
	}
	return auto, auto
}

// ---- weekly auto-update ----

type geoipFetchResult int

const (
	geoipFetchUnchanged geoipFetchResult = iota
	geoipFetchUpdated
)

// downloadGeoIPCountry performs one conditional download of url into
// destPath. The candidate lands in a temp file in the destination
// directory, is fully validated there, and only then renames over
// destPath (atomic on the same filesystem); any failure leaves
// destPath and the live database untouched.
func downloadGeoIPCountry(client *http.Client, url, destPath, etag string, ifModifiedSince time.Time) (geoipFetchResult, string, error) {
	req, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		return geoipFetchUnchanged, etag, err
	}
	if etag != "" {
		req.Header.Set("If-None-Match", etag)
	}
	if !ifModifiedSince.IsZero() {
		req.Header.Set("If-Modified-Since", ifModifiedSince.UTC().Format(http.TimeFormat))
	}
	resp, err := client.Do(req)
	if err != nil {
		return geoipFetchUnchanged, etag, err
	}
	defer func() { _ = resp.Body.Close() }()
	if resp.StatusCode == http.StatusNotModified {
		return geoipFetchUnchanged, etag, nil
	}
	if resp.StatusCode != http.StatusOK {
		return geoipFetchUnchanged, etag, fmt.Errorf("update server answered %s", resp.Status)
	}

	dir := filepath.Dir(destPath)
	tmp, err := os.CreateTemp(dir, ".geoip-country-*.tmp")
	if err != nil {
		return geoipFetchUnchanged, etag, err
	}
	tmpName := tmp.Name()
	defer func() { _ = os.Remove(tmpName) }()
	n, err := io.Copy(tmp, io.LimitReader(resp.Body, geoipCountryMaxBytes))
	if err := tmp.Close(); err != nil {
		return geoipFetchUnchanged, etag, err
	}
	if err != nil {
		return geoipFetchUnchanged, etag, err
	}
	if n == 0 {
		return geoipFetchUnchanged, etag, fmt.Errorf("update download was empty")
	}
	if err := verifyGeoIPFile(tmpName); err != nil {
		return geoipFetchUnchanged, etag, fmt.Errorf("downloaded database rejected: %w", err)
	}
	if err := os.Rename(tmpName, destPath); err != nil {
		return geoipFetchUnchanged, etag, err
	}
	return geoipFetchUpdated, resp.Header.Get("ETag"), nil
}

// UpdateGeoIPCountry runs one update pass: a conditional download
// from the configured URL, a validated atomic replace of the local
// copy, and a hot swap of the live reader. No configured URL or any
// failure keeps the current database; a failure logs a single
// warning (the hub discipline: a broken mirror must not spam).
func UpdateGeoIPCountry() {
	url := geoipCountryUpdateURL()
	if url == "" {
		return
	}
	dest := GeoIPCountryLocalPath()
	var lastMod time.Time
	if fi, err := os.Stat(dest); err == nil {
		lastMod = fi.ModTime()
	}
	etag, _ := probeSetting(geoipCountryETagSetting)
	client := &http.Client{Timeout: 2 * time.Minute}
	res, newETag, err := downloadGeoIPCountry(client, url, dest, etag, lastMod)
	if err != nil {
		logger.Warningf("geoip: country database update from %s failed: %v (keeping the current database)", url, err)
		return
	}
	if res != geoipFetchUpdated {
		return
	}
	if newETag != "" && newETag != etag {
		_ = setProbeSetting(geoipCountryETagSetting, newETag)
	}
	r, err := openGeoIPFile(dest)
	if err != nil {
		logger.Warningf("geoip: updated country database %s will not reopen: %v", dest, err)
		return
	}
	swapGeoIP(&geoipDB{reader: r, origin: "local", cache: map[string]string{}})
	logger.Infof("geoip: country database updated from %s", url)
}