Skip to content

Commit 966eeae

Browse files
wwqgtxxwwqgtxx
committed
chore: rewrite bbolt cachefile implements
never use returned byte slices outside the transaction, ref: https://pkg.go.dev/go.etcd.io/bbolt#hdr-Caveats
1 parent 150c6cc commit 966eeae

File tree

12 files changed

+212
-168
lines changed

12 files changed

+212
-168
lines changed

‎common/utils/hash.go

+45
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
package utils
2+
3+
import (
4+
"crypto/md5"
5+
"encoding/hex"
6+
)
7+
8+
// HashType warps hash array inside struct
9+
// someday can change to other hash algorithm simply
10+
type HashType struct {
11+
md5 [md5.Size]byte // MD5
12+
}
13+
14+
func MakeHash(data []byte) HashType {
15+
return HashType{md5.Sum(data)}
16+
}
17+
18+
func MakeHashFromBytes(hashBytes []byte) (h HashType) {
19+
if len(hashBytes) != md5.Size {
20+
return
21+
}
22+
copy(h.md5[:], hashBytes)
23+
return
24+
}
25+
26+
func (h HashType) Equal(hash HashType) bool {
27+
return h.md5 == hash.md5
28+
}
29+
30+
func (h HashType) Bytes() []byte {
31+
return h.md5[:]
32+
}
33+
34+
func (h HashType) String() string {
35+
return hex.EncodeToString(h.Bytes())
36+
}
37+
38+
func (h HashType) Len() int {
39+
return len(h.md5)
40+
}
41+
42+
func (h HashType) IsValid() bool {
43+
var zero HashType
44+
return h != zero
45+
}

‎component/fakeip/cachefile.go

+10-20
Original file line numberDiff line numberDiff line change
@@ -7,46 +7,32 @@ import (
77
)
88

99
type cachefileStore struct {
10-
cache *cachefile.CacheFile
10+
cache *cachefile.FakeIpStore
1111
}
1212

1313
// GetByHost implements store.GetByHost
1414
func (c *cachefileStore) GetByHost(host string) (netip.Addr, bool) {
15-
elm := c.cache.GetFakeip([]byte(host))
16-
if elm == nil {
17-
return netip.Addr{}, false
18-
}
19-
20-
if len(elm) == 4 {
21-
return netip.AddrFrom4(*(*[4]byte)(elm)), true
22-
} else {
23-
return netip.AddrFrom16(*(*[16]byte)(elm)), true
24-
}
15+
return c.cache.GetByHost(host)
2516
}
2617

2718
// PutByHost implements store.PutByHost
2819
func (c *cachefileStore) PutByHost(host string, ip netip.Addr) {
29-
c.cache.PutFakeip([]byte(host), ip.AsSlice())
20+
c.cache.PutByHost(host, ip)
3021
}
3122

3223
// GetByIP implements store.GetByIP
3324
func (c *cachefileStore) GetByIP(ip netip.Addr) (string, bool) {
34-
elm := c.cache.GetFakeip(ip.AsSlice())
35-
if elm == nil {
36-
return "", false
37-
}
38-
return string(elm), true
25+
return c.cache.GetByIP(ip)
3926
}
4027

4128
// PutByIP implements store.PutByIP
4229
func (c *cachefileStore) PutByIP(ip netip.Addr, host string) {
43-
c.cache.PutFakeip(ip.AsSlice(), []byte(host))
30+
c.cache.PutByIP(ip, host)
4431
}
4532

4633
// DelByIP implements store.DelByIP
4734
func (c *cachefileStore) DelByIP(ip netip.Addr) {
48-
addr := ip.AsSlice()
49-
c.cache.DelFakeipPair(addr, c.cache.GetFakeip(addr))
35+
c.cache.DelByIP(ip)
5036
}
5137

5238
// Exist implements store.Exist
@@ -63,3 +49,7 @@ func (c *cachefileStore) CloneTo(store store) {}
6349
func (c *cachefileStore) FlushFakeIP() error {
6450
return c.cache.FlushFakeIP()
6551
}
52+
53+
func newCachefileStore(cache *cachefile.CacheFile) *cachefileStore {
54+
return &cachefileStore{cache.FakeIpStore()}
55+
}

‎component/fakeip/pool.go

+1-3
Original file line numberDiff line numberDiff line change
@@ -201,9 +201,7 @@ func New(options Options) (*Pool, error) {
201201
ipnet: options.IPNet,
202202
}
203203
if options.Persistence {
204-
pool.store = &cachefileStore{
205-
cache: cachefile.Cache(),
206-
}
204+
pool.store = newCachefileStore(cachefile.Cache())
207205
} else {
208206
pool.store = newMemoryStore(options.Size)
209207
}

