pool.go 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201
  1. package fakeip
  2. import (
  3. "errors"
  4. "net/netip"
  5. "strings"
  6. "sync"
  7. "github.com/metacubex/mihomo/common/nnip"
  8. "github.com/metacubex/mihomo/component/profile/cachefile"
  9. "github.com/metacubex/mihomo/component/trie"
  10. )
  11. const (
  12. offsetKey = "key-offset-fake-ip"
  13. cycleKey = "key-cycle-fake-ip"
  14. )
  15. type store interface {
  16. GetByHost(host string) (netip.Addr, bool)
  17. PutByHost(host string, ip netip.Addr)
  18. GetByIP(ip netip.Addr) (string, bool)
  19. PutByIP(ip netip.Addr, host string)
  20. DelByIP(ip netip.Addr)
  21. Exist(ip netip.Addr) bool
  22. CloneTo(store)
  23. FlushFakeIP() error
  24. }
  25. // Pool is an implementation about fake ip generator without storage
  26. type Pool struct {
  27. gateway netip.Addr
  28. first netip.Addr
  29. last netip.Addr
  30. offset netip.Addr
  31. cycle bool
  32. mux sync.Mutex
  33. host *trie.DomainTrie[struct{}]
  34. ipnet netip.Prefix
  35. store store
  36. }
  37. // Lookup return a fake ip with host
  38. func (p *Pool) Lookup(host string) netip.Addr {
  39. p.mux.Lock()
  40. defer p.mux.Unlock()
  41. // RFC4343: DNS Case Insensitive, we SHOULD return result with all cases.
  42. host = strings.ToLower(host)
  43. if ip, exist := p.store.GetByHost(host); exist {
  44. return ip
  45. }
  46. ip := p.get(host)
  47. p.store.PutByHost(host, ip)
  48. return ip
  49. }
  50. // LookBack return host with the fake ip
  51. func (p *Pool) LookBack(ip netip.Addr) (string, bool) {
  52. p.mux.Lock()
  53. defer p.mux.Unlock()
  54. return p.store.GetByIP(ip)
  55. }
  56. // ShouldSkipped return if domain should be skipped
  57. func (p *Pool) ShouldSkipped(domain string) bool {
  58. if p.host == nil {
  59. return false
  60. }
  61. return p.host.Search(domain) != nil
  62. }
  63. // Exist returns if given ip exists in fake-ip pool
  64. func (p *Pool) Exist(ip netip.Addr) bool {
  65. p.mux.Lock()
  66. defer p.mux.Unlock()
  67. return p.store.Exist(ip)
  68. }
  69. // Gateway return gateway ip
  70. func (p *Pool) Gateway() netip.Addr {
  71. return p.gateway
  72. }
  73. // Broadcast return the last ip
  74. func (p *Pool) Broadcast() netip.Addr {
  75. return p.last
  76. }
  77. // IPNet return raw ipnet
  78. func (p *Pool) IPNet() netip.Prefix {
  79. return p.ipnet
  80. }
  81. // CloneFrom clone cache from old pool
  82. func (p *Pool) CloneFrom(o *Pool) {
  83. o.store.CloneTo(p.store)
  84. }
  85. func (p *Pool) get(host string) netip.Addr {
  86. p.offset = p.offset.Next()
  87. if !p.offset.Less(p.last) {
  88. p.cycle = true
  89. p.offset = p.first
  90. }
  91. if p.cycle || p.store.Exist(p.offset) {
  92. p.store.DelByIP(p.offset)
  93. }
  94. p.store.PutByIP(p.offset, host)
  95. return p.offset
  96. }
  97. func (p *Pool) FlushFakeIP() error {
  98. err := p.store.FlushFakeIP()
  99. if err == nil {
  100. p.cycle = false
  101. p.offset = p.first.Prev()
  102. }
  103. return err
  104. }
  105. func (p *Pool) StoreState() {
  106. if s, ok := p.store.(*cachefileStore); ok {
  107. s.PutByHost(offsetKey, p.offset)
  108. if p.cycle {
  109. s.PutByHost(cycleKey, p.offset)
  110. }
  111. }
  112. }
  113. func (p *Pool) restoreState() {
  114. if s, ok := p.store.(*cachefileStore); ok {
  115. if _, exist := s.GetByHost(cycleKey); exist {
  116. p.cycle = true
  117. }
  118. if offset, exist := s.GetByHost(offsetKey); exist {
  119. if p.ipnet.Contains(offset) {
  120. p.offset = offset
  121. } else {
  122. _ = p.FlushFakeIP()
  123. }
  124. } else if s.Exist(p.first) {
  125. _ = p.FlushFakeIP()
  126. }
  127. }
  128. }
  129. type Options struct {
  130. IPNet netip.Prefix
  131. Host *trie.DomainTrie[struct{}]
  132. // Size sets the maximum number of entries in memory
  133. // and does not work if Persistence is true
  134. Size int
  135. // Persistence will save the data to disk.
  136. // Size will not work and record will be fully stored.
  137. Persistence bool
  138. }
  139. // New return Pool instance
  140. func New(options Options) (*Pool, error) {
  141. var (
  142. hostAddr = options.IPNet.Masked().Addr()
  143. gateway = hostAddr.Next()
  144. first = gateway.Next().Next().Next() // default start with 198.18.0.4
  145. last = nnip.UnMasked(options.IPNet)
  146. )
  147. if !options.IPNet.IsValid() || !first.IsValid() || !first.Less(last) {
  148. return nil, errors.New("ipnet don't have valid ip")
  149. }
  150. pool := &Pool{
  151. gateway: gateway,
  152. first: first,
  153. last: last,
  154. offset: first.Prev(),
  155. cycle: false,
  156. host: options.Host,
  157. ipnet: options.IPNet,
  158. }
  159. if options.Persistence {
  160. pool.store = &cachefileStore{
  161. cache: cachefile.Cache(),
  162. }
  163. } else {
  164. pool.store = newMemoryStore(options.Size)
  165. }
  166. pool.restoreState()
  167. return pool, nil
  168. }