Mellaris/ruleset/builtins/geo/geo_matcher_test.go

package geo

import (
	"sync/atomic"
	"testing"

	"git.difuse.io/Difuse/Mellaris/ruleset/builtins/geo/v2geo"
)

type fakeGeoLoader struct {
	geoip   map[string]*v2geo.GeoIP
	geosite map[string]*v2geo.GeoSite
}

func (l *fakeGeoLoader) LoadGeoIP() (map[string]*v2geo.GeoIP, error) {
	return l.geoip, nil
}

func (l *fakeGeoLoader) LoadGeoSite() (map[string]*v2geo.GeoSite, error) {
	return l.geosite, nil
}

func TestGeoMatcher_MatchGeoIp_Cached(t *testing.T) {
	loader := &fakeGeoLoader{
		geoip: map[string]*v2geo.GeoIP{
			"us": {
				Cidr: []*v2geo.CIDR{
					{Ip: ipv4(8, 8, 8, 0), Prefix: 24},
				},
			},
		},
	}
	g := NewGeoMatcher("", "")
	g.geoLoader = loader

	if !g.MatchGeoIp("8.8.8.8", "US") {
		t.Error("MatchGeoIp should match 8.8.8.8 in US range")
	}
	if g.MatchGeoIp("9.9.9.9", "US") {
		t.Error("MatchGeoIp should not match 9.9.9.9 in US range")
	}
}

func TestGeoMatcher_MatchGeoIp_EmptyCondition(t *testing.T) {
	g := NewGeoMatcher("", "")
	if g.MatchGeoIp("1.2.3.4", "") {
		t.Error("MatchGeoIp with empty condition should return false")
	}
}

func TestGeoMatcher_MatchGeoIp_InvalidIP(t *testing.T) {
	g := NewGeoMatcher("", "")
	if g.MatchGeoIp("not-an-ip", "us") {
		t.Error("MatchGeoIp with invalid IP should return false")
	}
}

func TestGeoMatcher_MatchGeoIp_MissingCountry(t *testing.T) {
	loader := &fakeGeoLoader{
		geoip: map[string]*v2geo.GeoIP{},
	}
	g := NewGeoMatcher("", "")
	g.geoLoader = loader

	if g.MatchGeoIp("8.8.8.8", "us") {
		t.Error("MatchGeoIp for missing country should return false")
	}
}

func TestGeoMatcher_MatchGeoSite(t *testing.T) {
	loader := &fakeGeoLoader{
		geosite: map[string]*v2geo.GeoSite{
			"openai": {
				Domain: []*v2geo.Domain{
					{Type: v2geo.Domain_Plain, Value: "openai"},
					{Type: v2geo.Domain_Full, Value: "chatgpt.com"},
				},
			},
		},
	}
	g := NewGeoMatcher("", "")
	g.geoLoader = loader

	if !g.MatchGeoSite("api.openai.com", "openai") {
		t.Error("MatchGeoSite should match via plain domain")
	}
	if !g.MatchGeoSite("chatgpt.com", "openai") {
		t.Error("MatchGeoSite should match via full domain")
	}
	if g.MatchGeoSite("google.com", "openai") {
		t.Error("MatchGeoSite should not match unrelated host")
	}
}

func TestGeoMatcher_MatchGeoSite_EmptyCondition(t *testing.T) {
	g := NewGeoMatcher("", "")
	if g.MatchGeoSite("test.com", "") {
		t.Error("MatchGeoSite with empty condition should return false")
	}
}

func TestGeoMatcher_MatchGeoSite_MissingSite(t *testing.T) {
	loader := &fakeGeoLoader{
		geosite: map[string]*v2geo.GeoSite{},
	}
	g := NewGeoMatcher("", "")
	g.geoLoader = loader

	if g.MatchGeoSite("test.com", "nonexistent") {
		t.Error("MatchGeoSite for missing site should return false")
	}
}

func TestGeoMatcher_MatchGeoSiteSet(t *testing.T) {
	loader := &fakeGeoLoader{
		geosite: map[string]*v2geo.GeoSite{
			"openai": {
				Domain: []*v2geo.Domain{
					{Type: v2geo.Domain_Plain, Value: "openai"},
				},
			},
			"google": {
				Domain: []*v2geo.Domain{
					{Type: v2geo.Domain_RootDomain, Value: "google.com"},
				},
			},
		},
	}
	g := NewGeoMatcher("", "")
	g.geoLoader = loader

	set := &SiteConditionSet{Conditions: []string{" google ", "openai", "OPENAI"}}
	if !g.MatchGeoSiteSet("api.openai.com", set) {
		t.Error("MatchGeoSiteSet should match openai")
	}
	if !g.MatchGeoSiteSet("mail.google.com", set) {
		t.Error("MatchGeoSiteSet should match google")
	}
	if g.MatchGeoSiteSet("example.com", set) {
		t.Error("MatchGeoSiteSet should not match unrelated host")
	}
}

type countingMatcher struct {
	calls *atomic.Uint64
	match bool
}

func (m countingMatcher) Match(host HostInfo) bool {
	_ = host
	m.calls.Add(1)
	return m.match
}

func TestGeoMatcher_MatchGeoSite_UsesResultCache(t *testing.T) {
	g := NewGeoMatcher("", "")
	var calls atomic.Uint64
	g.geoSiteMatcher["openai"] = countingMatcher{calls: &calls, match: true}

	if !g.MatchGeoSite("api.openai.com", "openai") {
		t.Fatal("expected match")
	}
	if !g.MatchGeoSite("api.openai.com", "openai") {
		t.Fatal("expected cached match")
	}
	if got := calls.Load(); got != 1 {
		t.Fatalf("matcher calls=%d want=1", got)
	}
}

func TestGeoMatcher_MatchGeoSiteSet_UsesResultCache(t *testing.T) {
	g := NewGeoMatcher("", "")
	var calls atomic.Uint64
	g.geoSiteSets["openai\x1fyoutube"] = []hostMatcher{
		countingMatcher{calls: &calls, match: false},
		countingMatcher{calls: &calls, match: true},
	}

	set := &SiteConditionSet{Conditions: []string{"youtube", "openai"}}
	if !g.MatchGeoSiteSet("www.youtube.com", set) {
		t.Fatal("expected match")
	}
	if !g.MatchGeoSiteSet("www.youtube.com", set) {
		t.Fatal("expected cached match")
	}
	if got := calls.Load(); got != 2 {
		t.Fatalf("matcher calls=%d want=2", got)
	}
}

func ipv4(a, b, c, d byte) []byte {
	return []byte{a, b, c, d}
}