Refactor: fakeip pool use netip.Prefix, supports ipv6 range

component/fakeip/pool.go

@@ -2,39 +2,52 @@ package fakeip
 
 import (
 	"errors"
-	"net"
+	"math/bits"
+	"net/netip"
 	"sync"
+	_ "unsafe"
 
 	"github.com/Dreamacro/clash/component/profile/cachefile"
 	"github.com/Dreamacro/clash/component/trie"
 )
 
+//go:linkname beUint64 net/netip.beUint64
+func beUint64(b []byte) uint64
+
+//go:linkname bePutUint64 net/netip.bePutUint64
+func bePutUint64(b []byte, v uint64)
+
+type uint128 struct {
+	hi uint64
+	lo uint64
+}
+
 type store interface {
-	GetByHost(host string) (net.IP, bool)
-	PutByHost(host string, ip net.IP)
-	GetByIP(ip net.IP) (string, bool)
-	PutByIP(ip net.IP, host string)
-	DelByIP(ip net.IP)
-	Exist(ip net.IP) bool
+	GetByHost(host string) (netip.Addr, bool)
+	PutByHost(host string, ip netip.Addr)
+	GetByIP(ip netip.Addr) (string, bool)
+	PutByIP(ip netip.Addr, host string)
+	DelByIP(ip netip.Addr)
+	Exist(ip netip.Addr) bool
 	CloneTo(store)
 	FlushFakeIP() error
 }
 
 // Pool is a implementation about fake ip generator without storage
 type Pool struct {
-	max       uint32
-	min       uint32
-	gateway   uint32
-	broadcast uint32
-	offset    uint32
-	mux       sync.Mutex
-	host      *trie.DomainTrie[bool]
-	ipnet     *net.IPNet
-	store     store
+	gateway netip.Addr
+	first   netip.Addr
+	last    netip.Addr
+	offset  netip.Addr
+	cycle   bool
+	mux     sync.Mutex
+	host    *trie.DomainTrie[bool]
+	ipnet   *netip.Prefix
+	store   store
 }
 
 // Lookup return a fake ip with host
