github.com/igoogolx/clash@v1.19.8/component/trie/domain.go (about)

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