From c37f988a4f81e3f48fb938702e629e2c4179ef3e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Tue, 2 Jul 2024 15:49:37 +0800
Subject: [PATCH] Add dump for domain matcher

---
 common/domain/matcher.go      | 29 ++++++++++++++++++
 common/domain/matcher_test.go | 58 +++++++++++++++++++++++++++++++++++
 common/domain/set.go          | 28 +++++++++++++++++
 3 files changed, 115 insertions(+)
 create mode 100644 common/domain/matcher_test.go

diff --git a/common/domain/matcher.go b/common/domain/matcher.go
index 108a316..1f328b8 100644
--- a/common/domain/matcher.go
+++ b/common/domain/matcher.go
@@ -72,6 +72,35 @@ func (m *Matcher) Write(writer varbin.Writer) error {
 	})
 }
 
+func (m *Matcher) Dump() (domainList []string, prefixList []string) {
+	domainMap := make(map[string]bool)
+	prefixMap := make(map[string]bool)
+	for _, key := range m.set.keys() {
+		key = reverseDomain(key)
+		if key[0] == prefixLabel {
+			prefixMap[key[1:]] = true
+		} else {
+			domainMap[key] = true
+		}
+	}
+	for rawPrefix := range prefixMap {
+		if rawPrefix[0] == '.' {
+			if rootDomain := rawPrefix[1:]; domainMap[rootDomain] {
+				delete(domainMap, rootDomain)
+				prefixList = append(prefixList, rootDomain)
+				continue
+			}
+		}
+		prefixList = append(prefixList, rawPrefix)
+	}
+	for domain := range domainMap {
+		domainList = append(domainList, domain)
+	}
+	sort.Strings(domainList)
+	sort.Strings(prefixList)
+	return domainList, prefixList
+}
+
 func reverseDomain(domain string) string {
 	l := len(domain)
 	b := make([]byte, l)
diff --git a/common/domain/matcher_test.go b/common/domain/matcher_test.go
new file mode 100644
index 0000000..f7b4cc1
--- /dev/null
+++ b/common/domain/matcher_test.go
@@ -0,0 +1,58 @@
+package domain_test
+
+import (
+	"encoding/json"
+	"net/http"
+	"sort"
+	"testing"
+
+	"github.com/sagernet/sing/common/domain"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestMatcher(t *testing.T) {
+	testDomain := []string{"example.com", "example.org"}
+	testDomainSuffix := []string{".com.cn", ".org.cn", "sagernet.org"}
+	matcher := domain.NewMatcher(testDomain, testDomainSuffix)
+	require.NotNil(t, matcher)
+	require.True(t, matcher.Match("example.com"))
+	require.True(t, matcher.Match("example.org"))
+	require.False(t, matcher.Match("example.cn"))
+	require.True(t, matcher.Match("example.com.cn"))
+	require.True(t, matcher.Match("example.org.cn"))
+	require.False(t, matcher.Match("com.cn"))
+	require.False(t, matcher.Match("org.cn"))
+	require.True(t, matcher.Match("sagernet.org"))
+	require.True(t, matcher.Match("sing-box.sagernet.org"))
+	dDomain, dDomainSuffix := matcher.Dump()
+	require.Equal(t, testDomain, dDomain)
+	require.Equal(t, testDomainSuffix, dDomainSuffix)
+}
+
+type simpleRuleSet struct {
+	Rules []struct {
+		Domain       []string `json:"domain"`
+		DomainSuffix []string `json:"domain_suffix"`
+	}
+}
+
+func TestDumpLarge(t *testing.T) {
+	response, err := http.Get("https://raw.githubusercontent.com/MetaCubeX/meta-rules-dat/sing/geo/geosite/cn.json")
+	require.NoError(t, err)
+	defer response.Body.Close()
+	var ruleSet simpleRuleSet
+	err = json.NewDecoder(response.Body).Decode(&ruleSet)
+	require.NoError(t, err)
+	domainList := ruleSet.Rules[0].Domain
+	domainSuffixList := ruleSet.Rules[0].DomainSuffix
+	require.Len(t, ruleSet.Rules, 1)
+	require.True(t, len(domainList)+len(domainSuffixList) > 0)
+	sort.Strings(domainList)
+	sort.Strings(domainSuffixList)
+	matcher := domain.NewMatcher(domainList, domainSuffixList)
+	require.NotNil(t, matcher)
+	dDomain, dDomainSuffix := matcher.Dump()
+	require.Equal(t, domainList, dDomain)
+	require.Equal(t, domainSuffixList, dDomainSuffix)
+}
diff --git a/common/domain/set.go b/common/domain/set.go
index adf661a..ae1e00b 100644
--- a/common/domain/set.go
+++ b/common/domain/set.go
@@ -74,6 +74,34 @@ func (ss *succinctSet) Has(key string) bool {
 	}
 }
 
+func (ss *succinctSet) keys() []string {
+	var result []string
+	var currentKey []byte
+	var bmIdx, nodeId int
+
+	var traverse func(int, int)
+	traverse = func(nodeId, bmIdx int) {
+		if getBit(ss.leaves, nodeId) != 0 {
+			result = append(result, string(currentKey))
+		}
+
+		for ; ; bmIdx++ {
+			if getBit(ss.labelBitmap, bmIdx) != 0 {
+				return
+			}
+			nextLabel := ss.labels[bmIdx-nodeId]
+			currentKey = append(currentKey, nextLabel)
+			nextNodeId := countZeros(ss.labelBitmap, ss.ranks, bmIdx+1)
+			nextBmIdx := selectIthOne(ss.labelBitmap, ss.ranks, ss.selects, nextNodeId-1) + 1
+			traverse(nextNodeId, nextBmIdx)
+			currentKey = currentKey[:len(currentKey)-1]
+		}
+	}
+
+	traverse(nodeId, bmIdx)
+	return result
+}
+
 func setBit(bm *[]uint64, i int, v int) {
 	for i>>6 >= len(*bm) {
 		*bm = append(*bm, 0)