‎component/fakeip/pool_test.go

+1-3
Original file line numberDiff line numberDiff line change
@@ -43,9 +43,7 @@ func createCachefileStore(options Options) (*Pool, string, error) {
4343
return nil, "", err
4444
}
4545

46-
pool.store = &cachefileStore{
47-
cache: &cachefile.CacheFile{DB: db},
48-
}
46+
pool.store = newCachefileStore(&cachefile.CacheFile{DB: db})
4947
return pool, f.Name(), nil
5048
}
5149

‎component/profile/cachefile/cache.go

+15-89
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"sync"
77
"time"
88

9+
"github.com/metacubex/mihomo/common/utils"
910
"github.com/metacubex/mihomo/component/profile"
1011
C "github.com/metacubex/mihomo/constant"
1112
"github.com/metacubex/mihomo/log"
@@ -71,93 +72,19 @@ func (c *CacheFile) SelectedMap() map[string]string {
7172
return mapping
7273
}
7374

74-
func (c *CacheFile) PutFakeip(key, value []byte) error {
75-
if c.DB == nil {
76-
return nil
77-
}
78-
79-
err := c.DB.Batch(func(t *bbolt.Tx) error {
80-
bucket, err := t.CreateBucketIfNotExists(bucketFakeip)
81-
if err != nil {
82-
return err
83-
}
84-
return bucket.Put(key, value)
85-
})
86-
if err != nil {
87-
log.Warnln("[CacheFile] write cache to %s failed: %s", c.DB.Path(), err.Error())
88-
}
89-
90-
return err
91-
}
92-
93-
func (c *CacheFile) DelFakeipPair(ip, host []byte) error {
94-
if c.DB == nil {
95-
return nil
96-
}
97-
98-
err := c.DB.Batch(func(t *bbolt.Tx) error {
99-
bucket, err := t.CreateBucketIfNotExists(bucketFakeip)
100-
if err != nil {
101-
return err
102-
}
103-
err = bucket.Delete(ip)
104-
if len(host) > 0 {
105-
if err := bucket.Delete(host); err != nil {
106-
return err
107-
}
108-
}
109-
return err
110-
})
111-
if err != nil {
112-
log.Warnln("[CacheFile] write cache to %s failed: %s", c.DB.Path(), err.Error())
113-
}
114-
115-
return err
116-
}
117-
118-
func (c *CacheFile) GetFakeip(key []byte) []byte {
119-
if c.DB == nil {
120-
return nil
121-
}
122-
123-
tx, err := c.DB.Begin(false)
124-
if err != nil {
125-
return nil
126-
}
127-
defer tx.Rollback()
128-
129-
bucket := tx.Bucket(bucketFakeip)
130-
if bucket == nil {
131-
return nil
132-
}
133-
134-
return bucket.Get(key)
135-
}
136-
137-
func (c *CacheFile) FlushFakeIP() error {
138-
err := c.DB.Batch(func(t *bbolt.Tx) error {
139-
bucket := t.Bucket(bucketFakeip)
140-
if bucket == nil {
141-
return nil
142-
}
143-
return t.DeleteBucket(bucketFakeip)
144-
})
145-
return err
146-
}
147-
148-
func (c *CacheFile) SetETagWithHash(url string, hash []byte, etag string) {
75+
func (c *CacheFile) SetETagWithHash(url string, hash utils.HashType, etag string) {
14976
if c.DB == nil {
15077
return
15178
}
15279

153-
lenHash := len(hash)
80+
lenHash := hash.Len()
15481
if lenHash > math.MaxUint8 {
15582
return // maybe panic is better
15683
}
15784

15885
data := make([]byte, 1, 1+lenHash+len(etag))
15986
data[0] = uint8(lenHash)
160-
data = append(data, hash...)
87+
data = append(data, hash.Bytes()...)
16188
data = append(data, etag...)
16289

16390
err := c.DB.Batch(func(t *bbolt.Tx) error {
@@ -173,28 +100,27 @@ func (c *CacheFile) SetETagWithHash(url string, hash []byte, etag string) {
173100
return
174101
}
175102
}
176-
func (c *CacheFile) GetETagWithHash(key string) (hash []byte, etag string) {
103+
func (c *CacheFile) GetETagWithHash(key string) (hash utils.HashType, etag string) {
177104
if c.DB == nil {
178105
return
179106
}
180-
var value []byte
181107
c.DB.View(func(t *bbolt.Tx) error {
182108
if bucket := t.Bucket(bucketETag); bucket != nil {
183109
if v := bucket.Get([]byte(key)); v != nil {
184-
value = v
110+
if len(v) == 0 {
111+
return nil
112+
}
113+
lenHash := int(v[0])
114+
if len(v) < 1+lenHash {
115+
return nil
116+
}
117+
hash = utils.MakeHashFromBytes(v[1 : 1+lenHash])
118+
etag = string(v[1+lenHash:])
185119
}
186120
}
187121
return nil
188122
})
189-
if len(value) == 0 {
190-
return
191-
}
192-
lenHash := int(value[0])
193-
if len(value) < 1+lenHash {
194-
return
195-
}
196-
hash = value[1 : 1+lenHash]
197-
etag = string(value[1+lenHash:])
123+
198124
return
199125
}
200126

‎component/profile/cachefile/fakeip.go

+115
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
package cachefile
2+
3+
import (
4+
"net/netip"
5+
6+
"github.com/metacubex/mihomo/log"
7+
8+
"github.com/metacubex/bbolt"
9+
)
10+
11+
type FakeIpStore struct {
12+
*CacheFile
13+
}
14+
15+
func (c *CacheFile) FakeIpStore() *FakeIpStore {
16+
return &FakeIpStore{c}
17+
}
18+
19+
func (c *FakeIpStore) GetByHost(host string) (ip netip.Addr, exist bool) {
20+
if c.DB == nil {
21+
return
22+
}
23+
c.DB.View(func(t *bbolt.Tx) error {
24+
if bucket := t.Bucket(bucketFakeip); bucket != nil {
25+
if v := bucket.Get([]byte(host)); v != nil {
26+
ip, exist = netip.AddrFromSlice(v)
27+
}
28+
}
29+
return nil
30+
})
31+
return
32+
}
33+
34+
func (c *FakeIpStore) PutByHost(host string, ip netip.Addr) {
35+
if c.DB == nil {
36+
return
37+
}
38+
err := c.DB.Batch(func(t *bbolt.Tx) error {
39+
bucket, err := t.CreateBucketIfNotExists(bucketFakeip)
40+
if err != nil {
41+
return err
42+
}
43+
return bucket.Put([]byte(host), ip.AsSlice())
44+
})
45+
if err != nil {
46+
log.Warnln("[CacheFile] write cache to %s failed: %s", c.DB.Path(), err.Error())
47+
}
48+
}
49+
50+
func (c *FakeIpStore) GetByIP(ip netip.Addr) (host string, exist bool) {
51+
if c.DB == nil {
52+
return
53+
}
54+
c.DB.View(func(t *bbolt.Tx) error {
55+
if bucket := t.Bucket(bucketFakeip); bucket != nil {
56+
if v := bucket.Get(ip.AsSlice()); v != nil {
57+
host, exist = string(v), true
58+
}
59+
}
60+
return nil
61+
})
62+
return
63+
}
64+
65+
func (c *FakeIpStore) PutByIP(ip netip.Addr, host string) {
66+
if c.DB == nil {
67+
return
68+
}
69+
err := c.DB.Batch(func(t *bbolt.Tx) error {
70+
bucket, err := t.CreateBucketIfNotExists(bucketFakeip)
71+
if err != nil {
72+
return err
73+
}
74+
return bucket.Put(ip.AsSlice(), []byte(host))
75+
})
76+
if err != nil {
77+
log.Warnln("[CacheFile] write cache to %s failed: %s", c.DB.Path(), err.Error())
78+
}
79+
}
80+
81+
func (c *FakeIpStore) DelByIP(ip netip.Addr) {
82+
if c.DB == nil {
83+
return
84+
}
85+
86+
addr := ip.AsSlice()
87+
err := c.DB.Batch(func(t *bbolt.Tx) error {
88+
bucket, err := t.CreateBucketIfNotExists(bucketFakeip)
89+
if err != nil {
90+
return err
91+
}
92+
host := bucket.Get(addr)
93+
err = bucket.Delete(addr)
94+
if len(host) > 0 {
95+
if err = bucket.Delete(host); err != nil {
96+
return err
97+
}
98+
}
99+
return err
100+
})
101+
if err != nil {
102+
log.Warnln("[CacheFile] write cache to %s failed: %s", c.DB.Path(), err.Error())
103+
}
104+
}
105+
106+
func (c *FakeIpStore) FlushFakeIP() error {
107+
err := c.DB.Batch(func(t *bbolt.Tx) error {
108+
bucket := t.Bucket(bucketFakeip)
109+
if bucket == nil {
110+
return nil
111+
}
112+
return t.DeleteBucket(bucketFakeip)
113+
})
114+
return err
115+
}

0 commit comments

Comments
 (0)