mihomo/component/resource/fetcher.go
Larvan2 2f203330e4 feat: add override to proxy-providers
Co-authored-by: xishang0128 <xishang02@gmail.com>
2023-11-17 00:37:54 +08:00

214 lines
4 KiB
Go

package resource
import (
"bytes"
"crypto/md5"
"os"
"path/filepath"
"time"
types "github.com/metacubex/mihomo/constant/provider"
"github.com/metacubex/mihomo/log"
"github.com/samber/lo"
)
var (
fileMode os.FileMode = 0o666
dirMode os.FileMode = 0o755
)
type Parser[V any] func([]byte) (V, error)
type Fetcher[V any] struct {
resourceType string
name string
vehicle types.Vehicle
UpdatedAt *time.Time
ticker *time.Ticker
done chan struct{}
hash [16]byte
parser Parser[V]
interval time.Duration
OnUpdate func(V)
}
func (f *Fetcher[V]) Name() string {
return f.name
}
func (f *Fetcher[V]) Vehicle() types.Vehicle {
return f.vehicle
}
func (f *Fetcher[V]) VehicleType() types.VehicleType {
return f.vehicle.Type()
}
func (f *Fetcher[V]) Initial() (V, error) {
var (
buf []byte
err error
isLocal bool
forceUpdate bool
)
if stat, fErr := os.Stat(f.vehicle.Path()); fErr == nil {
buf, err = os.ReadFile(f.vehicle.Path())
modTime := stat.ModTime()
f.UpdatedAt = &modTime
isLocal = true
if f.interval != 0 && modTime.Add(f.interval).Before(time.Now()) {
log.Warnln("[Provider] %s not updated for %s, force update", f.Name(), time.Now().Sub(modTime))
forceUpdate = true
}
} else {
buf, err = f.vehicle.Read()
}
if err != nil {
return lo.Empty[V](), err
}
var contents V
if forceUpdate {
var forceBuf []byte
if forceBuf, err = f.vehicle.Read(); err == nil {
if contents, err = f.parser(forceBuf); err == nil {
isLocal = false
buf = forceBuf
}
}
}
if err != nil || !forceUpdate {
contents, err = f.parser(buf)
}
if err != nil {
if !isLocal {
return lo.Empty[V](), err
}
// parse local file error, fallback to remote
buf, err = f.vehicle.Read()
if err != nil {
return lo.Empty[V](), err
}
contents, err = f.parser(buf)
if err != nil {
return lo.Empty[V](), err
}
isLocal = false
}
if f.vehicle.Type() != types.File && !isLocal {
if err := safeWrite(f.vehicle.Path(), buf); err != nil {
return lo.Empty[V](), err
}
}
f.hash = md5.Sum(buf)
// pull contents automatically
if f.ticker != nil {
go f.pullLoop()
}
return contents, nil
}
func (f *Fetcher[V]) Update() (V, bool, error) {
buf, err := f.vehicle.Read()
if err != nil {
return lo.Empty[V](), false, err
}
now := time.Now()
hash := md5.Sum(buf)
if bytes.Equal(f.hash[:], hash[:]) {
f.UpdatedAt = &now
_ = os.Chtimes(f.vehicle.Path(), now, now)
return lo.Empty[V](), true, nil
}
contents, err := f.parser(buf)
if err != nil {
return lo.Empty[V](), false, err
}
if f.vehicle.Type() != types.File {
if err := safeWrite(f.vehicle.Path(), buf); err != nil {
return lo.Empty[V](), false, err
}
}
f.UpdatedAt = &now
f.hash = hash
return contents, false, nil
}
func (f *Fetcher[V]) Destroy() error {
if f.ticker != nil {
f.done <- struct{}{}
}
return nil
}
func (f *Fetcher[V]) pullLoop() {
for {
select {
case <-f.ticker.C:
elm, same, err := f.Update()
if err != nil {
log.Errorln("[Provider] %s pull error: %s", f.Name(), err.Error())
continue
}
if same {
log.Debugln("[Provider] %s's content doesn't change", f.Name())
continue
}
log.Infoln("[Provider] %s's content update", f.Name())
if f.OnUpdate != nil {
f.OnUpdate(elm)
}
case <-f.done:
f.ticker.Stop()
return
}
}
}
func safeWrite(path string, buf []byte) error {
dir := filepath.Dir(path)
if _, err := os.Stat(dir); os.IsNotExist(err) {
if err := os.MkdirAll(dir, dirMode); err != nil {
return err
}
}
return os.WriteFile(path, buf, fileMode)
}
func NewFetcher[V any](name string, interval time.Duration, vehicle types.Vehicle, parser Parser[V], onUpdate func(V)) *Fetcher[V] {
var ticker *time.Ticker
if interval != 0 {
ticker = time.NewTicker(interval)
}
return &Fetcher[V]{
name: name,
ticker: ticker,
vehicle: vehicle,
parser: parser,
done: make(chan struct{}, 1),
OnUpdate: onUpdate,
interval: interval,
}
}