Skip to content

Commit

Permalink
chore: rewrite bbolt cachefile implements
Browse files Browse the repository at this point in the history
never use returned byte slices outside the transaction, ref:
https://pkg.go.dev/go.etcd.io/bbolt#hdr-Caveats
  • Loading branch information
wwqgtxx committed Sep 23, 2024
1 parent 150c6cc commit 966eeae
Show file tree
Hide file tree
Showing 12 changed files with 212 additions and 168 deletions.
45 changes: 45 additions & 0 deletions common/utils/hash.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
package utils

import (
"crypto/md5"
"encoding/hex"
)

// HashType warps hash array inside struct
// someday can change to other hash algorithm simply
type HashType struct {
md5 [md5.Size]byte // MD5
}

func MakeHash(data []byte) HashType {
return HashType{md5.Sum(data)}
}

func MakeHashFromBytes(hashBytes []byte) (h HashType) {
if len(hashBytes) != md5.Size {
return
}
copy(h.md5[:], hashBytes)
return
}

func (h HashType) Equal(hash HashType) bool {
return h.md5 == hash.md5
}

func (h HashType) Bytes() []byte {
return h.md5[:]
}

func (h HashType) String() string {
return hex.EncodeToString(h.Bytes())
}

func (h HashType) Len() int {
return len(h.md5)
}

func (h HashType) IsValid() bool {
var zero HashType
return h != zero
}
30 changes: 10 additions & 20 deletions component/fakeip/cachefile.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,46 +7,32 @@ import (
)

type cachefileStore struct {
cache *cachefile.CacheFile
cache *cachefile.FakeIpStore
}

// GetByHost implements store.GetByHost
func (c *cachefileStore) GetByHost(host string) (netip.Addr, bool) {
elm := c.cache.GetFakeip([]byte(host))
if elm == nil {
return netip.Addr{}, false
}

if len(elm) == 4 {
return netip.AddrFrom4(*(*[4]byte)(elm)), true
} else {
return netip.AddrFrom16(*(*[16]byte)(elm)), true
}
return c.cache.GetByHost(host)
}

// PutByHost implements store.PutByHost
func (c *cachefileStore) PutByHost(host string, ip netip.Addr) {
c.cache.PutFakeip([]byte(host), ip.AsSlice())
c.cache.PutByHost(host, ip)
}

// GetByIP implements store.GetByIP
func (c *cachefileStore) GetByIP(ip netip.Addr) (string, bool) {
elm := c.cache.GetFakeip(ip.AsSlice())
if elm == nil {
return "", false
}
return string(elm), true
return c.cache.GetByIP(ip)
}

// PutByIP implements store.PutByIP
func (c *cachefileStore) PutByIP(ip netip.Addr, host string) {
c.cache.PutFakeip(ip.AsSlice(), []byte(host))
c.cache.PutByIP(ip, host)
}

// DelByIP implements store.DelByIP
func (c *cachefileStore) DelByIP(ip netip.Addr) {
addr := ip.AsSlice()
c.cache.DelFakeipPair(addr, c.cache.GetFakeip(addr))
c.cache.DelByIP(ip)
}

// Exist implements store.Exist
Expand All @@ -63,3 +49,7 @@ func (c *cachefileStore) CloneTo(store store) {}
func (c *cachefileStore) FlushFakeIP() error {
return c.cache.FlushFakeIP()
}

func newCachefileStore(cache *cachefile.CacheFile) *cachefileStore {
return &cachefileStore{cache.FakeIpStore()}
}
4 changes: 1 addition & 3 deletions component/fakeip/pool.go
Original file line number Diff line number Diff line change
Expand Up @@ -201,9 +201,7 @@ func New(options Options) (*Pool, error) {
ipnet: options.IPNet,
}
if options.Persistence {
pool.store = &cachefileStore{
cache: cachefile.Cache(),
}
pool.store = newCachefileStore(cachefile.Cache())
} else {
pool.store = newMemoryStore(options.Size)
}
Expand Down
4 changes: 1 addition & 3 deletions component/fakeip/pool_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -43,9 +43,7 @@ func createCachefileStore(options Options) (*Pool, string, error) {
return nil, "", err
}

pool.store = &cachefileStore{
cache: &cachefile.CacheFile{DB: db},
}
pool.store = newCachefileStore(&cachefile.CacheFile{DB: db})
return pool, f.Name(), nil
}

Expand Down
104 changes: 15 additions & 89 deletions component/profile/cachefile/cache.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"sync"
"time"

"github.com/metacubex/mihomo/common/utils"
"github.com/metacubex/mihomo/component/profile"
C "github.com/metacubex/mihomo/constant"
"github.com/metacubex/mihomo/log"
Expand Down Expand Up @@ -71,93 +72,19 @@ func (c *CacheFile) SelectedMap() map[string]string {
return mapping
}

func (c *CacheFile) PutFakeip(key, value []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
}
return bucket.Put(key, value)
})
if err != nil {
log.Warnln("[CacheFile] write cache to %s failed: %s", c.DB.Path(), err.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
}

tx, err := c.DB.Begin(false)
if err != nil {
return nil
}
defer tx.Rollback()

bucket := tx.Bucket(bucketFakeip)
if bucket == nil {
return nil
}

return bucket.Get(key)
}

