3x-ui-agent/internal/xray/geodata/geodata_test.go

package geodata

import (
	"encoding/json"
	"errors"
	"net/netip"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	"google.golang.org/protobuf/encoding/protowire"
)

// The fixtures are encoded with protowire directly: the databases' wire
// layout is the contract this package parses, and spelling it out keeps the
// tests independent of xray-core's generated types. Field numbers follow
// geodat.proto — list entry 1; entry code 1, payload 2; domain type 1,
// value 2, attribute 3; attribute key 1, bool value 2; cidr ip 1, prefix 2.
// proto3 zero-value omission is reproduced exactly (a domain of type 0 and
// a /0 prefix emit no field), because the reader's behavior depends on it.

// Domain match types, mirroring xray's Domain_Type enum.
const (
	DomainSubstr = 0
	DomainRegex  = 1
	DomainDomain = 2
	DomainFull   = 3
)

type fixtureDomain struct {
	typ   uint64
	value string
	attrs []string
}

type fixtureSite struct {
	code    string
	domains []fixtureDomain
}

type fixtureIP struct {
	code     string
	prefixes []string
}

func wireBytes(num protowire.Number, payload []byte) []byte {
	out := protowire.AppendTag(nil, num, protowire.BytesType)
	return protowire.AppendBytes(out, payload)
}

func wireString(num protowire.Number, value string) []byte {
	out := protowire.AppendTag(nil, num, protowire.BytesType)
	return protowire.AppendString(out, value)
}

func wireVarint(num protowire.Number, value uint64) []byte {
	out := protowire.AppendTag(nil, num, protowire.VarintType)
	return protowire.AppendVarint(out, value)
}

func encodeFixtureDomain(d fixtureDomain) []byte {
	var msg []byte
	if d.typ != 0 {
		msg = append(msg, wireVarint(1, d.typ)...)
	}
	msg = append(msg, wireString(2, d.value)...)
	for _, attr := range d.attrs {
		attrMsg := wireString(1, attr)
		attrMsg = append(attrMsg, wireVarint(2, 1)...) // bool_value = true
		msg = append(msg, wireBytes(3, attrMsg)...)
	}
	return msg
}

func writeSiteDB(t *testing.T, dir, name string, sites ...fixtureSite) string {
	t.Helper()
	var data []byte
	for _, s := range sites {
		entry := wireString(1, s.code)
		for _, d := range s.domains {
			entry = append(entry, wireBytes(2, encodeFixtureDomain(d))...)
		}
		data = append(data, wireBytes(1, entry)...)
	}
	return writeFile(t, dir, name, data)
}

func writeIPDB(t *testing.T, dir, name string, geoips ...fixtureIP) string {
	t.Helper()
	var data []byte
	for _, g := range geoips {
		entry := wireString(1, g.code)
		for _, raw := range g.prefixes {
			prefix := netip.MustParsePrefix(raw)
			cidr := wireBytes(1, prefix.Addr().AsSlice())
			if prefix.Bits() != 0 {
				cidr = append(cidr, wireVarint(2, uint64(prefix.Bits()))...)
			}
			entry = append(entry, wireBytes(2, cidr)...)
		}
		data = append(data, wireBytes(1, entry)...)
	}
	return writeFile(t, dir, name, data)
}

func writeFile(t *testing.T, dir, name string, data []byte) string {
	t.Helper()
	path := filepath.Join(dir, name)
	if err := os.WriteFile(path, data, 0o644); err != nil {
		t.Fatalf("write %s: %v", name, err)
	}
	return path
}

func site(code string, domains ...fixtureDomain) fixtureSite {
	return fixtureSite{code: code, domains: domains}
}

func domain(domainType uint64, value string, attributes ...string) fixtureDomain {
	return fixtureDomain{typ: domainType, value: value, attrs: attributes}
}

func geoip(code string, prefixes ...string) fixtureIP {
	return fixtureIP{code: code, prefixes: prefixes}
}

func sampleSiteDB(t *testing.T, dir string) {
	t.Helper()
	writeSiteDB(t, dir, "geosite.dat",
		site("google",
			domain(DomainDomain, "google.com"),
			domain(DomainFull, "ads.google.com", "ads"),
			domain(DomainSubstr, "googlevideo", "cn"),
			domain(DomainRegex, `^g.*\.cn$`),
		),
		site("CN",
			domain(DomainDomain, "baidu.com"),
			domain(DomainDomain, "qq.com"),
		),
	)
}

