From 075d8ed09417198913866a2820c0ee10f200258e Mon Sep 17 00:00:00 2001
From: Dreamacro <8615343+Dreamacro@users.noreply.github.com>
Date: Tue, 23 Nov 2021 22:01:49 +0800
Subject: [PATCH] Fix: fakeip pool cycle used

---
 component/fakeip/cachefile.go        |  6 ++++++
 component/fakeip/memory.go           |  9 +++++++++
 component/fakeip/pool.go             |  4 ++++
 component/fakeip/pool_test.go        | 15 +++++++++++----
 component/profile/cachefile/cache.go | 25 +++++++++++++++++++++++++
 5 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/component/fakeip/cachefile.go b/component/fakeip/cachefile.go
index 7ee88981..04bdc65c 100644
--- a/component/fakeip/cachefile.go
+++ b/component/fakeip/cachefile.go
@@ -38,6 +38,12 @@ func (c *cachefileStore) PutByIP(ip net.IP, host string) {
 	c.cache.PutFakeip(ip.To4(), []byte(host))
 }
 
+// DelByIP implements store.DelByIP
+func (c *cachefileStore) DelByIP(ip net.IP) {
+	ip = ip.To4()
+	c.cache.DelFakeipPair(ip, c.cache.GetFakeip(ip.To4()))
+}
+
 // Exist implements store.Exist
 func (c *cachefileStore) Exist(ip net.IP) bool {
 	_, exist := c.GetByIP(ip)
diff --git a/component/fakeip/memory.go b/component/fakeip/memory.go
index 75d4a3b2..c6c6873d 100644
--- a/component/fakeip/memory.go
+++ b/component/fakeip/memory.go
@@ -46,6 +46,15 @@ func (m *memoryStore) PutByIP(ip net.IP, host string) {
 	m.cache.Set(ipToUint(ip.To4()), host)
 }
 
+// DelByIP implements store.DelByIP
+func (m *memoryStore) DelByIP(ip net.IP) {
+	ipNum := ipToUint(ip.To4())
+	if elm, exist := m.cache.Get(ipNum); exist {
+		m.cache.Delete(elm.(string))
+	}
+	m.cache.Delete(ipNum)
+}
+
 // Exist implements store.Exist
 func (m *memoryStore) Exist(ip net.IP) bool {
 	return m.cache.Exist(ipToUint(ip.To4()))
diff --git a/component/fakeip/pool.go b/component/fakeip/pool.go
index 180d5eb1..925882bb 100644
--- a/component/fakeip/pool.go
+++ b/component/fakeip/pool.go
@@ -15,6 +15,7 @@ type store interface {
 	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
 	CloneTo(store)
 }
@@ -97,6 +98,9 @@ func (p *Pool) get(host string) net.IP {
 		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
 		}
 
diff --git a/component/fakeip/pool_test.go b/component/fakeip/pool_test.go
index bd4af4c9..bd034636 100644
--- a/component/fakeip/pool_test.go
+++ b/component/fakeip/pool_test.go
@@ -1,6 +1,7 @@
 package fakeip
 
 import (
+	"fmt"
 	"net"
 	"os"
 	"testing"
@@ -75,7 +76,7 @@ func TestPool_Basic(t *testing.T) {
 }
 
 func TestPool_CycleUsed(t *testing.T) {
-	_, ipnet, _ := net.ParseCIDR("192.168.0.1/30")
+	_, ipnet, _ := net.ParseCIDR("192.168.0.1/29")
 	pools, tempfile, err := createPools(Options{
 		IPNet: ipnet,
 		Size:  10,
@@ -84,9 +85,15 @@ func TestPool_CycleUsed(t *testing.T) {
 	defer os.Remove(tempfile)
 
 	for _, pool := range pools {
-		first := pool.Lookup("foo.com")
-		same := pool.Lookup("baz.com")
-		assert.True(t, first.Equal(same))
+		foo := pool.Lookup("foo.com")
+		bar := pool.Lookup("bar.com")
+		for i := 0; i < 3; i++ {
+			pool.Lookup(fmt.Sprintf("%d.com", i))
+		}
+		baz := pool.Lookup("baz.com")
+		next := pool.Lookup("foo.com")
+		assert.True(t, foo.Equal(baz))
+		assert.True(t, next.Equal(bar))
 	}
 }
 
diff --git a/component/profile/cachefile/cache.go b/component/profile/cachefile/cache.go
index 1ef47493..214e729b 100644
--- a/component/profile/cachefile/cache.go
+++ b/component/profile/cachefile/cache.go
@@ -90,6 +90,31 @@ func (c *CacheFile) PutFakeip(key, value []byte) error {
 	return err
 }
 
+func (c *CacheFile) DelFakeipPair(ip, host []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
+		}
+		err = bucket.Delete(ip)
+		if len(host) > 0 {
+			if err := bucket.Delete(host); err != nil {
+				return err
+			}
+		}
+		return err
+	})
+	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