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))
}
// 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)

View file

@ -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()))

View file

@ -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
}

View file

@ -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))
}
}

View file

@ -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