github.com/metacubex/mihomo@v1.18.5/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[T any] struct {
    21  	root *Node[T]
    22  }
    23  
    24  func ValidAndSplitDomain(domain string) ([]string, bool) {
    25  	if domain != "" && domain[len(domain)-1] == '.' {
    26  		return nil, false
    27  	}
    28  	domain = strings.ToLower(domain)
    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[T]) Insert(domain string, data T) 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[T]) insert(parts []string, data T) {
    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  		node = node.getOrNewChild(part)
    77  	}
    78  
    79  	node.setData(data)
    80  }
    81  
    82  // Search is the most important part of the Trie.
    83  // Priority as:
    84  // 1. static part
    85  // 2. wildcard domain
    86  // 2. dot wildcard domain
    87  func (t *DomainTrie[T]) Search(domain string) *Node[T] {
    88  	parts, valid := ValidAndSplitDomain(domain)
    89  	if !valid || parts[0] == "" {
    90  		return nil
    91  	}
    92  
    93  	n := t.search(t.root, parts)
    94  
    95  	if n.isEmpty() {
    96  		return nil
    97  	}
    98  
    99  	return n
   100  }
   101  
   102  func (t *DomainTrie[T]) search(node *Node[T], parts []string) *Node[T] {
   103  	if len(parts) == 0 {
   104  		return node
   105  	}
   106  
   107  	if c := node.getChild(parts[len(parts)-1]); c != nil {
   108  		if n := t.search(c, parts[:len(parts)-1]); !n.isEmpty() {
   109  			return n
   110  		}
   111  	}
   112  
   113  	if c := node.getChild(wildcard); c != nil {
   114  		if n := t.search(c, parts[:len(parts)-1]); !n.isEmpty() {
   115  			return n
   116  		}
   117  	}
   118  
   119  	return node.getChild(dotWildcard)
   120  }
   121  
   122  func (t *DomainTrie[T]) Optimize() {
   123  	t.root.optimize()
   124  }
   125  
   126  func (t *DomainTrie[T]) Foreach(print func(domain string, data T)) {
   127  	for key, data := range t.root.getChildren() {
   128  		recursion([]string{key}, data, print)
   129  		if data != nil && data.inited {
   130  			print(joinDomain([]string{key}), data.data)
   131  		}
   132  	}
   133  }
   134  
   135  func recursion[T any](items []string, node *Node[T], fn func(domain string, data T)) {
   136  	for key, data := range node.getChildren() {
   137  		newItems := append([]string{key}, items...)
   138  		if data != nil && data.inited {
   139  			domain := joinDomain(newItems)
   140  			if domain[0] == domainStepByte {
   141  				domain = complexWildcard + domain
   142  			}
   143  			fn(domain, data.Data())
   144  		}
   145  		recursion(newItems, data, fn)
   146  	}
   147  }
   148  
   149  func joinDomain(items []string) string {
   150  	return strings.Join(items, domainStep)
   151  }
   152  
   153  // New returns a new, empty Trie.
   154  func New[T any]() *DomainTrie[T] {
   155  	return &DomainTrie[T]{root: newNode[T]()}
   156  }