国家、省份数据不再每个小时更新一次；WAF增加国家/地区、省份、城市、ISP等参数

internal/iplibrary/manager_provider.go

@@ -0,0 +1,148 @@
+package iplibrary
+
+import (
+	"crypto/md5"
+	"encoding/json"
+	"fmt"
+	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
+	"github.com/TeaOSLab/EdgeNode/internal/events"
+	"github.com/TeaOSLab/EdgeNode/internal/goman"
+	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
+	"github.com/TeaOSLab/EdgeNode/internal/rpc"
+	"github.com/TeaOSLab/EdgeNode/internal/utils"
+	"github.com/iwind/TeaGo/Tea"
+	_ "github.com/iwind/TeaGo/bootstrap"
+	"io/ioutil"
+	"os"
+	"sync"
+	"time"
+)
+
+var SharedProviderManager = NewProviderManager()
+
+func init() {
+	events.On(events.EventLoaded, func() {
+		goman.New(func() {
+			SharedProviderManager.Start()
+		})
+	})
+}
+
+// ProviderManager 中国省份信息管理
+type ProviderManager struct {
+	cacheFile string
+
+	providerMap map[string]int64 // name => id
+	dataHash    string           // 国家JSON的md5
+
+	locker sync.RWMutex
+
+	isUpdated bool
+}
+
+func NewProviderManager() *ProviderManager {
+	return &ProviderManager{
+		cacheFile:   Tea.Root + "/configs/region_provider.json.cache",
+		providerMap: map[string]int64{},
+	}
+}
+
+func (this *ProviderManager) Start() {
+	// 从缓存中读取
+	err := this.load()
+	if err != nil {
+		remotelogs.ErrorObject("PROVIDER_MANAGER", err)
+	}
+
+	// 第一次更新
+	err = this.loop()
+	if err != nil {
+		remotelogs.ErrorObject("PROVIDER_MANAGER", err)
+	}
+
+	// 定时更新
+	ticker := utils.NewTicker(4 * time.Hour)
+	events.On(events.EventQuit, func() {
+		ticker.Stop()
+	})
+	for ticker.Next() {
+		err := this.loop()
+		if err != nil {
+			remotelogs.ErrorObject("PROVIDER_MANAGER", err)
+		}
+	}
+}
+
+func (this *ProviderManager) Lookup(providerName string) (providerId int64) {
+	this.locker.RLock()
+	providerId, _ = this.providerMap[providerName]
+	this.locker.RUnlock()
+	return
+}
+
+// 从缓存中读取
+func (this *ProviderManager) load() error {
+	data, err := ioutil.ReadFile(this.cacheFile)
+	if err != nil {
+		if os.IsNotExist(err) {
+			return nil
+		}
+		return err
+	}
+	m := map[string]int64{}
+	err = json.Unmarshal(data, &m)
+	if err != nil {
+		return err
+	}
+	if m != nil && len(m) > 0 {
+		this.providerMap = m
+	}
+
+	return nil
+}
+
+// 更新服务商信息
+func (this *ProviderManager) loop() error {
+	if this.isUpdated {
+		return nil
+	}
+
+	rpcClient, err := rpc.SharedRPC()
+	if err != nil {
+		return err
+	}
+	resp, err := rpcClient.RegionProviderRPC().FindAllEnabledRegionProviders(rpcClient.Context(), &pb.FindAllEnabledRegionProvidersRequest{})
+	if err != nil {
+		return err
+	}
+
+	m := map[string]int64{}
+	for _, provider := range resp.RegionProviders {
+		for _, code := range provider.Codes {
+			m[code] = provider.Id
+		}
+	}
+
+	// 检查是否有更新
+	data, err := json.Marshal(m)
+	if err != nil {
+		return err
+	}
+	hash := md5.New()
+	hash.Write(data)
+	dataHash := fmt.Sprintf("%x", hash.Sum(nil))
+	if this.dataHash == dataHash {
+		return nil
+	}
+	this.dataHash = dataHash
+
+	this.locker.Lock()
+	this.providerMap = m
+	this.isUpdated = true
+	this.locker.Unlock()
+
+	// 保存到本地缓存
+
+	err = ioutil.WriteFile(this.cacheFile, data, 0666)
+	return err
+}