From 3d5681cffd371ecdb9b340fd963872e8416d32e9 Mon Sep 17 00:00:00 2001
From: Dreamacro <chuainian@gmail.com>
Date: Mon, 11 Oct 2021 20:48:58 +0800
Subject: [PATCH] Feature: persistence fakeip (#1662)

---
 component/fakeip/cachefile.go        |  49 +++++++++
 component/fakeip/memory.go           |  60 +++++++++++
 component/fakeip/pool.go             |  93 ++++++++++-------
 component/fakeip/pool_test.go        | 147 +++++++++++++++++++++++----
 component/profile/cachefile/cache.go |  75 +++++++++++---
 config/config.go                     |  13 ++-
 dns/enhancer.go                      |   2 +-
 dns/middleware.go                    |   2 +-
 8 files changed, 366 insertions(+), 75 deletions(-)
 create mode 100644 component/fakeip/cachefile.go
 create mode 100644 component/fakeip/memory.go

diff --git a/component/fakeip/cachefile.go b/component/fakeip/cachefile.go
new file mode 100644
index 0000000000..7ee889812a
--- /dev/null
+++ b/component/fakeip/cachefile.go
@@ -0,0 +1,49 @@
+package fakeip
+
+import (
+	"net"
+
+	"github.com/Dreamacro/clash/component/profile/cachefile"
+)
+
+type cachefileStore struct {
+	cache *cachefile.CacheFile
+}
+
+// GetByHost implements store.GetByHost
+func (c *cachefileStore) GetByHost(host string) (net.IP, bool) {
+	elm := c.cache.GetFakeip([]byte(host))
+	if elm == nil {
+		return nil, false
+	}
+	return net.IP(elm), true
+}
+
+// PutByHost implements store.PutByHost
+func (c *cachefileStore) PutByHost(host string, ip net.IP) {
+	c.cache.PutFakeip([]byte(host), ip)
+}
+
+// GetByIP implements store.GetByIP
+func (c *cachefileStore) GetByIP(ip net.IP) (string, bool) {
+	elm := c.cache.GetFakeip(ip.To4())
+	if elm == nil {
+		return "", false
+	}
+	return string(elm), true
+}
+
+// PutByIP implements store.PutByIP
+func (c *cachefileStore) PutByIP(ip net.IP, host string) {
+	c.cache.PutFakeip(ip.To4(), []byte(host))
+}
+
+// Exist implements store.Exist
+func (c *cachefileStore) Exist(ip net.IP) bool {
+	_, exist := c.GetByIP(ip)
+	return exist
+}
+
+// CloneTo implements store.CloneTo
+// already persistence
+func (c *cachefileStore) CloneTo(store store) {}
diff --git a/component/fakeip/memory.go b/component/fakeip/memory.go
new file mode 100644
index 0000000000..75d4a3b243
--- /dev/null
+++ b/component/fakeip/memory.go
@@ -0,0 +1,60 @@
+package fakeip
+
+import (
+	"net"
+
+	"github.com/Dreamacro/clash/common/cache"
+)
+
+type memoryStore struct {
+	cache *cache.LruCache
+}
+
+// GetByHost implements store.GetByHost
+func (m *memoryStore) GetByHost(host string) (net.IP, bool) {
+	if elm, exist := m.cache.Get(host); exist {
+		ip := elm.(net.IP)
+
+		// ensure ip --> host on head of linked list
+		m.cache.Get(ipToUint(ip.To4()))
+		return ip, true
+	}
+
+	return nil, false
+}
+
+// PutByHost implements store.PutByHost
+func (m *memoryStore) PutByHost(host string, ip net.IP) {
+	m.cache.Set(host, ip)
+}
+
+// GetByIP implements store.GetByIP
+func (m *memoryStore) GetByIP(ip net.IP) (string, bool) {
+	if elm, exist := m.cache.Get(ipToUint(ip.To4())); exist {
+		host := elm.(string)
+
+		// ensure host --> ip on head of linked list
+		m.cache.Get(host)
+		return host, true
+	}
+
+	return "", false
+}
+
+// PutByIP implements store.PutByIP
+func (m *memoryStore) PutByIP(ip net.IP, host string) {
+	m.cache.Set(ipToUint(ip.To4()), host)
+}
+
+// Exist implements store.Exist
+func (m *memoryStore) Exist(ip net.IP) bool {
+	return m.cache.Exist(ipToUint(ip.To4()))
+}
+
+// CloneTo implements store.CloneTo
+// only for memoryStore to memoryStore
+func (m *memoryStore) CloneTo(store store) {
+	if ms, ok := store.(*memoryStore); ok {
+		m.cache.CloneTo(ms.cache)
+	}
+}
diff --git a/component/fakeip/pool.go b/component/fakeip/pool.go
index 97d812ac3f..180d5eb189 100644
--- a/component/fakeip/pool.go
+++ b/component/fakeip/pool.go
@@ -6,9 +6,19 @@ import (
 	"sync"
 
 	"github.com/Dreamacro/clash/common/cache"
+	"github.com/Dreamacro/clash/component/profile/cachefile"
 	"github.com/Dreamacro/clash/component/trie"
 )
 
+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)
+	Exist(ip net.IP) bool
+	CloneTo(store)
+}
+
 // Pool is a implementation about fake ip generator without storage
 type Pool struct {
 	max     uint32
@@ -18,25 +28,19 @@ type Pool struct {
 	mux     sync.Mutex
 	host    *trie.DomainTrie
 	ipnet   *net.IPNet
-	cache   *cache.LruCache
+	store   store
 }
 
 // Lookup return a fake ip with host
 func (p *Pool) Lookup(host string) net.IP {
 	p.mux.Lock()
 	defer p.mux.Unlock()
-	if elm, exist := p.cache.Get(host); exist {
-		ip := elm.(net.IP)
-
-		// ensure ip --> host on head of linked list
-		n := ipToUint(ip.To4())
-		offset := n - p.min + 1
-		p.cache.Get(offset)
+	if ip, exist := p.store.GetByHost(host); exist {
 		return ip
 	}
 
 	ip := p.get(host)
-	p.cache.Set(host, ip)
+	p.store.PutByHost(host, ip)
 	return ip
 }
 
@@ -49,22 +53,11 @@ func (p *Pool) LookBack(ip net.IP) (string, bool) {
 		return "", false
 	}
 
-	n := ipToUint(ip.To4())
-	offset := n - p.min + 1
-
-	if elm, exist := p.cache.Get(offset); exist {
-		host := elm.(string)
-
-		// ensure host --> ip on head of linked list
-		p.cache.Get(host)
-		return host, true
-	}
-
-	return "", false
+	return p.store.GetByIP(ip)
 }
 
-// LookupHost return if domain in host
-func (p *Pool) LookupHost(domain string) bool {
+// ShouldSkipped return if domain should be skipped
+func (p *Pool) ShouldSkipped(domain string) bool {
 	if p.host == nil {
 		return false
 	}
@@ -80,9 +73,7 @@ func (p *Pool) Exist(ip net.IP) bool {
 		return false
 	}
 
-	n := ipToUint(ip.To4())
-	offset := n - p.min + 1
-	return p.cache.Exist(offset)
+	return p.store.Exist(ip)
 }
 
 // Gateway return gateway ip
@@ -95,9 +86,9 @@ func (p *Pool) IPNet() *net.IPNet {
 	return p.ipnet
 }
 
-// PatchFrom clone cache from old pool
-func (p *Pool) PatchFrom(o *Pool) {
-	o.cache.CloneTo(p.cache)
+// CloneFrom clone cache from old pool
+func (p *Pool) CloneFrom(o *Pool) {
+	o.store.CloneTo(p.store)
 }
 
 func (p *Pool) get(host string) net.IP {
@@ -109,12 +100,13 @@ func (p *Pool) get(host string) net.IP {
 			break
 		}
 
-		if !p.cache.Exist(p.offset) {
+		ip := uintToIP(p.min + p.offset - 1)
+		if !p.store.Exist(ip) {
 			break
 		}
 	}
 	ip := uintToIP(p.min + p.offset - 1)
-	p.cache.Set(p.offset, host)
+	p.store.PutByIP(ip, host)
 	return ip
 }
 
@@ -130,11 +122,24 @@ 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
+	Host  *trie.DomainTrie
+
+	// Size sets the maximum number of entries in memory
+	// and does not work if Persistence is true
+	Size int
+
+	// Persistence will save the data to disk.
+	// Size will not work and record will be fully stored.
+	Persistence bool
+}
+
 // New return Pool instance
-func New(ipnet *net.IPNet, size int, host *trie.DomainTrie) (*Pool, error) {
-	min := ipToUint(ipnet.IP) + 2
+func New(options Options) (*Pool, error) {
+	min := ipToUint(options.IPNet.IP) + 2
 
-	ones, bits := ipnet.Mask.Size()
+	ones, bits := options.IPNet.Mask.Size()
 	total := 1<<uint(bits-ones) - 2
 
 	if total <= 0 {
@@ -142,12 +147,22 @@ func New(ipnet *net.IPNet, size int, host *trie.DomainTrie) (*Pool, error) {
 	}
 
 	max := min + uint32(total) - 1
-	return &Pool{
+	pool := &Pool{
 		min:     min,
 		max:     max,
 		gateway: min - 1,
-		host:    host,
-		ipnet:   ipnet,
-		cache:   cache.NewLRUCache(cache.WithSize(size * 2)),
-	}, nil
+		host:    options.Host,
+		ipnet:   options.IPNet,
+	}
+	if options.Persistence {
+		pool.store = &cachefileStore{
+			cache: cachefile.Cache(),
+		}
+	} else {
+		pool.store = &memoryStore{
+			cache: cache.NewLRUCache(cache.WithSize(options.Size * 2)),
+		}
+	}
+
+	return pool, nil
 }
diff --git a/component/fakeip/pool_test.go b/component/fakeip/pool_test.go
index 89b161b89f..bd4af4c949 100644
--- a/component/fakeip/pool_test.go
+++ b/component/fakeip/pool_test.go
@@ -2,38 +2,118 @@ package fakeip
 
 import (
 	"net"
+	"os"
 	"testing"
+	"time"
+
+	"github.com/Dreamacro/clash/component/profile/cachefile"
+	"github.com/Dreamacro/clash/component/trie"
 
 	"github.com/stretchr/testify/assert"
+	"go.etcd.io/bbolt"
 )
 
-func TestPool_Basic(t *testing.T) {
-	_, ipnet, _ := net.ParseCIDR("192.168.0.1/29")
-	pool, _ := New(ipnet, 10, nil)
+func createPools(options Options) ([]*Pool, string, error) {
+	pool, err := New(options)
+	if err != nil {
+		return nil, "", err
+	}
+	filePool, tempfile, err := createCachefileStore(options)
+	if err != nil {
+		return nil, "", err
+	}
+
+	return []*Pool{pool, filePool}, tempfile, nil
+}
 
-	first := pool.Lookup("foo.com")
-	last := pool.Lookup("bar.com")
-	bar, exist := pool.LookBack(last)
+func createCachefileStore(options Options) (*Pool, string, error) {
+	pool, err := New(options)
+	if err != nil {
+		return nil, "", err
+	}
+	f, err := os.CreateTemp("", "clash")
+	if err != nil {
+		return nil, "", err
+	}
+
+	db, err := bbolt.Open(f.Name(), 0o666, &bbolt.Options{Timeout: time.Second})
+	if err != nil {
+		return nil, "", err
+	}
+
+	pool.store = &cachefileStore{
+		cache: &cachefile.CacheFile{DB: db},
+	}
+	return pool, f.Name(), nil
+}
 
-	assert.True(t, first.Equal(net.IP{192, 168, 0, 2}))
-	assert.True(t, last.Equal(net.IP{192, 168, 0, 3}))
-	assert.True(t, exist)
-	assert.Equal(t, bar, "bar.com")
+func TestPool_Basic(t *testing.T) {
+	_, ipnet, _ := net.ParseCIDR("192.168.0.1/29")
+	pools, tempfile, err := createPools(Options{
+		IPNet: ipnet,
+		Size:  10,
+	})
+	assert.Nil(t, err)
+	defer os.Remove(tempfile)
+
+	for _, pool := range pools {
+		first := pool.Lookup("foo.com")
+		last := pool.Lookup("bar.com")
+		bar, exist := pool.LookBack(last)
+
+		assert.True(t, first.Equal(net.IP{192, 168, 0, 2}))
+		assert.Equal(t, pool.Lookup("foo.com"), net.IP{192, 168, 0, 2})
+		assert.True(t, last.Equal(net.IP{192, 168, 0, 3}))
+		assert.True(t, exist)
+		assert.Equal(t, bar, "bar.com")
+		assert.Equal(t, pool.Gateway(), net.IP{192, 168, 0, 1})
+		assert.Equal(t, pool.IPNet().String(), ipnet.String())
+		assert.True(t, pool.Exist(net.IP{192, 168, 0, 3}))
+		assert.False(t, pool.Exist(net.IP{192, 168, 0, 4}))
+		assert.False(t, pool.Exist(net.ParseIP("::1")))
+	}
 }
 
-func TestPool_Cycle(t *testing.T) {
+func TestPool_CycleUsed(t *testing.T) {
 	_, ipnet, _ := net.ParseCIDR("192.168.0.1/30")
-	pool, _ := New(ipnet, 10, nil)
-
-	first := pool.Lookup("foo.com")
-	same := pool.Lookup("baz.com")
+	pools, tempfile, err := createPools(Options{
+		IPNet: ipnet,
+		Size:  10,
+	})
+	assert.Nil(t, err)
+	defer os.Remove(tempfile)
+
+	for _, pool := range pools {
+		first := pool.Lookup("foo.com")
+		same := pool.Lookup("baz.com")
+		assert.True(t, first.Equal(same))
+	}
+}
 
-	assert.True(t, first.Equal(same))
+func TestPool_Skip(t *testing.T) {
+	_, ipnet, _ := net.ParseCIDR("192.168.0.1/30")
+	tree := trie.New()
+	tree.Insert("example.com", tree)
+	pools, tempfile, err := createPools(Options{
+		IPNet: ipnet,
+		Size:  10,
+		Host:  tree,
+	})
+	assert.Nil(t, err)
+	defer os.Remove(tempfile)
+
+	for _, pool := range pools {
+		assert.True(t, pool.ShouldSkipped("example.com"))
+		assert.False(t, pool.ShouldSkipped("foo.com"))
+	}
 }
 
 func TestPool_MaxCacheSize(t *testing.T) {
 	_, ipnet, _ := net.ParseCIDR("192.168.0.1/24")
-	pool, _ := New(ipnet, 2, nil)
+	pool, _ := New(Options{
+		IPNet: ipnet,
+		Size:  2,
+	})
 
 	first := pool.Lookup("foo.com")
 	pool.Lookup("bar.com")
@@ -45,7 +125,10 @@ func TestPool_MaxCacheSize(t *testing.T) {
 
 func TestPool_DoubleMapping(t *testing.T) {
 	_, ipnet, _ := net.ParseCIDR("192.168.0.1/24")
-	pool, _ := New(ipnet, 2, nil)
+	pool, _ := New(Options{
+		IPNet: ipnet,
+		Size:  2,
+	})
 
 	// fill cache
 	fooIP := pool.Lookup("foo.com")
@@ -70,9 +153,35 @@ func TestPool_DoubleMapping(t *testing.T) {
 	assert.False(t, bazIP.Equal(newBazIP))
 }
 
+func TestPool_Clone(t *testing.T) {
+	_, ipnet, _ := net.ParseCIDR("192.168.0.1/24")
+	pool, _ := New(Options{
+		IPNet: ipnet,
+		Size:  2,
+	})
+
+	first := pool.Lookup("foo.com")
+	last := pool.Lookup("bar.com")
+	assert.True(t, first.Equal(net.IP{192, 168, 0, 2}))
+	assert.True(t, last.Equal(net.IP{192, 168, 0, 3}))
+
+	newPool, _ := New(Options{
+		IPNet: ipnet,
+		Size:  2,
+	})
+	newPool.CloneFrom(pool)
+	_, firstExist := newPool.LookBack(first)
+	_, lastExist := newPool.LookBack(last)
+	assert.True(t, firstExist)
+	assert.True(t, lastExist)
+}
+
 func TestPool_Error(t *testing.T) {
 	_, ipnet, _ := net.ParseCIDR("192.168.0.1/31")
-	_, err := New(ipnet, 10, nil)
+	_, err := New(Options{
+		IPNet: ipnet,
+		Size:  10,
+	})
 
 	assert.Error(t, err)
 }
diff --git a/component/profile/cachefile/cache.go b/component/profile/cachefile/cache.go
index 03f389a3a0..1ef47493a7 100644
--- a/component/profile/cachefile/cache.go
+++ b/component/profile/cachefile/cache.go
@@ -11,7 +11,7 @@ import (
 	C "github.com/Dreamacro/clash/constant"
 	"github.com/Dreamacro/clash/log"
 
-	bolt "go.etcd.io/bbolt"
+	"go.etcd.io/bbolt"
 )
 
 var (
@@ -20,21 +20,22 @@ var (
 	defaultCache *CacheFile
 
 	bucketSelected = []byte("selected")
+	bucketFakeip   = []byte("fakeip")
 )
 
 // CacheFile store and update the cache file
 type CacheFile struct {
-	db *bolt.DB
+	DB *bbolt.DB
 }
 
 func (c *CacheFile) SetSelected(group, selected string) {
 	if !profile.StoreSelected.Load() {
 		return
-	} else if c.db == nil {
+	} else if c.DB == nil {
 		return
 	}
 
-	err := c.db.Batch(func(t *bolt.Tx) error {
+	err := c.DB.Batch(func(t *bbolt.Tx) error {
 		bucket, err := t.CreateBucketIfNotExists(bucketSelected)
 		if err != nil {
 			return err
@@ -42,7 +43,7 @@ func (c *CacheFile) SetSelected(group, selected string) {
 		return bucket.Put([]byte(group), []byte(selected))
 	})
 	if err != nil {
-		log.Warnln("[CacheFile] write cache to %s failed: %s", c.db.Path(), err.Error())
+		log.Warnln("[CacheFile] write cache to %s failed: %s", c.DB.Path(), err.Error())
 		return
 	}
 }
@@ -50,12 +51,12 @@ func (c *CacheFile) SetSelected(group, selected string) {
 func (c *CacheFile) SelectedMap() map[string]string {
 	if !profile.StoreSelected.Load() {
 		return nil
-	} else if c.db == nil {
+	} else if c.DB == nil {
 		return nil
 	}
 
 	mapping := map[string]string{}
-	c.db.View(func(t *bolt.Tx) error {
+	c.DB.View(func(t *bbolt.Tx) error {
 		bucket := t.Bucket(bucketSelected)
 		if bucket == nil {
 			return nil
@@ -70,19 +71,68 @@ func (c *CacheFile) SelectedMap() map[string]string {
 	return mapping
 }
 
+func (c *CacheFile) PutFakeip(key, value []byte) error {
+	if c.DB == nil {
+		return nil
+	}
+
+	err := c.DB.Batch(func(t *bbolt.Tx) error {
+		bucket, err := t.CreateBucketIfNotExists(bucketFakeip)
+		if err != nil {
+			return err
+		}
+		return bucket.Put(key, value)
+	})
+	if err != nil {
+		log.Warnln("[CacheFile] write cache to %s failed: %s", c.DB.Path(), err.Error())
+	}
+
+	return err
+}
+
+func (c *CacheFile) GetFakeip(key []byte) []byte {
+	if c.DB == nil {
+		return nil
+	}
+
+	tx, err := c.DB.Begin(false)
+	if err != nil {
+		return nil
+	}
+	defer tx.Rollback()
+
+	bucket := tx.Bucket(bucketFakeip)
+	if bucket == nil {
+		return nil
+	}
+
+	return bucket.Get(key)
+}
+
 func (c *CacheFile) Close() error {
-	return c.db.Close()
+	return c.DB.Close()
 }
 
 // TODO: remove migrateCache until 2022
 func migrateCache() {
 	defer func() {
-		db, err := bolt.Open(C.Path.Cache(), fileMode, &bolt.Options{Timeout: time.Second})
+		options := bbolt.Options{Timeout: time.Second}
+		db, err := bbolt.Open(C.Path.Cache(), fileMode, &options)
+		switch err {
+		case bbolt.ErrInvalid, bbolt.ErrChecksum, bbolt.ErrVersionMismatch:
+			if err = os.Remove(C.Path.Cache()); err != nil {
+				log.Warnln("[CacheFile] remove invalid cache file error: %s", err.Error())
+				break
+			}
+			log.Infoln("[CacheFile] remove invalid cache file and create new one")
+			db, err = bbolt.Open(C.Path.Cache(), fileMode, &options)
+		}
 		if err != nil {
 			log.Warnln("[CacheFile] can't open cache file: %s", err.Error())
 		}
+
 		defaultCache = &CacheFile{
-			db: db,
+			DB: db,
 		}
 	}()
 
@@ -103,12 +153,13 @@ func migrateCache() {
 	gob.NewDecoder(bufReader).Decode(model)
 
 	// write to new cache file
-	db, err := bolt.Open(C.Path.Cache(), fileMode, nil)
+	db, err := bbolt.Open(C.Path.Cache(), fileMode, nil)
 	if err != nil {
 		return
 	}
 	defer db.Close()
-	db.Batch(func(t *bolt.Tx) error {
+
+	db.Batch(func(t *bbolt.Tx) error {
 		bucket, err := t.CreateBucketIfNotExists(bucketSelected)
 		if err != nil {
 			return err
diff --git a/config/config.go b/config/config.go
index 004a935f0c..f532c61a14 100644
--- a/config/config.go
+++ b/config/config.go
@@ -80,6 +80,7 @@ type FallbackFilter struct {
 // Profile config
 type Profile struct {
 	StoreSelected bool `yaml:"store-selected"`
+	StoreFakeIP   bool `yaml:"store-fakeip"`
 }
 
 // Experimental config
@@ -226,7 +227,7 @@ func ParseRawConfig(rawCfg *RawConfig) (*Config, error) {
 	}
 	config.Hosts = hosts
 
-	dnsCfg, err := parseDNS(rawCfg.DNS, hosts)
+	dnsCfg, err := parseDNS(rawCfg, hosts)
 	if err != nil {
 		return nil, err
 	}
@@ -541,7 +542,8 @@ func parseFallbackIPCIDR(ips []string) ([]*net.IPNet, error) {
 	return ipNets, nil
 }
 
-func parseDNS(cfg RawDNS, hosts *trie.DomainTrie) (*DNS, error) {
+func parseDNS(rawCfg *RawConfig, hosts *trie.DomainTrie) (*DNS, error) {
+	cfg := rawCfg.DNS
 	if cfg.Enable && len(cfg.NameServer) == 0 {
 		return nil, fmt.Errorf("if DNS configuration is turned on, NameServer cannot be empty")
 	}
@@ -597,7 +599,12 @@ func parseDNS(cfg RawDNS, hosts *trie.DomainTrie) (*DNS, error) {
 			}
 		}
 
-		pool, err := fakeip.New(ipnet, 1000, host)
+		pool, err := fakeip.New(fakeip.Options{
+			IPNet:       ipnet,
+			Size:        1000,
+			Host:        host,
+			Persistence: rawCfg.Profile.StoreFakeIP,
+		})
 		if err != nil {
 			return nil, err
 		}
diff --git a/dns/enhancer.go b/dns/enhancer.go
index 0b2570059d..d2592e6e9e 100644
--- a/dns/enhancer.go
+++ b/dns/enhancer.go
@@ -67,7 +67,7 @@ func (h *ResolverEnhancer) PatchFrom(o *ResolverEnhancer) {
 	}
 
 	if h.fakePool != nil && o.fakePool != nil {
-		h.fakePool.PatchFrom(o.fakePool)
+		h.fakePool.CloneFrom(o.fakePool)
 	}
 }
 
diff --git a/dns/middleware.go b/dns/middleware.go
index da84788d70..de0b3aff47 100644
--- a/dns/middleware.go
+++ b/dns/middleware.go
@@ -107,7 +107,7 @@ func withFakeIP(fakePool *fakeip.Pool) middleware {
 			q := r.Question[0]
 
 			host := strings.TrimRight(q.Name, ".")
-			if fakePool.LookupHost(host) {
+			if fakePool.ShouldSkipped(host) {
 				return next(ctx, r)
 			}