Skip to content

Commit

Permalink
chore: save etag in bbolt by msgpack
Browse files Browse the repository at this point in the history
  • Loading branch information
wwqgtxx committed Sep 23, 2024
1 parent 966eeae commit 59a2b24
Show file tree
Hide file tree
Showing 6 changed files with 97 additions and 66 deletions.
33 changes: 25 additions & 8 deletions common/utils/hash.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package utils
import (
"crypto/md5"
"encoding/hex"
"errors"
)

// HashType warps hash array inside struct
Expand All @@ -15,14 +16,6 @@ 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
}
Expand All @@ -35,6 +28,30 @@ func (h HashType) String() string {
return hex.EncodeToString(h.Bytes())
}

func (h HashType) MarshalText() ([]byte, error) {
return []byte(h.String()), nil
}

func (h *HashType) UnmarshalText(data []byte) error {
if hex.DecodedLen(len(data)) != md5.Size {
return errors.New("invalid hash length")
}
_, err := hex.Decode(h.md5[:], data)
return err
}

func (h HashType) MarshalBinary() ([]byte, error) {
return h.md5[:], nil
}

func (h *HashType) UnmarshalBinary(data []byte) error {
if len(data) != md5.Size {
return errors.New("invalid hash length")
}
copy(h.md5[:], data)
return nil
}

func (h HashType) Len() int {
return len(h.md5)
}
Expand Down
54 changes: 0 additions & 54 deletions component/profile/cachefile/cache.go
Original file line number Diff line number Diff line change
@@ -1,12 +1,10 @@
package cachefile

import (
"math"
"os"
"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 @@ -72,58 +70,6 @@ func (c *CacheFile) SelectedMap() map[string]string {
return mapping
}

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

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.Bytes()...)
data = append(data, etag...)

