enhancer.go 2.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120
  1. package dns
  2. import (
  3. "net/netip"
  4. "github.com/metacubex/mihomo/common/lru"
  5. "github.com/metacubex/mihomo/component/fakeip"
  6. C "github.com/metacubex/mihomo/constant"
  7. )
  8. type ResolverEnhancer struct {
  9. mode C.DNSMode
  10. fakePool *fakeip.Pool
  11. mapping *lru.LruCache[netip.Addr, string]
  12. }
  13. func (h *ResolverEnhancer) FakeIPEnabled() bool {
  14. return h.mode == C.DNSFakeIP
  15. }
  16. func (h *ResolverEnhancer) MappingEnabled() bool {
  17. return h.mode == C.DNSFakeIP || h.mode == C.DNSMapping
  18. }
  19. func (h *ResolverEnhancer) IsExistFakeIP(ip netip.Addr) bool {
  20. if !h.FakeIPEnabled() {
  21. return false
  22. }
  23. if pool := h.fakePool; pool != nil {
  24. return pool.Exist(ip)
  25. }
  26. return false
  27. }
  28. func (h *ResolverEnhancer) IsFakeIP(ip netip.Addr) bool {
  29. if !h.FakeIPEnabled() {
  30. return false
  31. }
  32. if pool := h.fakePool; pool != nil {
  33. return pool.IPNet().Contains(ip) && ip != pool.Gateway() && ip != pool.Broadcast()
  34. }
  35. return false
  36. }
  37. func (h *ResolverEnhancer) IsFakeBroadcastIP(ip netip.Addr) bool {
  38. if !h.FakeIPEnabled() {
  39. return false
  40. }
  41. if pool := h.fakePool; pool != nil {
  42. return pool.Broadcast() == ip
  43. }
  44. return false
  45. }
  46. func (h *ResolverEnhancer) FindHostByIP(ip netip.Addr) (string, bool) {
  47. if pool := h.fakePool; pool != nil {
  48. if host, existed := pool.LookBack(ip); existed {
  49. return host, true
  50. }
  51. }
  52. if mapping := h.mapping; mapping != nil {
  53. if host, existed := h.mapping.Get(ip); existed {
  54. return host, true
  55. }
  56. }
  57. return "", false
  58. }
  59. func (h *ResolverEnhancer) InsertHostByIP(ip netip.Addr, host string) {
  60. if mapping := h.mapping; mapping != nil {
  61. h.mapping.Set(ip, host)
  62. }
  63. }
  64. func (h *ResolverEnhancer) FlushFakeIP() error {
  65. if h.fakePool != nil {
  66. return h.fakePool.FlushFakeIP()
  67. }
  68. return nil
  69. }
  70. func (h *ResolverEnhancer) PatchFrom(o *ResolverEnhancer) {
  71. if h.mapping != nil && o.mapping != nil {
  72. o.mapping.CloneTo(h.mapping)
  73. }
  74. if h.fakePool != nil && o.fakePool != nil {
  75. h.fakePool.CloneFrom(o.fakePool)
  76. }
  77. }
  78. func (h *ResolverEnhancer) StoreFakePoolState() {
  79. if h.fakePool != nil {
  80. h.fakePool.StoreState()
  81. }
  82. }
  83. func NewEnhancer(cfg Config) *ResolverEnhancer {
  84. var fakePool *fakeip.Pool
  85. var mapping *lru.LruCache[netip.Addr, string]
  86. if cfg.EnhancedMode != C.DNSNormal {
  87. fakePool = cfg.Pool
  88. mapping = lru.New(lru.WithSize[netip.Addr, string](4096))
  89. }
  90. return &ResolverEnhancer{
  91. mode: cfg.EnhancedMode,
  92. fakePool: fakePool,
  93. mapping: mapping,
  94. }
  95. }