func (c *CacheFile) FlushFakeIP() error {
err := c.DB.Batch(func(t *bbolt.Tx) error {
bucket := t.Bucket(bucketFakeip)
if bucket == nil {
return nil
}
return t.DeleteBucket(bucketFakeip)
})
return err
}

func (c *CacheFile) SetETagWithHash(url string, hash []byte, etag string) {
func (c *CacheFile) SetETagWithHash(url string, hash utils.HashType, etag string) {
if c.DB == nil {
return
}

lenHash := len(hash)
lenHash := hash.Len()
if lenHash > math.MaxUint8 {
return // maybe panic is better
}

data := make([]byte, 1, 1+lenHash+len(etag))
data[0] = uint8(lenHash)
data = append(data, hash...)
data = append(data, hash.Bytes()...)
data = append(data, etag...)

err := c.DB.Batch(func(t *bbolt.Tx) error {
Expand All @@ -173,28 +100,27 @@ func (c *CacheFile) SetETagWithHash(url string, hash []byte, etag string) {
return
}
}
func (c *CacheFile) GetETagWithHash(key string) (hash []byte, etag string) {
func (c *CacheFile) GetETagWithHash(key string) (hash utils.HashType, etag string) {
if c.DB == nil {
return
}
var value []byte
c.DB.View(func(t *bbolt.Tx) error {
if bucket := t.Bucket(bucketETag); bucket != nil {
if v := bucket.Get([]byte(key)); v != nil {
value = v
if len(v) == 0 {
return nil
}
lenHash := int(v[0])
if len(v) < 1+lenHash {
return nil
}
hash = utils.MakeHashFromBytes(v[1 : 1+lenHash])
etag = string(v[1+lenHash:])
}
}
return nil
})
if len(value) == 0 {
return
}
lenHash := int(value[0])
if len(value) < 1+lenHash {
return
}
hash = value[1 : 1+lenHash]
etag = string(value[1+lenHash:])

return
}

Expand Down
115 changes: 115 additions & 0 deletions component/profile/cachefile/fakeip.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,115 @@
package cachefile

import (
"net/netip"

"github.com/metacubex/mihomo/log"

"github.com/metacubex/bbolt"
)

type FakeIpStore struct {
*CacheFile
}

func (c *CacheFile) FakeIpStore() *FakeIpStore {
return &FakeIpStore{c}
}

func (c *FakeIpStore) GetByHost(host string) (ip netip.Addr, exist bool) {
if c.DB == nil {
return
}
c.DB.View(func(t *bbolt.Tx) error {
if bucket := t.Bucket(bucketFakeip); bucket != nil {
if v := bucket.Get([]byte(host)); v != nil {
ip, exist = netip.AddrFromSlice(v)
}
}
return nil
})
return
}

func (c *FakeIpStore) PutByHost(host string, ip netip.Addr) {
if c.DB == nil {
return
}
err := c.DB.Batch(func(t *bbolt.Tx) error {
bucket, err := t.CreateBucketIfNotExists(bucketFakeip)
if err != nil {
return err
}
return bucket.Put([]byte(host), ip.AsSlice())
})
if err != nil {
log.Warnln("[CacheFile] write cache to %s failed: %s", c.DB.Path(), err.Error())
}
}

func (c *FakeIpStore) GetByIP(ip netip.Addr) (host string, exist bool) {
if c.DB == nil {
return
}
c.DB.View(func(t *bbolt.Tx) error {
if bucket := t.Bucket(bucketFakeip); bucket != nil {
if v := bucket.Get(ip.AsSlice()); v != nil {
host, exist = string(v), true
}
}
return nil
})
return
}

func (c *FakeIpStore) PutByIP(ip netip.Addr, host string) {
if c.DB == nil {
return
}
err := c.DB.Batch(func(t *bbolt.Tx) error {
bucket, err := t.CreateBucketIfNotExists(bucketFakeip)
if err != nil {
return err
}
return bucket.Put(ip.AsSlice(), []byte(host))
})
if err != nil {
log.Warnln("[CacheFile] write cache to %s failed: %s", c.DB.Path(), err.Error())
}
}

func (c *FakeIpStore) DelByIP(ip netip.Addr) {
if c.DB == nil {
return
}

addr := ip.AsSlice()
err := c.DB.Batch(func(t *bbolt.Tx) error {
bucket, err := t.CreateBucketIfNotExists(bucketFakeip)
if err != nil {
return err
}
host := bucket.Get(addr)
err = bucket.Delete(addr)
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())
}
}

func (c *FakeIpStore) FlushFakeIP() error {
err := c.DB.Batch(func(t *bbolt.Tx) error {
bucket := t.Bucket(bucketFakeip)
if bucket == nil {
return nil
}
return t.DeleteBucket(bucketFakeip)
})
return err
}
Loading

0 comments on commit 966eeae

Please sign in to comment.