err := c.DB.Batch(func(t *bbolt.Tx) error {
bucket, err := t.CreateBucketIfNotExists(bucketETag)
if err != nil {
return err
}

return bucket.Put([]byte(url), data)
})
if err != nil {
log.Warnln("[CacheFile] write cache to %s failed: %s", c.DB.Path(), err.Error())
return
}
}
func (c *CacheFile) GetETagWithHash(key string) (hash utils.HashType, etag string) {
if c.DB == nil {
return
}
c.DB.View(func(t *bbolt.Tx) error {
if bucket := t.Bucket(bucketETag); bucket != nil {
if v := bucket.Get([]byte(key)); v != nil {
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
})

return
}

func (c *CacheFile) Close() error {
return c.DB.Close()
}
Expand Down
58 changes: 58 additions & 0 deletions component/profile/cachefile/etag.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,58 @@
package cachefile

import (
"time"

"github.com/metacubex/mihomo/common/utils"
"github.com/metacubex/mihomo/log"

"github.com/metacubex/bbolt"
"github.com/vmihailenco/msgpack/v5"
)

type EtagWithHash struct {
Hash utils.HashType
ETag string
Time time.Time
}

func (c *CacheFile) SetETagWithHash(url string, etagWithHash EtagWithHash) {
if c.DB == nil {
return
}

data, err := msgpack.Marshal(etagWithHash)
if err != nil {
return // maybe panic is better
}

err = c.DB.Batch(func(t *bbolt.Tx) error {
bucket, err := t.CreateBucketIfNotExists(bucketETag)
if err != nil {
return err
}

return bucket.Put([]byte(url), data)
})
if err != nil {
log.Warnln("[CacheFile] write cache to %s failed: %s", c.DB.Path(), err.Error())
return
}
}
func (c *CacheFile) GetETagWithHash(key string) (etagWithHash EtagWithHash) {
if c.DB == nil {
return
}
c.DB.View(func(t *bbolt.Tx) error {
if bucket := t.Bucket(bucketETag); bucket != nil {
if v := bucket.Get([]byte(key)); v != nil {
if err := msgpack.Unmarshal(v, &etagWithHash); err != nil {
return err
}
}
}
return nil
})

return
}
12 changes: 8 additions & 4 deletions component/resource/vehicle.go
Original file line number Diff line number Diff line change
Expand Up @@ -117,14 +117,14 @@ func (h *HTTPVehicle) Read(ctx context.Context, oldHash utils.HashType) (buf []b
header := h.header
setIfNoneMatch := false
if etag && oldHash.IsValid() {
hashBytes, etag := cachefile.Cache().GetETagWithHash(h.url)
if oldHash.Equal(hashBytes) && etag != "" {
etagWithHash := cachefile.Cache().GetETagWithHash(h.url)
if oldHash.Equal(etagWithHash.Hash) && etagWithHash.ETag != "" {
if header == nil {
header = http.Header{}
} else {
header = header.Clone()
}
header.Set("If-None-Match", etag)
header.Set("If-None-Match", etagWithHash.ETag)
setIfNoneMatch = true
}
}
Expand All @@ -146,7 +146,11 @@ func (h *HTTPVehicle) Read(ctx context.Context, oldHash utils.HashType) (buf []b
}
hash = utils.MakeHash(buf)
if etag {
cachefile.Cache().SetETagWithHash(h.url, hash, resp.Header.Get("ETag"))
cachefile.Cache().SetETagWithHash(h.url, cachefile.EtagWithHash{
Hash: hash,
ETag: resp.Header.Get("ETag"),
Time: time.Now(),
})
}
return
}
Expand Down
2 changes: 2 additions & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ require (
github.com/shirou/gopsutil/v3 v3.24.5
github.com/sirupsen/logrus v1.9.3
github.com/stretchr/testify v1.9.0
github.com/vmihailenco/msgpack/v5 v5.4.1
github.com/wk8/go-ordered-map/v2 v2.1.8
gitlab.com/go-extension/aes-ccm v0.0.0-20230221065045-e58665ef23c7
go.uber.org/automaxprocs v1.5.3
Expand Down Expand Up @@ -104,6 +105,7 @@ require (
github.com/tklauser/numcpus v0.6.1 // indirect
github.com/u-root/uio v0.0.0-20230220225925-ffce2a382923 // indirect
github.com/vishvananda/netns v0.0.4 // indirect
github.com/vmihailenco/tagparser/v2 v2.0.0 // indirect
github.com/yusufpapurcu/wmi v1.2.4 // indirect
gitlab.com/yawning/bsaes.git v0.0.0-20190805113838-0a714cd429ec // indirect
go.uber.org/mock v0.4.0 // indirect
Expand Down
4 changes: 4 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -210,6 +210,10 @@ github.com/vishvananda/netns v0.0.0-20191106174202-0a2b9b5464df/go.mod h1:JP3t17
github.com/vishvananda/netns v0.0.0-20210104183010-2eb08e3e575f/go.mod h1:DD4vA1DwXk04H54A1oHXtwZmA0grkVMdPxx/VGLCah0=
github.com/vishvananda/netns v0.0.4 h1:Oeaw1EM2JMxD51g9uhtC0D7erkIjgmj8+JZc26m1YX8=
github.com/vishvananda/netns v0.0.4/go.mod h1:SpkAiCQRtJ6TvvxPnOSyH3BMl6unz3xZlaprSwhNNJM=
github.com/vmihailenco/msgpack/v5 v5.4.1 h1:cQriyiUvjTwOHg8QZaPihLWeRAAVoCpE00IUPn0Bjt8=
github.com/vmihailenco/msgpack/v5 v5.4.1/go.mod h1:GaZTsDaehaPpQVyxrf5mtQlH+pc21PIudVV/E3rRQok=
github.com/vmihailenco/tagparser/v2 v2.0.0 h1:y09buUbR+b5aycVFQs/g70pqKVZNBmxwAhO7/IwNM9g=
github.com/vmihailenco/tagparser/v2 v2.0.0/go.mod h1:Wri+At7QHww0WTrCBeu4J6bNtoV6mEfg5OIWRZA9qds=
github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc=
github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw=
github.com/yusufpapurcu/wmi v1.2.4 h1:zFUKzehAFReQwLys1b/iSMl+JQGSCSjtVqQn9bBrPo0=
Expand Down

0 comments on commit 59a2b24

Please sign in to comment.