260 lines
8.7 KiB
Go
260 lines
8.7 KiB
Go
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")
|
|
}
|