github.com/yaling888/clash@v1.53.0/component/trie/domain.go (about)

     1  package trie
     2  
     3  import (
     4  	"errors"
     5  	"strings"
     6  
     7  	"github.com/samber/lo"
     8  )
     9  
    10  const (
    11  	wildcard        = "*"
    12  	dotWildcard     = ""
    13  	complexWildcard = "+"
    14  	domainStep      = "."
    15  )
    16  
    17  // ErrInvalidDomain means insert domain is invalid
    18  var ErrInvalidDomain = errors.New("invalid domain")
    19  
    20  // DomainTrie contains the main logic for adding and searching nodes for domain segments.
    21  // support wildcard domain (e.g *.google.com)
    22  type DomainTrie[T comparable] struct {
    23  	root *Node[T]
    24  }
    25  
    26  func ValidAndSplitDomain(domain string) ([]string, bool) {
    27  	if domain != "" && domain[len(domain)-1] == '.' {
    28  		return nil, false
    29  	}
    30  
    31  	parts := strings.Split(domain, domainStep)
    32  	if len(parts) == 1 {
    33  		if parts[0] == "" {
    34  			return nil, false
    35  		}
    36  
    37  		return parts, true
    38  	}
    39  
    40  	for _, part := range parts[1:] {
    41  		if part == "" {
    42  			return nil, false
    43  		}
    44  	}
    45  
    46  	return parts, true
    47  }
    48  
    49  // Insert adds a node to the trie.
    50  // Support
    51  // 1. www.example.com
    52  // 2. *.example.com
    53  // 3. subdomain.*.example.com
    54  // 4. .example.com
    55  // 5. +.example.com
    56  func (t *DomainTrie[T]) Insert(domain string, data T) error {
    57  	parts, valid := ValidAndSplitDomain(domain)
    58  	if !valid {
    59  		return ErrInvalidDomain
    60  	}
    61  
    62  	if parts[0] == complexWildcard {
    63  		t.insert(parts[1:], data)
    64  		parts[0] = dotWildcard
    65  		t.insert(parts, data)
    66  	} else {
    67  		t.insert(parts, data)
    68  	}
    69  
    70  	return nil
    71  }
    72  
    73  func (t *DomainTrie[T]) insert(parts []string, data T) {
    74  	node := t.root
    75  	// reverse storage domain part to save space
    76  	for i := len(parts) - 1; i >= 0; i-- {
    77  		part := parts[i]
    78  		if !node.hasChild(part) {
    79  			node.addChild(part, newNode(lo.Empty[T]()))
    80  		}
    81  
    82  		node = node.getChild(part)
    83  	}
    84  
    85  	node.Data = data
    86  }
    87  
    88  // Search is the most important part of the Trie.
    89  // Priority as:
    90  // 1. static part
    91  // 2. wildcard domain
    92  // 2. dot wildcard domain
    93  func (t *DomainTrie[T]) Search(domain string) *Node[T] {
    94  	parts, valid := ValidAndSplitDomain(domain)
    95  	if !valid || parts[0] == "" {
    96  		return nil
    97  	}
    98  
    99  	n := t.search(t.root, parts)
   100  
   101  	if n == nil || n.Data == lo.Empty[T]() {
   102  		return nil
   103  	}
   104  
   105  	return n
   106  }
   107  
   108  func (t *DomainTrie[T]) search(node *Node[T], parts []string) *Node[T] {
   109  	if len(parts) == 0 {
   110  		return node
   111  	}
   112  
   113  	if c := node.getChild(parts[len(parts)-1]); c != nil {
   114  		if n := t.search(c, parts[:len(parts)-1]); n != nil && n.Data != lo.Empty[T]() {
   115  			return n
   116  		}
   117  	}
   118  
   119  	if c := node.getChild(wildcard); c != nil {
   120  		if n := t.search(c, parts[:len(parts)-1]); n != nil && n.Data != lo.Empty[T]() {
   121  			return n
   122  		}
   123  	}
   124  
   125  	return node.getChild(dotWildcard)
   126  }
   127  
   128  // New returns a new, empty Trie.
   129  func New[T comparable]() *DomainTrie[T] {
   130  	return &DomainTrie[T]{root: newNode[T](lo.Empty[T]())}
   131  }