From 52125a3992ec2c3dd18afd33b4689b82680fd2ea Mon Sep 17 00:00:00 2001
From: Dreamacro <305009791@qq.com>
Date: Fri, 11 Oct 2019 14:01:16 +0800
Subject: [PATCH] Fix: fakeip missing host

---
 common/cache/lrucache.go      |  9 +++++++++
 component/fakeip/pool.go      | 20 +++++++++++++++-----
 component/fakeip/pool_test.go | 27 +++++++++++++++++++++++++++
 3 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/common/cache/lrucache.go b/common/cache/lrucache.go
index 5a139bf7..1cb278d9 100644
--- a/common/cache/lrucache.go
+++ b/common/cache/lrucache.go
@@ -86,6 +86,15 @@ func (c *LruCache) Get(key interface{}) (interface{}, bool) {
 	return value, true
 }
 
+// Exist returns if key exist in cache but not put item to the head of linked list
+func (c *LruCache) Exist(key interface{}) bool {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	_, ok := c.cache[key]
+	return ok
+}
+
 // Set stores the interface{} representation of a response for a given key.
 func (c *LruCache) Set(key interface{}, value interface{}) {
 	c.mu.Lock()
diff --git a/component/fakeip/pool.go b/component/fakeip/pool.go
index 8b7a2766..e71ddfaa 100644
--- a/component/fakeip/pool.go
+++ b/component/fakeip/pool.go
@@ -22,8 +22,14 @@ type Pool struct {
 func (p *Pool) Lookup(host string) net.IP {
 	p.mux.Lock()
 	defer p.mux.Unlock()
-	if ip, exist := p.cache.Get(host); exist {
-		return ip.(net.IP)
+	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)
+		return ip
 	}
 
 	ip := p.get(host)
@@ -43,8 +49,12 @@ func (p *Pool) LookBack(ip net.IP) (string, bool) {
 	n := ipToUint(ip.To4())
 	offset := n - p.min + 1
 
-	if host, exist := p.cache.Get(offset); exist {
-		return host.(string), true
+	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
@@ -64,7 +74,7 @@ func (p *Pool) get(host string) net.IP {
 			break
 		}
 
-		if _, exist := p.cache.Get(p.offset); !exist {
+		if !p.cache.Exist(p.offset) {
 			break
 		}
 	}
diff --git a/component/fakeip/pool_test.go b/component/fakeip/pool_test.go
index fd62349e..d68c1011 100644
--- a/component/fakeip/pool_test.go
+++ b/component/fakeip/pool_test.go
@@ -43,6 +43,33 @@ func TestPool_MaxCacheSize(t *testing.T) {
 	assert.False(t, first.Equal(next))
 }
 
+func TestPool_DoubleMapping(t *testing.T) {
+	_, ipnet, _ := net.ParseCIDR("192.168.0.1/24")
+	pool, _ := New(ipnet, 2)
+
+	// fill cache
+	fooIP := pool.Lookup("foo.com")
+	bazIP := pool.Lookup("baz.com")
+
+	// make foo.com hot
+	pool.Lookup("foo.com")
+
+	// should drop baz.com
+	barIP := pool.Lookup("bar.com")
+
+	_, fooExist := pool.LookBack(fooIP)
+	_, bazExist := pool.LookBack(bazIP)
+	_, barExist := pool.LookBack(barIP)
+
+	newBazIP := pool.Lookup("baz.com")
+
+	assert.True(t, fooExist)
+	assert.False(t, bazExist)
+	assert.True(t, barExist)
+
+	assert.False(t, bazIP.Equal(newBazIP))
+}
+
 func TestPool_Error(t *testing.T) {
 	_, ipnet, _ := net.ParseCIDR("192.168.0.1/31")
 	_, err := New(ipnet, 10)