domain_strategy.go 2.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107
  1. package provider
  2. import (
  3. "errors"
  4. "io"
  5. "strings"
  6. "github.com/metacubex/mihomo/component/trie"
  7. C "github.com/metacubex/mihomo/constant"
  8. P "github.com/metacubex/mihomo/constant/provider"
  9. "github.com/metacubex/mihomo/log"
  10. "golang.org/x/exp/slices"
  11. )
  12. type domainStrategy struct {
  13. count int
  14. domainTrie *trie.DomainTrie[struct{}]
  15. domainSet *trie.DomainSet
  16. }
  17. func (d *domainStrategy) Behavior() P.RuleBehavior {
  18. return P.Domain
  19. }
  20. func (d *domainStrategy) ShouldFindProcess() bool {
  21. return false
  22. }
  23. func (d *domainStrategy) Match(metadata *C.Metadata) bool {
  24. return d.domainSet != nil && d.domainSet.Has(metadata.RuleHost())
  25. }
  26. func (d *domainStrategy) Count() int {
  27. return d.count
  28. }
  29. func (d *domainStrategy) ShouldResolveIP() bool {
  30. return false
  31. }
  32. func (d *domainStrategy) Reset() {
  33. d.domainTrie = trie.New[struct{}]()
  34. d.domainSet = nil
  35. d.count = 0
  36. }
  37. func (d *domainStrategy) Insert(rule string) {
  38. if strings.ContainsRune(rule, '/') {
  39. log.Warnln("invalid domain:[%s]", rule)
  40. return
  41. }
  42. err := d.domainTrie.Insert(rule, struct{}{})
  43. if err != nil {
  44. log.Warnln("invalid domain:[%s]", rule)
  45. } else {
  46. d.count++
  47. }
  48. }
  49. func (d *domainStrategy) FinishInsert() {
  50. d.domainSet = d.domainTrie.NewDomainSet()
  51. d.domainTrie = nil
  52. }
  53. func (d *domainStrategy) FromMrs(r io.Reader, count int) error {
  54. domainSet, err := trie.ReadDomainSetBin(r)
  55. if err != nil {
  56. return err
  57. }
  58. d.count = count
  59. d.domainSet = domainSet
  60. return nil
  61. }
  62. func (d *domainStrategy) WriteMrs(w io.Writer) error {
  63. if d.domainSet == nil {
  64. return errors.New("nil domainSet")
  65. }
  66. return d.domainSet.WriteBin(w)
  67. }
  68. func (d *domainStrategy) DumpMrs(f func(key string) bool) {
  69. if d.domainSet != nil {
  70. var keys []string
  71. d.domainSet.Foreach(func(key string) bool {
  72. keys = append(keys, key)
  73. return true
  74. })
  75. slices.Sort(keys)
  76. for _, key := range keys {
  77. if _, ok := slices.BinarySearch(keys, "+."+key); ok {
  78. continue // ignore the rules added by trie internal processing
  79. }
  80. if !f(key) {
  81. return
  82. }
  83. }
  84. }
  85. }
  86. var _ mrsRuleStrategy = (*domainStrategy)(nil)
  87. func NewDomainStrategy() *domainStrategy {
  88. return &domainStrategy{}
  89. }