-func (p *Pool) Lookup(host string) net.IP {
+func (p *Pool) Lookup(host string) netip.Addr {
 	p.mux.Lock()
 	defer p.mux.Unlock()
 	if ip, exist := p.store.GetByHost(host); exist {
@@ -47,14 +60,10 @@ func (p *Pool) Lookup(host string) net.IP {
 }
 
 // LookBack return host with the fake ip
-func (p *Pool) LookBack(ip net.IP) (string, bool) {
+func (p *Pool) LookBack(ip netip.Addr) (string, bool) {
 	p.mux.Lock()
 	defer p.mux.Unlock()
 
-	if ip = ip.To4(); ip == nil {
-		return "", false
-	}
-
 	return p.store.GetByIP(ip)
 }
 
@@ -67,29 +76,25 @@ func (p *Pool) ShouldSkipped(domain string) bool {
 }
 
 // Exist returns if given ip exists in fake-ip pool
-func (p *Pool) Exist(ip net.IP) bool {
+func (p *Pool) Exist(ip netip.Addr) bool {
 	p.mux.Lock()
 	defer p.mux.Unlock()
 
-	if ip = ip.To4(); ip == nil {
-		return false
-	}
-
 	return p.store.Exist(ip)
 }
 
 // Gateway return gateway ip
-func (p *Pool) Gateway() net.IP {
-	return uintToIP(p.gateway)
+func (p *Pool) Gateway() netip.Addr {
+	return p.gateway
 }
 
-// Broadcast return broadcast ip
-func (p *Pool) Broadcast() net.IP {
-	return uintToIP(p.broadcast)
+// Broadcast return the last ip
+func (p *Pool) Broadcast() netip.Addr {
+	return p.last
 }
 
 // IPNet return raw ipnet
-func (p *Pool) IPNet() *net.IPNet {
+func (p *Pool) IPNet() *netip.Prefix {
 	return p.ipnet
 }
 
@@ -98,46 +103,28 @@ func (p *Pool) CloneFrom(o *Pool) {
 	o.store.CloneTo(p.store)
 }
 
-func (p *Pool) get(host string) net.IP {
-	current := p.offset
-	for {
-		p.offset = (p.offset + 1) % (p.max - p.min)
-		// Avoid infinite loops
-		if p.offset == current {
-			p.offset = (p.offset + 1) % (p.max - p.min)
-			ip := uintToIP(p.min + p.offset - 1)
-			p.store.DelByIP(ip)
-			break
-		}
+func (p *Pool) get(host string) netip.Addr {
+	p.offset = p.offset.Next()
 
-		ip := uintToIP(p.min + p.offset - 1)
-		if !p.store.Exist(ip) {
-			break
-		}
+	if !p.offset.Less(p.last) {
+		p.cycle = true
+		p.offset = p.first
 	}
-	ip := uintToIP(p.min + p.offset - 1)
-	p.store.PutByIP(ip, host)
-	return ip
+
+	if p.cycle {
+		p.store.DelByIP(p.offset)
+	}
+
+	p.store.PutByIP(p.offset, host)
+	return p.offset
 }
 
 func (p *Pool) FlushFakeIP() error {
 	return p.store.FlushFakeIP()
 }
 
-func ipToUint(ip net.IP) uint32 {
-	v := uint32(ip[0]) << 24
-	v += uint32(ip[1]) << 16
-	v += uint32(ip[2]) << 8
-	v += uint32(ip[3])
-	return v
-}
-
-func uintToIP(v uint32) net.IP {
-	return net.IP{byte(v >> 24), byte(v >> 16), byte(v >> 8), byte(v)}
-}
-
 type Options struct {
-	IPNet *net.IPNet
+	IPNet *netip.Prefix
 	Host  *trie.DomainTrie[bool]
 
 	// Size sets the maximum number of entries in memory
@@ -151,23 +138,25 @@ type Options struct {
 
 // New return Pool instance
 func New(options Options) (*Pool, error) {
-	min := ipToUint(options.IPNet.IP) + 3
+	var (
+		hostAddr = options.IPNet.Masked().Addr()
+		gateway  = hostAddr.Next()
+		first    = gateway.Next().Next()
+		last     = add(hostAddr, 1<<uint64(hostAddr.BitLen()-options.IPNet.Bits())-1)
+	)
 
-	ones, bits := options.IPNet.Mask.Size()
-	total := 1<<uint(bits-ones) - 4
-
-	if total <= 0 {
+	if !options.IPNet.IsValid() || !first.Less(last) || !options.IPNet.Contains(last) {
 		return nil, errors.New("ipnet don't have valid ip")
 	}
 
-	max := min + uint32(total) - 1
 	pool := &Pool{
-		min:       min,
-		max:       max,
-		gateway:   min - 2,
-		broadcast: max + 1,
-		host:      options.Host,
-		ipnet:     options.IPNet,
+		gateway: gateway,
+		first:   first,
+		last:    last,
+		offset:  first.Prev(),
+		cycle:   false,
+		host:    options.Host,
+		ipnet:   options.IPNet,
 	}
 	if options.Persistence {
 		pool.store = &cachefileStore{
@@ -179,3 +168,29 @@ func New(options Options) (*Pool, error) {
 
 	return pool, nil
 }
+
+// add returns addr + n.
+func add(addr netip.Addr, n uint64) netip.Addr {
+	buf := addr.As16()
+
+	u := uint128{
+		beUint64(buf[:8]),
+		beUint64(buf[8:]),
+	}
+
+	lo, carry := bits.Add64(u.lo, n, 0)
+
+	u.hi = u.hi + carry
+	u.lo = lo
+
+	bePutUint64(buf[:8], u.hi)
+	bePutUint64(buf[8:], u.lo)
+
+	a := netip.AddrFrom16(buf)
+
+	if addr.Is4() {
+		return a.Unmap()
+	}
+
+	return a
+}