domain.go 3.2 KB

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