github.com/database64128/shadowsocks-go@v1.10.2-0.20240315062903-143a773533f1/domainset/matcher_suffix_trie.go (about)

     1  package domainset
     2  
     3  // DomainSuffixTrie is a trie of domain parts segmented by '.'.
     4  type DomainSuffixTrie struct {
     5  	Included bool
     6  	Children map[string]*DomainSuffixTrie
     7  }
     8  
     9  // Insert inserts a domain suffix to the trie.
    10  // Insertion purges the leaf node's children.
    11  // If say, we insert "www.google.com" and then "google.com",
    12  // The children of node "google" will be purged.
    13  func (dst *DomainSuffixTrie) Insert(domain string) {
    14  	cdst := dst
    15  
    16  	for i := len(domain) - 1; i >= 0; i-- {
    17  		if domain[i] != '.' {
    18  			continue
    19  		}
    20  
    21  		part := domain[i+1:]
    22  		domain = domain[:i]
    23  		if cdst.Children == nil {
    24  			var ndst DomainSuffixTrie
    25  			cdst.Children = map[string]*DomainSuffixTrie{
    26  				part: &ndst,
    27  			}
    28  			cdst = &ndst
    29  		} else {
    30  			ndst, ok := cdst.Children[part]
    31  			switch {
    32  			case !ok:
    33  				ndst = &DomainSuffixTrie{}
    34  				cdst.Children[part] = ndst
    35  				cdst = ndst
    36  			case ndst.Included:
    37  				return
    38  			default:
    39  				cdst = ndst
    40  			}
    41  		}
    42  	}
    43  
    44  	if cdst.Children == nil {
    45  		cdst.Children = map[string]*DomainSuffixTrie{
    46  			domain: {
    47  				Included: true,
    48  			},
    49  		}
    50  	} else {
    51  		ndst, ok := cdst.Children[domain]
    52  		if !ok {
    53  			cdst.Children[domain] = &DomainSuffixTrie{
    54  				Included: true,
    55  			}
    56  		} else {
    57  			ndst.Included = true
    58  			ndst.Children = nil
    59  		}
    60  	}
    61  }
    62  
    63  // Match implements the Matcher Match method.
    64  func (dst *DomainSuffixTrie) Match(domain string) bool {
    65  	cdst := dst
    66  
    67  	for i := len(domain) - 1; i >= 0; i-- {
    68  		if domain[i] != '.' {
    69  			continue
    70  		}
    71  
    72  		if cdst.Children == nil {
    73  			return false
    74  		}
    75  
    76  		ndst, ok := cdst.Children[domain[i+1:]]
    77  		if !ok {
    78  			return false
    79  		}
    80  		if ndst.Included {
    81  			return true
    82  		}
    83  		cdst = ndst
    84  		domain = domain[:i]
    85  	}
    86  
    87  	ndst, ok := cdst.Children[domain]
    88  	if !ok {
    89  		return false
    90  	}
    91  	return ndst.Included
    92  }
    93  
    94  // Keys returns the keys of the trie.
    95  func (dst *DomainSuffixTrie) Keys() (keys []string) {
    96  	for s, c := range dst.Children {
    97  		keys = c.keys(s, keys)
    98  	}
    99  	return
   100  }
   101  
   102  func (dst *DomainSuffixTrie) keys(suffix string, keys []string) []string {
   103  	if dst.Included {
   104  		keys = append(keys, suffix)
   105  	}
   106  	for s, c := range dst.Children {
   107  		keys = c.keys(s+"."+suffix, keys)
   108  	}
   109  	return keys
   110  }
   111  
   112  // Rules implements the MatcherBuilder Rules method.
   113  func (dst *DomainSuffixTrie) Rules() []string {
   114  	return dst.Keys()
   115  }
   116  
   117  // MatcherCount implements the MatcherBuilder MatcherCount method.
   118  func (dst *DomainSuffixTrie) MatcherCount() int {
   119  	if dst.Children == nil {
   120  		return 0
   121  	}
   122  	return 1
   123  }
   124  
   125  // AppendTo implements the MatcherBuilder AppendTo method.
   126  func (dst *DomainSuffixTrie) AppendTo(matchers []Matcher) ([]Matcher, error) {
   127  	if dst.Children == nil {
   128  		return matchers, nil
   129  	}
   130  	return append(matchers, dst), nil
   131  }
   132  
   133  func NewDomainSuffixTrie(capacity int) MatcherBuilder {
   134  	return &DomainSuffixTrie{}
   135  }
   136  
   137  func DomainSuffixTrieFromSlice(suffixes []string) *DomainSuffixTrie {
   138  	var dst DomainSuffixTrie
   139  	for _, s := range suffixes {
   140  		dst.Insert(s)
   141  	}
   142  	return &dst
   143  }