history-museum/Clash.Meta
Archived
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Refactor: DomainTrie use generics

Browse Source
This commit is contained in:
yaling888
2022-04-06 04:25:53 +08:00
committed by Meta
parent 0c65f6962a
commit a6eb11ce18
10 changed files with 100 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
component/fakeip/pool.go
View File

@ -28,7 +28,7 @@ type Pool struct {
broadcast uint32
offset uint32
mux sync.Mutex
host *trie.DomainTrie
host *trie.DomainTrie[bool]
ipnet *net.IPNet
store store
}
@ -138,7 +138,7 @@ func uintToIP(v uint32) net.IP {
type Options struct {
IPNet *net.IPNet
Host *trie.DomainTrie
Host *trie.DomainTrie[bool]
// Size sets the maximum number of entries in memory
// and does not work if Persistence is true

4
component/fakeip/pool_test.go
View File

@ -100,8 +100,8 @@ func TestPool_CycleUsed(t *testing.T) {
func TestPool_Skip(t *testing.T) {
_, ipnet, _ := net.ParseCIDR("192.168.0.1/29")
tree := trie.New()
tree.Insert("example.com", tree)
tree := trie.New[bool]()
tree.Insert("example.com", true)
pools, tempfile, err := createPools(Options{
IPNet: ipnet,
Size: 10,

14
component/resolver/resolver.go
View File

@ -5,6 +5,7 @@ import (
"errors"
"math/rand"
"net"
"net/netip"
"strings"
"time"
@ -23,7 +24,7 @@ var (
DisableIPv6 = true
// DefaultHosts aim to resolve hosts
DefaultHosts = trie.New()
DefaultHosts = trie.New[netip.Addr]()
// DefaultDNSTimeout defined the default dns request timeout
DefaultDNSTimeout = time.Second * 5
@ -48,8 +49,8 @@ func ResolveIPv4(host string) (net.IP, error) {
func ResolveIPv4WithResolver(host string, r Resolver) (net.IP, error) {
if node := DefaultHosts.Search(host); node != nil {
if ip := node.Data.(net.IP).To4(); ip != nil {
return ip, nil
if ip := node.Data; ip.Is4() {
return ip.AsSlice(), nil
}
}
@ -92,8 +93,8 @@ func ResolveIPv6WithResolver(host string, r Resolver) (net.IP, error) {
}
if node := DefaultHosts.Search(host); node != nil {
if ip := node.Data.(net.IP).To16(); ip != nil {
return ip, nil
if ip := node.Data; ip.Is6() {
return ip.AsSlice(), nil
}
}
@ -128,7 +129,8 @@ func ResolveIPv6WithResolver(host string, r Resolver) (net.IP, error) {
// ResolveIPWithResolver same as ResolveIP, but with a resolver
func ResolveIPWithResolver(host string, r Resolver) (net.IP, error) {
if node := DefaultHosts.Search(host); node != nil {
return node.Data.(net.IP), nil
ip := node.Data
return ip.Unmap().AsSlice(), nil
}
if r != nil {

24
component/trie/domain.go
View File

@ -17,8 +17,8 @@ var ErrInvalidDomain = errors.New("invalid domain")
// DomainTrie contains the main logic for adding and searching nodes for domain segments.
// support wildcard domain (e.g *.google.com)
type DomainTrie struct {
root *Node
type DomainTrie[T comparable] struct {
root *Node[T]
}
func ValidAndSplitDomain(domain string) ([]string, bool) {
@ -51,7 +51,7 @@ func ValidAndSplitDomain(domain string) ([]string, bool) {
// 3. subdomain.*.example.com
// 4. .example.com
// 5. +.example.com
func (t *DomainTrie) Insert(domain string, data any) error {
func (t *DomainTrie[T]) Insert(domain string, data T) error {
parts, valid := ValidAndSplitDomain(domain)
if !valid {
return ErrInvalidDomain
@ -68,13 +68,13 @@ func (t *DomainTrie) Insert(domain string, data any) error {
return nil
}
func (t *DomainTrie) insert(parts []string, data any) {
func (t *DomainTrie[T]) insert(parts []string, data T) {
node := t.root
// reverse storage domain part to save space
for i := len(parts) - 1; i >= 0; i-- {
part := parts[i]
if !node.hasChild(part) {
node.addChild(part, newNode(nil))
node.addChild(part, newNode(getZero[T]()))
}
node = node.getChild(part)
@ -88,7 +88,7 @@ func (t *DomainTrie) insert(parts []string, data any) {
// 1. static part
// 2. wildcard domain
// 2. dot wildcard domain
func (t *DomainTrie) Search(domain string) *Node {
func (t *DomainTrie[T]) Search(domain string) *Node[T] {
parts, valid := ValidAndSplitDomain(domain)
if !valid || parts[0] == "" {
return nil
@ -96,26 +96,26 @@ func (t *DomainTrie) Search(domain string) *Node {
n := t.search(t.root, parts)
if n == nil || n.Data == nil {
if n == nil || n.Data == getZero[T]() {
return nil
}
return n
}
func (t *DomainTrie) search(node *Node, parts []string) *Node {
func (t *DomainTrie[T]) search(node *Node[T], parts []string) *Node[T] {
if len(parts) == 0 {
return node
}
if c := node.getChild(parts[len(parts)-1]); c != nil {
if n := t.search(c, parts[:len(parts)-1]); n != nil && n.Data != nil {
if n := t.search(c, parts[:len(parts)-1]); n != nil && n.Data != getZero[T]() {
return n
}
}
if c := node.getChild(wildcard); c != nil {
if n := t.search(c, parts[:len(parts)-1]); n != nil && n.Data != nil {
if n := t.search(c, parts[:len(parts)-1]); n != nil && n.Data != getZero[T]() {
return n
}
}
@ -124,6 +124,6 @@ func (t *DomainTrie) search(node *Node, parts []string) *Node {
}
// New returns a new, empty Trie.
func New() *DomainTrie {
return &DomainTrie{root: newNode(nil)}
func New[T comparable]() *DomainTrie[T] {
return &DomainTrie[T]{root: newNode[T](getZero[T]())}
}

28
component/trie/domain_test.go
View File

@ -1,16 +1,16 @@
package trie
import (
"net"
"net/netip"
"testing"
"github.com/stretchr/testify/assert"
)
var localIP = net.IP{127, 0, 0, 1}
var localIP = netip.AddrFrom4([4]byte{127, 0, 0, 1})
func TestTrie_Basic(t *testing.T) {
tree := New()
tree := New[netip.Addr]()
domains := []string{
"example.com",
"google.com",
@ -23,7 +23,7 @@ func TestTrie_Basic(t *testing.T) {
node := tree.Search("example.com")
assert.NotNil(t, node)
assert.True(t, node.Data.(net.IP).Equal(localIP))
assert.True(t, node.Data == localIP)
assert.NotNil(t, tree.Insert("", localIP))
assert.Nil(t, tree.Search(""))
assert.NotNil(t, tree.Search("localhost"))
@ -31,7 +31,7 @@ func TestTrie_Basic(t *testing.T) {
}
func TestTrie_Wildcard(t *testing.T) {
tree := New()
tree := New[netip.Addr]()
domains := []string{
"*.example.com",
"sub.*.example.com",
@ -64,7 +64,7 @@ func TestTrie_Wildcard(t *testing.T) {
}
func TestTrie_Priority(t *testing.T) {
tree := New()
tree := New[int]()
domains := []string{
".dev",
"example.dev",
@ -79,18 +79,18 @@ func TestTrie_Priority(t *testing.T) {
}
for idx, domain := range domains {
tree.Insert(domain, idx)
tree.Insert(domain, idx+1)
}
assertFn("test.dev", 0)
assertFn("foo.bar.dev", 0)
assertFn("example.dev", 1)
assertFn("foo.example.dev", 2)
assertFn("test.example.dev", 3)
assertFn("test.dev", 1)
assertFn("foo.bar.dev", 1)
assertFn("example.dev", 2)
assertFn("foo.example.dev", 3)
assertFn("test.example.dev", 4)
}
func TestTrie_Boundary(t *testing.T) {
tree := New()
tree := New[netip.Addr]()
tree.Insert("*.dev", localIP)
assert.NotNil(t, tree.Insert(".", localIP))
@ -99,7 +99,7 @@ func TestTrie_Boundary(t *testing.T) {
}
func TestTrie_WildcardBoundary(t *testing.T) {
tree := New()
tree := New[netip.Addr]()
tree.Insert("+.*", localIP)
tree.Insert("stun.*.*.*", localIP)

31
component/trie/node.go
View File

@ -1,34 +1,31 @@
package trie
// Node is the trie's node
type Node struct {
children map[string]*Node
Data any
type Node[T comparable] struct {
children map[string]*Node[T]
Data T
}
func (n *Node) getChild(s string) *Node {
if n.children == nil {
return nil
}
func (n *Node[T]) getChild(s string) *Node[T] {
return n.children[s]
}
func (n *Node) hasChild(s string) bool {
func (n *Node[T]) hasChild(s string) bool {
return n.getChild(s) != nil
}
func (n *Node) addChild(s string, child *Node) {
if n.children == nil {
n.children = map[string]*Node{}
}
func (n *Node[T]) addChild(s string, child *Node[T]) {
n.children[s] = child
}
func newNode(data any) *Node {
return &Node{
func newNode[T comparable](data T) *Node[T] {
return &Node[T]{
Data: data,
children: nil,
children: map[string]*Node[T]{},
}
}
func getZero[T comparable]() T {
var result T
return result
}