github.com/chwjbn/xclash@v0.2.0/component/fakeip/pool.go (about)

     1  package fakeip
     2  
     3  import (
     4  	"errors"
     5  	"net"
     6  	"sync"
     7  
     8  	"github.com/chwjbn/xclash/common/cache"
     9  	"github.com/chwjbn/xclash/component/profile/cachefile"
    10  	"github.com/chwjbn/xclash/component/trie"
    11  )
    12  
    13  type store interface {
    14  	GetByHost(host string) (net.IP, bool)
    15  	PutByHost(host string, ip net.IP)
    16  	GetByIP(ip net.IP) (string, bool)
    17  	PutByIP(ip net.IP, host string)
    18  	DelByIP(ip net.IP)
    19  	Exist(ip net.IP) bool
    20  	CloneTo(store)
    21  }
    22  
    23  // Pool is a implementation about fake ip generator without storage
    24  type Pool struct {
    25  	max     uint32
    26  	min     uint32
    27  	gateway uint32
    28  	offset  uint32
    29  	mux     sync.Mutex
    30  	host    *trie.DomainTrie
    31  	ipnet   *net.IPNet
    32  	store   store
    33  }
    34  
    35  // Lookup return a fake ip with host
    36  func (p *Pool) Lookup(host string) net.IP {
    37  	p.mux.Lock()
    38  	defer p.mux.Unlock()
    39  	if ip, exist := p.store.GetByHost(host); exist {
    40  		return ip
    41  	}
    42  
    43  	ip := p.get(host)
    44  	p.store.PutByHost(host, ip)
    45  	return ip
    46  }
    47  
    48  // LookBack return host with the fake ip
    49  func (p *Pool) LookBack(ip net.IP) (string, bool) {
    50  	p.mux.Lock()
    51  	defer p.mux.Unlock()
    52  
    53  	if ip = ip.To4(); ip == nil {
    54  		return "", false
    55  	}
    56  
    57  	return p.store.GetByIP(ip)
    58  }
    59  
    60  // ShouldSkipped return if domain should be skipped
    61  func (p *Pool) ShouldSkipped(domain string) bool {
    62  	if p.host == nil {
    63  		return false
    64  	}
    65  	return p.host.Search(domain) != nil
    66  }
    67  
    68  // Exist returns if given ip exists in fake-ip pool
    69  func (p *Pool) Exist(ip net.IP) bool {
    70  	p.mux.Lock()
    71  	defer p.mux.Unlock()
    72  
    73  	if ip = ip.To4(); ip == nil {
    74  		return false
    75  	}
    76  
    77  	return p.store.Exist(ip)
    78  }
    79  
    80  // Gateway return gateway ip
    81  func (p *Pool) Gateway() net.IP {
    82  	return uintToIP(p.gateway)
    83  }
    84  
    85  // IPNet return raw ipnet
    86  func (p *Pool) IPNet() *net.IPNet {
    87  	return p.ipnet
    88  }
    89  
    90  // CloneFrom clone cache from old pool
    91  func (p *Pool) CloneFrom(o *Pool) {
    92  	o.store.CloneTo(p.store)
    93  }
    94  
    95  func (p *Pool) get(host string) net.IP {
    96  	current := p.offset
    97  	for {
    98  		p.offset = (p.offset + 1) % (p.max - p.min)
    99  		// Avoid infinite loops
   100  		if p.offset == current {
   101  			p.offset = (p.offset + 1) % (p.max - p.min)
   102  			ip := uintToIP(p.min + p.offset - 1)
   103  			p.store.DelByIP(ip)
   104  			break
   105  		}
   106  
   107  		ip := uintToIP(p.min + p.offset - 1)
   108  		if !p.store.Exist(ip) {
   109  			break
   110  		}
   111  	}
   112  	ip := uintToIP(p.min + p.offset - 1)
   113  	p.store.PutByIP(ip, host)
   114  	return ip
   115  }
   116  
   117  func ipToUint(ip net.IP) uint32 {
   118  	v := uint32(ip[0]) << 24
   119  	v += uint32(ip[1]) << 16
   120  	v += uint32(ip[2]) << 8
   121  	v += uint32(ip[3])
   122  	return v
   123  }
   124  
   125  func uintToIP(v uint32) net.IP {
   126  	return net.IP{byte(v >> 24), byte(v >> 16), byte(v >> 8), byte(v)}
   127  }
   128  
   129  type Options struct {
   130  	IPNet *net.IPNet
   131  	Host  *trie.DomainTrie
   132  
   133  	// Size sets the maximum number of entries in memory
   134  	// and does not work if Persistence is true
   135  	Size int
   136  
   137  	// Persistence will save the data to disk.
   138  	// Size will not work and record will be fully stored.
   139  	Persistence bool
   140  }
   141  
   142  // New return Pool instance
   143  func New(options Options) (*Pool, error) {
   144  	min := ipToUint(options.IPNet.IP) + 2
   145  
   146  	ones, bits := options.IPNet.Mask.Size()
   147  	total := 1<<uint(bits-ones) - 2
   148  
   149  	if total <= 0 {
   150  		return nil, errors.New("ipnet don't have valid ip")
   151  	}
   152  
   153  	max := min + uint32(total) - 1
   154  	pool := &Pool{
   155  		min:     min,
   156  		max:     max,
   157  		gateway: min - 1,
   158  		host:    options.Host,
   159  		ipnet:   options.IPNet,
   160  	}
   161  	if options.Persistence {
   162  		pool.store = &cachefileStore{
   163  			cache: cachefile.Cache(),
   164  		}
   165  	} else {
   166  		pool.store = &memoryStore{
   167  			cache: cache.NewLRUCache(cache.WithSize(options.Size * 2)),
   168  		}
   169  	}
   170  
   171  	return pool, nil
   172  }