func TestListFilesReportsKindAndCategories(t *testing.T) {
	dir := t.TempDir()
	sampleSiteDB(t, dir)
	writeIPDB(t, dir, "geoip.dat", geoip("cn", "1.0.1.0/24"), geoip("private", "10.0.0.0/8", "fc00::/7"))

	files, err := NewStore(dir).ListFiles()
	if err != nil {
		t.Fatalf("ListFiles() error = %v", err)
	}
	if len(files) != 2 {
		t.Fatalf("ListFiles() returned %d files, want 2", len(files))
	}

	byName := make(map[string]GeoFile, len(files))
	for _, file := range files {
		byName[file.Name] = file
	}

	geosite := byName["geosite.dat"]
	if geosite.Kind != KindSite {
		t.Errorf("geosite.dat kind = %q, want %q", geosite.Kind, KindSite)
	}
	if geosite.Categories != 2 {
		t.Errorf("geosite.dat categories = %d, want 2", geosite.Categories)
	}
	if geosite.Error != "" {
		t.Errorf("geosite.dat error = %q, want empty", geosite.Error)
	}

	geoipFile := byName["geoip.dat"]
	if geoipFile.Kind != KindIP {
		t.Errorf("geoip.dat kind = %q, want %q", geoipFile.Kind, KindIP)
	}
	if geoipFile.Categories != 2 {
		t.Errorf("geoip.dat categories = %d, want 2", geoipFile.Categories)
	}
}

func TestKindDetectedFromContentsNotName(t *testing.T) {
	dir := t.TempDir()
	writeSiteDB(t, dir, "my_ip_rules.dat", site("corp", domain(DomainDomain, "intranet.corp.local")))
	writeIPDB(t, dir, "custom_sites.dat", geoip("office", "192.168.7.0/24"))

	store := NewStore(dir)

	sitePage, err := store.Categories("my_ip_rules.dat", "", 0, 10)
	if err != nil {
		t.Fatalf("Categories(my_ip_rules.dat) error = %v", err)
	}
	if sitePage.Total != 1 || sitePage.Items[0].Code != "corp" {
		t.Fatalf("Categories(my_ip_rules.dat) = %+v, want single category corp", sitePage)
	}

	entries, err := store.Entries("custom_sites.dat", "office", "", 0, 10)
	if err != nil {
		t.Fatalf("Entries(custom_sites.dat) error = %v", err)
	}
	if len(entries.Items) != 1 {
		t.Fatalf("Entries(custom_sites.dat) returned %d items, want 1", len(entries.Items))
	}
	if got := entries.Items[0]; got.Kind != "cidr" || got.Value != "192.168.7.0/24" {
		t.Errorf("entry = %+v, want cidr 192.168.7.0/24", got)
	}
}

func TestEntriesMapDomainTypesAndAttributes(t *testing.T) {
	dir := t.TempDir()
	sampleSiteDB(t, dir)
	store := NewStore(dir)

	page, err := store.Entries("geosite.dat", "google", "", 0, 10)
	if err != nil {
		t.Fatalf("Entries() error = %v", err)
	}
	want := []GeoEntry{
		{Kind: "domain", Value: "google.com"},
		{Kind: "full", Value: "ads.google.com"},
		{Kind: "keyword", Value: "googlevideo"},
		{Kind: "regexp", Value: `^g.*\.cn$`},
	}
	if page.Total != len(want) {
		t.Fatalf("Entries() total = %d, want %d", page.Total, len(want))
	}
	for i, entry := range want {
		if page.Items[i] != entry {
			t.Errorf("entry %d = %+v, want %+v", i, page.Items[i], entry)
		}
	}

	category, err := store.Lookup("geosite.dat", "google")
	if err != nil {
		t.Fatalf("Lookup() error = %v", err)
	}
	if len(category.Attributes) != 2 || category.Attributes[0] != "ads" || category.Attributes[1] != "cn" {
		t.Errorf("attributes = %v, want [ads cn]", category.Attributes)
	}
}

func TestCategoriesWithoutAttributesMarshalAsEmptyArray(t *testing.T) {
	dir := t.TempDir()
	sampleSiteDB(t, dir)

	page, err := NewStore(dir).Categories("geosite.dat", "cn", 0, 10)
	if err != nil {
		t.Fatalf("Categories() error = %v", err)
	}
	if page.Items[0].Attributes == nil {
		t.Fatal("attributes are nil, want an empty slice so the JSON stays an array")
	}
	encoded, err := json.Marshal(page.Items[0])
	if err != nil {
		t.Fatalf("marshal category: %v", err)
	}
	if !strings.Contains(string(encoded), `"attributes":[]`) {
		t.Errorf("encoded category = %s, want an empty attributes array", encoded)
	}
}

