Fix: fakeip pool cycle used

This commit is contained in:
Dreamacro 2021-11-23 22:01:49 +08:00
parent b1bed7623d
commit 075d8ed094
5 changed files with 55 additions and 4 deletions

View file

@ -38,6 +38,12 @@ func (c *cachefileStore) PutByIP(ip net.IP, host string) {
c.cache.PutFakeip(ip.To4(), []byte(host)) 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 // Exist implements store.Exist
func (c *cachefileStore) Exist(ip net.IP) bool { func (c *cachefileStore) Exist(ip net.IP) bool {
_, exist := c.GetByIP(ip) _, exist := c.GetByIP(ip)

View file

@ -46,6 +46,15 @@ func (m *memoryStore) PutByIP(ip net.IP, host string) {
m.cache.Set(ipToUint(ip.To4()), host) 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 // Exist implements store.Exist
func (m *memoryStore) Exist(ip net.IP) bool { func (m *memoryStore) Exist(ip net.IP) bool {
return m.cache.Exist(ipToUint(ip.To4())) return m.cache.Exist(ipToUint(ip.To4()))

View file

@ -15,6 +15,7 @@ type store interface {
PutByHost(host string, ip net.IP) PutByHost(host string, ip net.IP)
GetByIP(ip net.IP) (string, bool) GetByIP(ip net.IP) (string, bool)
PutByIP(ip net.IP, host string) PutByIP(ip net.IP, host string)
DelByIP(ip net.IP)
Exist(ip net.IP) bool Exist(ip net.IP) bool
CloneTo(store) CloneTo(store)
} }
@ -97,6 +98,9 @@ func (p *Pool) get(host string) net.IP {
p.offset = (p.offset + 1) % (p.max - p.min) p.offset = (p.offset + 1) % (p.max - p.min)
// Avoid infinite loops // Avoid infinite loops
if p.offset == current { 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 break
} }

View file

@ -1,6 +1,7 @@
package fakeip package fakeip
import ( import (
"fmt"
"net" "net"
"os" "os"
"testing" "testing"
@ -75,7 +76,7 @@ func TestPool_Basic(t *testing.T) {
} }
func TestPool_CycleUsed(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{ pools, tempfile, err := createPools(Options{
IPNet: ipnet, IPNet: ipnet,
Size: 10, Size: 10,
@ -84,9 +85,15 @@ func TestPool_CycleUsed(t *testing.T) {
defer os.Remove(tempfile) defer os.Remove(tempfile)
for _, pool := range pools { for _, pool := range pools {
first := pool.Lookup("foo.com") foo := pool.Lookup("foo.com")
same := pool.Lookup("baz.com") bar := pool.Lookup("bar.com")
assert.True(t, first.Equal(same)) 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))
} }
} }

View file

@ -90,6 +90,31 @@ func (c *CacheFile) PutFakeip(key, value []byte) error {
return err 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 { func (c *CacheFile) GetFakeip(key []byte) []byte {
if c.DB == nil { if c.DB == nil {
return nil return nil