From 3ea3653d7a341654dbddfc18a60a3a91dc303969 Mon Sep 17 00:00:00 2001
From: yaling888 <73897884+yaling888@users.noreply.github.com>
Date: Wed, 13 Apr 2022 05:55:08 +0800
Subject: [PATCH] Chore: persistence fakeip pool state

---
 component/fakeip/pool.go       | 62 ++++++++++++++++++++++++++--------
 component/resolver/enhancer.go |  7 ++++
 dns/enhancer.go                | 12 +++++--
 hub/executor/executor.go       |  1 +
 4 files changed, 64 insertions(+), 18 deletions(-)

diff --git a/component/fakeip/pool.go b/component/fakeip/pool.go
index e1a8150a..57613a48 100644
--- a/component/fakeip/pool.go
+++ b/component/fakeip/pool.go
@@ -11,6 +11,11 @@ import (
 	"github.com/Dreamacro/clash/component/trie"
 )
 
+const (
+	offsetKey = "key-offset-fake-ip"
+	cycleKey  = "key-cycle-fake-ip"
+)
+
 type uint128 struct {
 	hi uint64
 	lo uint64
@@ -98,22 +103,15 @@ func (p *Pool) CloneFrom(o *Pool) {
 }
 
 func (p *Pool) get(host string) netip.Addr {
-	for {
-		p.offset = p.offset.Next()
+	p.offset = p.offset.Next()
 
-		if !p.offset.Less(p.last) {
-			p.cycle = true
-			p.offset = p.first
-		}
+	if !p.offset.Less(p.last) {
+		p.cycle = true
+		p.offset = p.first
+	}
 
-		if p.cycle {
-			p.store.DelByIP(p.offset)
-			break
-		}
-
-		if !p.store.Exist(p.offset) {
-			break
-		}
+	if p.cycle || p.store.Exist(p.offset) {
+		p.store.DelByIP(p.offset)
 	}
 
 	p.store.PutByIP(p.offset, host)
@@ -121,7 +119,39 @@ func (p *Pool) get(host string) netip.Addr {
 }
 
 func (p *Pool) FlushFakeIP() error {
-	return p.store.FlushFakeIP()
+	err := p.store.FlushFakeIP()
+	if err == nil {
+		p.cycle = false
+		p.offset = p.first
+	}
+	return err
+}
+
+func (p *Pool) StoreState() {
+	if s, ok := p.store.(*cachefileStore); ok {
+		s.PutByHost(offsetKey, p.offset)
+		if p.cycle {
+			s.PutByHost(cycleKey, p.offset)
+		}
+	}
+}
+
+func (p *Pool) restoreState() {
+	if s, ok := p.store.(*cachefileStore); ok {
+		if _, exist := s.GetByHost(cycleKey); exist {
+			p.cycle = true
+		}
+
+		if offset, exist := s.GetByHost(offsetKey); exist {
+			if p.ipnet.Contains(offset) {
+				p.offset = offset
+			} else {
+				_ = p.FlushFakeIP()
+			}
+		} else if s.Exist(p.first) {
+			_ = p.FlushFakeIP()
+		}
+	}
 }
 
 type Options struct {
@@ -167,6 +197,8 @@ func New(options Options) (*Pool, error) {
 		pool.store = newMemoryStore(options.Size)
 	}
 
+	pool.restoreState()
+
 	return pool, nil
 }
 
diff --git a/component/resolver/enhancer.go b/component/resolver/enhancer.go
index 77f18374..503211c4 100644
--- a/component/resolver/enhancer.go
+++ b/component/resolver/enhancer.go
@@ -15,6 +15,7 @@ type Enhancer interface {
 	FindHostByIP(net.IP) (string, bool)
 	FlushFakeIP() error
 	InsertHostByIP(net.IP, string)
+	StoreFakePoolSate()
 }
 
 func FakeIPEnabled() bool {
@@ -77,3 +78,9 @@ func FlushFakeIP() error {
 	}
 	return nil
 }
+
+func StoreFakePoolSate() {
+	if mapper := DefaultHostMapper; mapper != nil {
+		mapper.StoreFakePoolSate()
+	}
+}
diff --git a/dns/enhancer.go b/dns/enhancer.go
index 2190f6aa..4e828987 100644
--- a/dns/enhancer.go
+++ b/dns/enhancer.go
@@ -84,6 +84,13 @@ func (h *ResolverEnhancer) InsertHostByIP(ip net.IP, host string) {
 	}
 }
 
+func (h *ResolverEnhancer) FlushFakeIP() error {
+	if h.fakePool != nil {
+		return h.fakePool.FlushFakeIP()
+	}
+	return nil
+}
+
 func (h *ResolverEnhancer) PatchFrom(o *ResolverEnhancer) {
 	if h.mapping != nil && o.mapping != nil {
 		o.mapping.CloneTo(h.mapping)
@@ -94,11 +101,10 @@ func (h *ResolverEnhancer) PatchFrom(o *ResolverEnhancer) {
 	}
 }
 
-func (h *ResolverEnhancer) FlushFakeIP() error {
+func (h *ResolverEnhancer) StoreFakePoolSate() {
 	if h.fakePool != nil {
-		return h.fakePool.FlushFakeIP()
+		h.fakePool.StoreState()
 	}
-	return nil
 }
 
 func NewEnhancer(cfg Config) *ResolverEnhancer {
diff --git a/hub/executor/executor.go b/hub/executor/executor.go
index 01222ce1..debf173c 100644
--- a/hub/executor/executor.go
+++ b/hub/executor/executor.go
@@ -395,6 +395,7 @@ func updateIPTables(cfg *config.Config) {
 func Shutdown() {
 	P.Cleanup()
 	tproxy.CleanupTProxyIPTables()
+	resolver.StoreFakePoolSate()
 
 	log.Warnln("Clash shutting down")
 }