func TestCategoryCodesAreLowercasedAndSorted(t *testing.T) {
	dir := t.TempDir()
	sampleSiteDB(t, dir)

	page, err := NewStore(dir).Categories("geosite.dat", "", 0, 10)
	if err != nil {
		t.Fatalf("Categories() error = %v", err)
	}
	if page.Items[0].Code != "cn" || page.Items[1].Code != "google" {
		t.Errorf("codes = %q, %q; want cn, google", page.Items[0].Code, page.Items[1].Code)
	}
}

func TestSearchFilters(t *testing.T) {
	dir := t.TempDir()
	sampleSiteDB(t, dir)
	store := NewStore(dir)

	categories, err := store.Categories("geosite.dat", "OOG", 0, 10)
	if err != nil {
		t.Fatalf("Categories() error = %v", err)
	}
	if categories.Total != 1 || categories.Items[0].Code != "google" {
		t.Errorf("Categories(OOG) = %+v, want only google", categories)
	}

	entries, err := store.Entries("geosite.dat", "google", "ADS.", 0, 10)
	if err != nil {
		t.Fatalf("Entries() error = %v", err)
	}
	if entries.Total != 1 || entries.Items[0].Value != "ads.google.com" {
		t.Errorf("Entries(ADS.) = %+v, want only ads.google.com", entries)
	}
}

func TestPagination(t *testing.T) {
	dir := t.TempDir()
	domains := make([]fixtureDomain, 0, 250)
	for i := range 250 {
		domains = append(domains, domain(DomainDomain, "host"+strconv.Itoa(i)+".example.com"))
	}
	writeSiteDB(t, dir, "geosite.dat", site("bulk", domains...))
	store := NewStore(dir)

	tests := []struct {
		name      string
		offset    int
		limit     int
		wantCount int
		wantFirst string
	}{
		{name: "first page", offset: 0, limit: 10, wantCount: 10, wantFirst: "host0.example.com"},
		{name: "middle page", offset: 20, limit: 5, wantCount: 5, wantFirst: "host20.example.com"},
		{name: "negative offset clamps to start", offset: -5, limit: 3, wantCount: 3, wantFirst: "host0.example.com"},
		{name: "tail shorter than limit", offset: 245, limit: 50, wantCount: 5, wantFirst: "host245.example.com"},
		{name: "offset past end", offset: 900, limit: 10, wantCount: 0},
		{name: "limit above cap", offset: 0, limit: 5000, wantCount: 250, wantFirst: "host0.example.com"},
		{name: "zero limit uses cap", offset: 0, limit: 0, wantCount: 250, wantFirst: "host0.example.com"},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			page, err := store.Entries("geosite.dat", "bulk", "", tt.offset, tt.limit)
			if err != nil {
				t.Fatalf("Entries() error = %v", err)
			}
			if page.Total != 250 {
				t.Errorf("total = %d, want 250", page.Total)
			}
			if len(page.Items) != tt.wantCount {
				t.Fatalf("items = %d, want %d", len(page.Items), tt.wantCount)
			}
			if tt.wantFirst != "" && page.Items[0].Value != tt.wantFirst {
				t.Errorf("first item = %q, want %q", page.Items[0].Value, tt.wantFirst)
			}
		})
	}
}

func TestCategoriesReturnEverythingWithoutLimit(t *testing.T) {
	dir := t.TempDir()
	sites := make([]fixtureSite, 0, MaxPageSize+20)
	for i := range MaxPageSize + 20 {
		sites = append(sites, site("cat"+strconv.Itoa(i), domain(DomainDomain, "example.com")))
	}
	writeSiteDB(t, dir, "geosite.dat", sites...)
	store := NewStore(dir)

	all, err := store.Categories("geosite.dat", "", 0, 0)
	if err != nil {
		t.Fatalf("Categories() error = %v", err)
	}
	if len(all.Items) != MaxPageSize+20 {
		t.Errorf("items without a limit = %d, want %d", len(all.Items), MaxPageSize+20)
	}

	capped, err := store.Categories("geosite.dat", "", 0, 10)
	if err != nil {
		t.Fatalf("Categories() error = %v", err)
	}
	if len(capped.Items) != 10 || capped.Total != MaxPageSize+20 {
		t.Errorf("explicit limit gave %d items with total %d, want 10 and %d", len(capped.Items), capped.Total, MaxPageSize+20)
	}

	entries, err := store.Entries("geosite.dat", "cat0", "", 0, 0)
	if err != nil {
		t.Fatalf("Entries() error = %v", err)
	}
	if len(entries.Items) != 1 {
		t.Errorf("entries = %d, want 1", len(entries.Items))
	}
}

