github.com/v2fly/v2ray-core/v5@v5.16.2-0.20240507031116-8191faa6e095/common/strmatcher/matchergroup_domain.go (about)

     1  package strmatcher
     2  
     3  type trieNode struct {
     4  	values   []uint32
     5  	children map[string]*trieNode
     6  }
     7  
     8  // DomainMatcherGroup is an implementation of MatcherGroup.
     9  // It uses trie to optimize both memory consumption and lookup speed. Trie node is domain label based.
    10  type DomainMatcherGroup struct {
    11  	root *trieNode
    12  }
    13  
    14  func NewDomainMatcherGroup() *DomainMatcherGroup {
    15  	return &DomainMatcherGroup{
    16  		root: new(trieNode),
    17  	}
    18  }
    19  
    20  // AddDomainMatcher implements MatcherGroupForDomain.AddDomainMatcher.
    21  func (g *DomainMatcherGroup) AddDomainMatcher(matcher DomainMatcher, value uint32) {
    22  	node := g.root
    23  	pattern := matcher.Pattern()
    24  	for i := len(pattern); i > 0; {
    25  		var part string
    26  		for j := i - 1; ; j-- {
    27  			if pattern[j] == '.' {
    28  				part = pattern[j+1 : i]
    29  				i = j
    30  				break
    31  			}
    32  			if j == 0 {
    33  				part = pattern[j:i]
    34  				i = j
    35  				break
    36  			}
    37  		}
    38  		if node.children == nil {
    39  			node.children = make(map[string]*trieNode)
    40  		}
    41  		next := node.children[part]
    42  		if next == nil {
    43  			next = new(trieNode)
    44  			node.children[part] = next
    45  		}
    46  		node = next
    47  	}
    48  
    49  	node.values = append(node.values, value)
    50  }
    51  
    52  // Match implements MatcherGroup.Match.
    53  func (g *DomainMatcherGroup) Match(input string) []uint32 {
    54  	matches := make([][]uint32, 0, 5)
    55  	node := g.root
    56  	for i := len(input); i > 0; {
    57  		for j := i - 1; ; j-- {
    58  			if input[j] == '.' { // Domain label found
    59  				node = node.children[input[j+1:i]]
    60  				i = j
    61  				break
    62  			}
    63  			if j == 0 { // The last part of domain label
    64  				node = node.children[input[j:i]]
    65  				i = j
    66  				break
    67  			}
    68  		}
    69  		if node == nil { // No more match if no trie edge transition
    70  			break
    71  		}
    72  		if len(node.values) > 0 { // Found matched matchers
    73  			matches = append(matches, node.values)
    74  		}
    75  		if node.children == nil { // No more match if leaf node reached
    76  			break
    77  		}
    78  	}
    79  	return CompositeMatchesReverse(matches)
    80  }
    81  
    82  // MatchAny implements MatcherGroup.MatchAny.
    83  func (g *DomainMatcherGroup) MatchAny(input string) bool {
    84  	node := g.root
    85  	for i := len(input); i > 0; {
    86  		for j := i - 1; ; j-- {
    87  			if input[j] == '.' {
    88  				node = node.children[input[j+1:i]]
    89  				i = j
    90  				break
    91  			}
    92  			if j == 0 {
    93  				node = node.children[input[j:i]]
    94  				i = j
    95  				break
    96  			}
    97  		}
    98  		if node == nil {
    99  			return false
   100  		}
   101  		if len(node.values) > 0 {
   102  			return true
   103  		}
   104  		if node.children == nil {
   105  			return false
   106  		}
   107  	}
   108  	return false
   109  }