geoip: tests — known-address calibration, source order, pin precedence, updater download/validate/replace, contract frame

2026-10-23 16:42:28 +00:00
parent c94136c8a1
commit d16b888146
2 changed files with 317 additions and 0 deletions
@@ -601,3 +601,61 @@ func TestProbeWebSocket(t *testing.T) {
 func itoa(v int64) string {
 	return strconv.FormatInt(v, 10)
 }
+
+// Q7: the resolved country reaches both frames. The seed's held
+// address is a documentation range (unlisted in any country DB), so
+// repoint it at a real public address whose country the embedded
+// DB-IP snapshot knows — 8.8.8.8 is US there.
+func TestProbePublicCountryFlag(t *testing.T) {
+	rig := newProbeTestRig(t)
+	pubID, _ := rig.seedAgent(t, "FlagNode", true)
+
+	db := database.GetDB()
+	var agent model.AgentNode
+	if err := db.First(&agent, pubID).Error; err != nil {
+		t.Fatalf("load seeded agent: %v", err)
+	}
+	agent.IPv4 = "8.8.8.8"
+	if err := db.Save(&agent).Error; err != nil {
+		t.Fatalf("repoint agent address: %v", err)
+	}
+
+	rig.setSetting(t, "public_page", "on")
+	service.InvalidateProbeSnapshot()
+
+	find := func(rec *httptest.ResponseRecorder) map[string]any {
+		t.Helper()
+		var frame struct {
+			Nodes []map[string]any `json:"nodes"`
+			Admin bool             `json:"admin"`
+		}
+		if err := json.Unmarshal(rec.Body.Bytes(), &frame); err != nil {
+			t.Fatalf("frame decode: %v", err)
+		}
+		for _, n := range frame.Nodes {
+			if n["id"] == float64(pubID) {
+				return n
+			}
+		}
+		t.Fatalf("seeded node missing from frame: %s", rec.Body.String())
+		return nil
+	}
+
+	rec := rig.do(http.MethodGet, "/api/nodes", nil, false)
+	if rec.Code != 200 {
+		t.Fatalf("nodes anon: %d", rec.Code)
+	}
+	if got := find(rec)["country"]; got != "US" {
+		t.Fatalf("anonymous country = %v, want US", got)
+	}
+
+	rec = rig.do(http.MethodGet, "/api/nodes", nil, true)
+	if rec.Code != 200 {
+		t.Fatalf("nodes admin: %d", rec.Code)
+	}
+	admin := find(rec)
+	if admin["country"] != "US" || admin["country_auto"] != "US" || admin["country_pin"] != "" {
+		t.Fatalf("admin country triple = (%v, %v, %v), want (US, US, empty)",
+			admin["country"], admin["country_auto"], admin["country_pin"])
+	}
+}
@@ -0,0 +1,259 @@
+package service
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+
+	"github.com/oschwald/maxminddb-golang"
+
+	"github.com/mhsanaei/3x-ui/v3/internal/database/model"
+	"github.com/mhsanaei/3x-ui/v3/internal/web/service/probe"
+)
+
+// Q7 node country flags: local MMDB resolution on the master. The
+// known-address expectations below are calibrated against the
+// embedded factory snapshot (DB-IP Lite country database, retrieved
+// 2026-10-24; provenance in geoipdata/NOTICE): these are long-lived
+// anycast/ISP addresses whose country attribution is stable across
+// editions — 8.8.8.8 (Google DNS) US, 223.5.5.5 (AliDNS) CN,
+// 1.1.1.1 (Cloudflare/APNIC) AU, 2400:3200::1 (AliDNS v6) CN.
+
+// useEmbeddedGeoIP swaps the embedded snapshot in as the live
+// database for one test, restoring whatever was live before.
+func useEmbeddedGeoIP(t *testing.T) {
+	t.Helper()
+	prev := currentGeoIP() // fire the lazy load so the Once is spent
+	r, err := maxminddb.FromBytes(geoipCountryEmbedded)
+	if err != nil {
+		t.Fatalf("embedded snapshot will not open: %v", err)
+	}
+	swapGeoIP(&geoipDB{reader: r, origin: "embedded", cache: map[string]string{}})
+	t.Cleanup(func() { swapGeoIP(prev) })
+}
+
+func TestCountryCodeForIPKnownAddresses(t *testing.T) {
+	useEmbeddedGeoIP(t)
+	for ip, want := range map[string]string{
+		"8.8.8.8":      "US",
+		"223.5.5.5":    "CN",
+		"1.1.1.1":      "AU",
+		"2400:3200::1": "CN",
+		" 8.8.8.8 ":    "US", // surrounding space is trimmed
+	} {
+		if got := countryCodeForIP(ip); got != want {
+			t.Errorf("countryCodeForIP(%q) = %q, want %q", ip, got, want)
+		}
+	}
+	// Non-public and unlisted addresses have no country.
+	for _, ip := range []string{
+		"192.168.0.1", // RFC1918
+		"10.0.0.5",    // RFC1918
+		"127.0.0.1",   // loopback
+		"100.64.0.1",  // CGNAT shared range
+		"203.0.113.9", // public documentation range: unlisted in the DB
+		"",            // empty
+		"example.com", // not a literal
+		"999.1.1.1",   // not an address
+	} {
+		if got := countryCodeForIP(ip); got != "" {
+			t.Errorf("countryCodeForIP(%q) = %q, want empty", ip, got)
+		}
+	}
+}
+
+func TestProbeCountryAutoSourceOrder(t *testing.T) {
+	useEmbeddedGeoIP(t)
+	// Held v4 leads over held v6 (hub order).
+	a := &model.AgentNode{IPv4: "8.8.8.8", IPv6: "2400:3200::1", SeenIPv4: "223.5.5.5"}
+	if got := probeCountryAuto(a); got != "US" {
+		t.Errorf("held v4 first: got %q, want US", got)
+	}
+	// A private held v4 is skipped; the held public v6 beats the
+	// observed exit v4.
+	a = &model.AgentNode{IPv4: "10.0.0.5", IPv6: "2400:3200::1", SeenIPv4: "8.8.8.8"}
+	if got := probeCountryAuto(a); got != "CN" {
+		t.Errorf("held v6 over exit v4: got %q, want CN", got)
+	}
+	// Exit addresses are the fallback when the node holds none.
+	a = &model.AgentNode{SeenIPv4: "223.5.5.5"}
+	if got := probeCountryAuto(a); got != "CN" {
+		t.Errorf("exit v4 fallback: got %q, want CN", got)
+	}
+	a = &model.AgentNode{SeenIPv6: "2400:3200::1"}
+	if got := probeCountryAuto(a); got != "CN" {
+		t.Errorf("exit v6 fallback: got %q, want CN", got)
+	}
+	// No public address anywhere: no country.
+	a = &model.AgentNode{IPv4: "192.168.1.2", SeenIPv4: "10.1.1.1"}
+	if got := probeCountryAuto(a); got != "" {
+		t.Errorf("private only: got %q, want empty", got)
+	}
+}
+
+func TestProbeCountryPinWins(t *testing.T) {
+	useEmbeddedGeoIP(t)
+	a := &model.AgentNode{IPv4: "223.5.5.5"}
+	country, auto := probeCountry("JP", a)
+	if country != "JP" || auto != "CN" {
+		t.Errorf("pin set: got (%q, %q), want (JP, CN)", country, auto)
+	}
+	country, auto = probeCountry("", a)
+	if country != "CN" || auto != "CN" {
+		t.Errorf("no pin: got (%q, %q), want (CN, CN)", country, auto)
+	}
+	country, auto = probeCountry("JP", &model.AgentNode{})
+	if country != "JP" || auto != "" {
+		t.Errorf("pin without address: got (%q, %q), want (JP, empty)", country, auto)
+	}
+}
+
+func TestBuildProbeNodeViewCountry(t *testing.T) {
+	useEmbeddedGeoIP(t)
+	src := probeNodeSource{
+		name:  "flagged",
+		agent: model.AgentNode{IPv4: "8.8.8.8"},
+		meta:  probe.Meta{Public: true},
+	}
+	anon := buildProbeNodeView(src, false)
+	if anon["country"] != "US" {
+		t.Errorf("anonymous country = %v, want US", anon["country"])
+	}
+	admin := buildProbeNodeView(src, true)
+	if admin["country"] != "US" || admin["country_auto"] != "US" || admin["country_pin"] != "" {
+		t.Errorf("admin pair = (%v, %v, pin %v), want (US, US, empty)",
+			admin["country"], admin["country_auto"], admin["country_pin"])
+	}
+	src.meta.CountryPin = "DE"
+	if got := buildProbeNodeView(src, false)["country"]; got != "DE" {
+		t.Errorf("pinned country = %v, want DE", got)
+	}
+	if got := buildProbeNodeView(src, true)["country_auto"]; got != "US" {
+		t.Errorf("country_auto under a pin = %v, want US", got)
+	}
+	// A node with no public address keeps the pre-flag empty country.
+	src = probeNodeSource{name: "dark", agent: model.AgentNode{IPv4: "192.168.0.9"}, meta: probe.Meta{Public: true}}
+	if got := buildProbeNodeView(src, false)["country"]; got != "" {
+		t.Errorf("private node country = %v, want empty", got)
+	}
+}
+
+func TestOpenGeoIPFileValidation(t *testing.T) {
+	dir := t.TempDir()
+	bad := filepath.Join(dir, "bad.mmdb")
+	if err := os.WriteFile(bad, []byte("this is not an mmdb"), 0o644); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := openGeoIPFile(bad); err == nil {
+		t.Fatal("garbage file opened as a country database")
+	}
+	good := filepath.Join(dir, "good.mmdb")
+	if err := os.WriteFile(good, geoipCountryEmbedded, 0o644); err != nil {
+		t.Fatal(err)
+	}
+	r, err := openGeoIPFile(good)
+	if err != nil {
+		t.Fatalf("embedded bytes as a file rejected: %v", err)
+	}
+	defer func() { _ = r.Close() }()
+	db := &geoipDB{reader: r, cache: map[string]string{}}
+	if got := db.lookup("223.5.5.5"); got != "CN" {
+		t.Errorf("file-backed lookup = %q, want CN", got)
+	}
+}
+
+// ---- the weekly updater's download/validate/replace core ----
+
+func TestDownloadGeoIPCountryUpdates(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("ETag", `"db-v1"`)
+		_, _ = w.Write(geoipCountryEmbedded)
+	}))
+	defer srv.Close()
+
+	dest := filepath.Join(t.TempDir(), "geoip-country.mmdb")
+	res, etag, err := downloadGeoIPCountry(srv.Client(), srv.URL, dest, "", time.Time{})
+	if err != nil {
+		t.Fatalf("download: %v", err)
+	}
+	if res != geoipFetchUpdated || etag != `"db-v1"` {
+		t.Fatalf("got (res %v, etag %q), want (updated, \"db-v1\")", res, etag)
+	}
+	fi, err := os.Stat(dest)
+	if err != nil || fi.Size() != int64(len(geoipCountryEmbedded)) {
+		t.Fatalf("dest stat = %v (err %v), want %d bytes", fi, err, len(geoipCountryEmbedded))
+	}
+	// The replaced file is a working database.
+	r, err := openGeoIPFile(dest)
+	if err != nil {
+		t.Fatalf("replaced file will not open: %v", err)
+	}
+	_ = r.Close()
+}
+
+func TestDownloadGeoIPCountryNotModified(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Header.Get("If-None-Match") == `"db-v1"` || r.Header.Get("If-Modified-Since") != "" {
+			w.WriteHeader(http.StatusNotModified)
+			return
+		}
+		_, _ = w.Write(geoipCountryEmbedded)
+	}))
+	defer srv.Close()
+
+	dir := t.TempDir()
+	dest := filepath.Join(dir, "geoip-country.mmdb")
+	if err := os.WriteFile(dest, []byte("ORIGINAL"), 0o644); err != nil {
+		t.Fatal(err)
+	}
+	// ETag path.
+	res, _, err := downloadGeoIPCountry(srv.Client(), srv.URL, dest, `"db-v1"`, time.Time{})
+	if err != nil || res != geoipFetchUnchanged {
+		t.Fatalf("etag path: got (res %v, err %v), want unchanged", res, err)
+	}
+	// If-Modified-Since path.
+	res, _, err = downloadGeoIPCountry(srv.Client(), srv.URL, dest, "", time.Now())
+	if err != nil || res != geoipFetchUnchanged {
+		t.Fatalf("modified-since path: got (res %v, err %v), want unchanged", res, err)
+	}
+	body, _ := os.ReadFile(dest)
+	if string(body) != "ORIGINAL" {
+		t.Fatalf("dest rewritten on a 304: %q", body)
+	}
+}
+
+func TestDownloadGeoIPCountryRejectsBad(t *testing.T) {
+	dir := t.TempDir()
+	dest := filepath.Join(dir, "geoip-country.mmdb")
+	if err := os.WriteFile(dest, []byte("ORIGINAL"), 0o644); err != nil {
+		t.Fatal(err)
+	}
+	assertKept := func(name string) {
+		t.Helper()
+		body, _ := os.ReadFile(dest)
+		if string(body) != "ORIGINAL" {
+			t.Fatalf("%s: dest replaced despite the failure: %q", name, body)
+		}
+	}
+
+	garbage := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		_, _ = w.Write([]byte("definitely not an mmdb, just bytes"))
+	}))
+	defer garbage.Close()
+	if _, _, err := downloadGeoIPCountry(garbage.Client(), garbage.URL, dest, "", time.Time{}); err == nil {
+		t.Fatal("garbage download accepted")
+	}
+	assertKept("garbage")
+
+	broken := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusInternalServerError)
+	}))
+	defer broken.Close()
+	if _, _, err := downloadGeoIPCountry(broken.Client(), broken.URL, dest, "", time.Time{}); err == nil {
+		t.Fatal("500 download accepted")
+	}
+	assertKept("500")
+}