func TestErrors(t *testing.T) {
	dir := t.TempDir()
	sampleSiteDB(t, dir)
	writeFile(t, dir, "broken.dat", []byte("this is not a protobuf message at all"))
	store := NewStore(dir)

	tests := []struct {
		name string
		call func() error
		want error
	}{
		{
			name: "unknown category",
			call: func() error { _, err := store.Entries("geosite.dat", "nope", "", 0, 10); return err },
			want: ErrUnknownCategory,
		},
		{
			name: "lookup of unknown category",
			call: func() error { _, err := store.Lookup("geosite.dat", "nope"); return err },
			want: ErrUnknownCategory,
		},
		{
			name: "path traversal",
			call: func() error { _, err := store.Categories("../geosite.dat", "", 0, 10); return err },
			want: ErrInvalidName,
		},
		{
			name: "non dat extension",
			call: func() error { _, err := store.Categories("x-ui.db", "", 0, 10); return err },
			want: ErrInvalidName,
		},
		{
			name: "empty name",
			call: func() error { _, err := store.Categories("", "", 0, 10); return err },
			want: ErrInvalidName,
		},
		{
			name: "unparsable file",
			call: func() error { _, err := store.Categories("broken.dat", "", 0, 10); return err },
			want: ErrUnrecognized,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if err := tt.call(); !errors.Is(err, tt.want) {
				t.Errorf("error = %v, want %v", err, tt.want)
			}
		})
	}
}

func TestBrokenFileIsListedWithReason(t *testing.T) {
	dir := t.TempDir()
	writeFile(t, dir, "broken.dat", []byte("not a database"))

	files, err := NewStore(dir).ListFiles()
	if err != nil {
		t.Fatalf("ListFiles() error = %v", err)
	}
	if len(files) != 1 {
		t.Fatalf("ListFiles() returned %d files, want 1", len(files))
	}
	if !strings.HasPrefix(files[0].Error, ErrUnrecognized.Error()) {
		t.Errorf("error = %q, want it to start with %q", files[0].Error, ErrUnrecognized.Error())
	}
	if files[0].Kind != "" {
		t.Errorf("kind = %q, want empty", files[0].Kind)
	}
}

func TestFileAboveSizeLimitIsRejected(t *testing.T) {
	dir := t.TempDir()
	path := writeFile(t, dir, "huge.dat", []byte("x"))
	if err := os.Truncate(path, MaxFileSize+1); err != nil {
		t.Fatalf("truncate: %v", err)
	}

	store := NewStore(dir)
	if _, err := store.Categories("huge.dat", "", 0, 10); !errors.Is(err, ErrFileTooLarge) {
		t.Errorf("error = %v, want %v", err, ErrFileTooLarge)
	}

	files, err := store.ListFiles()
	if err != nil {
		t.Fatalf("ListFiles() error = %v", err)
	}
	if len(files) != 1 || files[0].Error != ErrFileTooLarge.Error() {
		t.Errorf("ListFiles() = %+v, want the file listed with a too-large error", files)
	}
}

func TestIndexCacheInvalidatedWhenFileChanges(t *testing.T) {
	dir := t.TempDir()
	sampleSiteDB(t, dir)
	store := NewStore(dir)

	before, err := store.Categories("geosite.dat", "", 0, 10)
	if err != nil {
		t.Fatalf("Categories() error = %v", err)
	}
	if before.Total != 2 {
		t.Fatalf("total before rewrite = %d, want 2", before.Total)
	}

	path := writeSiteDB(t, dir, "geosite.dat",
		site("google", domain(DomainDomain, "google.com")),
		site("cn", domain(DomainDomain, "baidu.com")),
		site("telegram", domain(DomainDomain, "t.me")),
	)
	touch(t, path, time.Now().Add(time.Second))

	after, err := store.Categories("geosite.dat", "", 0, 10)
	if err != nil {
		t.Fatalf("Categories() after rewrite error = %v", err)
	}
	if after.Total != 3 {
		t.Errorf("total after rewrite = %d, want 3", after.Total)
	}
	if len(store.indexes) != 1 {
		t.Errorf("cached indexes = %d, want 1 after the stale entry is dropped", len(store.indexes))
	}
}

