domain_matcher.go 1.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798
  1. package strmatcher
  2. import "strings"
  3. func breakDomain(domain string) []string {
  4. return strings.Split(domain, ".")
  5. }
  6. type node struct {
  7. values []uint32
  8. sub map[string]*node
  9. }
  10. // DomainMatcherGroup is a IndexMatcher for a large set of Domain matchers.
  11. // Visible for testing only.
  12. type DomainMatcherGroup struct {
  13. root *node
  14. }
  15. func (g *DomainMatcherGroup) Add(domain string, value uint32) {
  16. if g.root == nil {
  17. g.root = new(node)
  18. }
  19. current := g.root
  20. parts := breakDomain(domain)
  21. for i := len(parts) - 1; i >= 0; i-- {
  22. part := parts[i]
  23. if current.sub == nil {
  24. current.sub = make(map[string]*node)
  25. }
  26. next := current.sub[part]
  27. if next == nil {
  28. next = new(node)
  29. current.sub[part] = next
  30. }
  31. current = next
  32. }
  33. current.values = append(current.values, value)
  34. }
  35. func (g *DomainMatcherGroup) addMatcher(m domainMatcher, value uint32) {
  36. g.Add(string(m), value)
  37. }
  38. func (g *DomainMatcherGroup) Match(domain string) []uint32 {
  39. if domain == "" {
  40. return nil
  41. }
  42. current := g.root
  43. if current == nil {
  44. return nil
  45. }
  46. nextPart := func(idx int) int {
  47. for i := idx - 1; i >= 0; i-- {
  48. if domain[i] == '.' {
  49. return i
  50. }
  51. }
  52. return -1
  53. }
  54. matches := [][]uint32{}
  55. idx := len(domain)
  56. for {
  57. if idx == -1 || current.sub == nil {
  58. break
  59. }
  60. nidx := nextPart(idx)
  61. part := domain[nidx+1 : idx]
  62. next := current.sub[part]
  63. if next == nil {
  64. break
  65. }
  66. current = next
  67. idx = nidx
  68. if len(current.values) > 0 {
  69. matches = append(matches, current.values)
  70. }
  71. }
  72. switch len(matches) {
  73. case 0:
  74. return nil
  75. case 1:
  76. return matches[0]
  77. default:
  78. result := []uint32{}
  79. for idx := range matches {
  80. // Insert reversely, the subdomain that matches further ranks higher
  81. result = append(result, matches[len(matches)-1-idx]...)
  82. }
  83. return result
  84. }
  85. }