func touch(t *testing.T, path string, when time.Time) {
	t.Helper()
	if err := os.Chtimes(path, when, when); err != nil {
		t.Fatalf("chtimes %s: %v", path, err)
	}
}

func TestDefaultRouteCIDRSurvives(t *testing.T) {
	dir := t.TempDir()
	writeIPDB(t, dir, "geoip.dat", geoip("any", "0.0.0.0/0", "::/0"), geoip("cn", "1.0.1.0/24"))

	page, err := NewStore(dir).Entries("geoip.dat", "any", "", 0, 10)
	if err != nil {
		t.Fatalf("Entries() error = %v", err)
	}
	if page.Total != 2 {
		t.Fatalf("total = %d, want 2 — a zero prefix is omitted by proto3 and must not be dropped", page.Total)
	}
	if page.Items[0].Value != "0.0.0.0/0" || page.Items[1].Value != "::/0" {
		t.Errorf("items = %+v, want the two default routes", page.Items)
	}
}

func TestBrokenFileIsParsedOnlyOnce(t *testing.T) {
	dir := t.TempDir()
	writeFile(t, dir, "broken.dat", []byte("not a database"))
	store := NewStore(dir)

	for range 3 {
		if _, err := store.Categories("broken.dat", "", 0, 10); !errors.Is(err, ErrUnrecognized) {
			t.Fatalf("error = %v, want %v", err, ErrUnrecognized)
		}
	}
	if len(store.indexes) != 1 {
		t.Errorf("cached indexes = %d, want the failure cached once", len(store.indexes))
	}
}

func TestConcurrentReadsAreConsistent(t *testing.T) {
	dir := t.TempDir()
	sampleSiteDB(t, dir)
	writeIPDB(t, dir, "geoip.dat", geoip("private", "10.0.0.0/8"))
	store := NewStore(dir)

	var wg sync.WaitGroup
	for i := range 24 {
		wg.Add(1)
		go func(worker int) {
			defer wg.Done()
			switch worker % 3 {
			case 0:
				page, err := store.Categories("geosite.dat", "", 0, 0)
				if err != nil || page.Total != 2 {
					t.Errorf("Categories() = %+v, err = %v; want 2 categories", page, err)
				}
			case 1:
				page, err := store.Entries("geosite.dat", "google", "", 0, 10)
				if err != nil || page.Total != 4 {
					t.Errorf("Entries() = %+v, err = %v; want 4 entries", page, err)
				}
			default:
				files, err := store.ListFiles()
				if err != nil || len(files) != 2 {
					t.Errorf("ListFiles() = %d files, err = %v; want 2 files", len(files), err)
				}
			}
		}(i)
	}
	wg.Wait()
}

func TestLookupDoesNotForgiveStraySpaces(t *testing.T) {
	dir := t.TempDir()
	sampleSiteDB(t, dir)
	store := NewStore(dir)

	if _, err := store.Lookup("geosite.dat", "google"); err != nil {
		t.Fatalf("Lookup(google) error = %v", err)
	}
	for _, code := range []string{" google", "google ", "goo gle"} {
		if _, err := store.Lookup("geosite.dat", code); !errors.Is(err, ErrUnknownCategory) {
			t.Errorf("Lookup(%q) error = %v, want %v — the core does not trim either", code, err, ErrUnknownCategory)
		}
	}
}

func TestSymlinkOutOfTheAssetFolderIsRefused(t *testing.T) {
	outside := t.TempDir()
	secret := filepath.Join(outside, "secret.dat")
	if err := os.WriteFile(secret, []byte("not yours"), 0o644); err != nil {
		t.Fatalf("write secret: %v", err)
	}

	dir := t.TempDir()
	sampleSiteDB(t, dir)
	if err := os.Symlink(secret, filepath.Join(dir, "escape.dat")); err != nil {
		t.Skipf("symlinks unavailable: %v", err)
	}

	store := NewStore(dir)
	if _, err := store.Categories("escape.dat", "", 0, 10); err == nil {
		t.Error("Categories() read through a symlink pointing outside the asset folder")
	}
	if _, err := store.Entries("escape.dat", "google", "", 0, 10); err == nil {
		t.Error("Entries() read through a symlink pointing outside the asset folder")
	}

	files, err := store.ListFiles()
	if err != nil {
		t.Fatalf("ListFiles() error = %v", err)
	}
	for _, file := range files {
		if file.Name == "escape.dat" && file.Error == "" {
			t.Error("ListFiles() reported an escaping symlink as a usable database")
		}
	}
}