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

internal/iplibrary/manager_city.go

@@ -0,0 +1,149 @@
+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"
+	"github.com/iwind/TeaGo/types"
+	"io/ioutil"
+	"os"
+	"sync"
+	"time"
+)
+
+var SharedCityManager = NewCityManager()
+
+func init() {
+	events.On(events.EventLoaded, func() {
+		goman.New(func() {
+			SharedCityManager.Start()
+		})
+	})
+}
+
+// CityManager 中国省份信息管理
+type CityManager struct {
+	cacheFile string
+
+	cityMap  map[string]int64 // provinceName_cityName => cityName
+	dataHash string           // 国家JSON的md5
+
+	locker sync.RWMutex
+
+	isUpdated bool
+}
+
+func NewCityManager() *CityManager {
+	return &CityManager{
+		cacheFile: Tea.Root + "/configs/region_city.json.cache",
+		cityMap:   map[string]int64{},
+	}
+}
+
+func (this *CityManager) Start() {
+	// 从缓存中读取
+	err := this.load()
+	if err != nil {
+		remotelogs.ErrorObject("CITY_MANAGER", err)
+	}
+
+	// 第一次更新
+	err = this.loop()
+	if err != nil {
+		remotelogs.ErrorObject("City_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("CITY_MANAGER", err)
+		}
+	}
+}
+
+func (this *CityManager) Lookup(provinceId int64, cityName string) (cityId int64) {
+	this.locker.RLock()
+	cityId, _ = this.cityMap[types.String(provinceId)+"_"+cityName]
+	this.locker.RUnlock()
+	return
+}
+
+// 从缓存中读取
+func (this *CityManager) 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.cityMap = m
+	}
+
+	return nil
+}
+
+// 更新城市信息
+func (this *CityManager) loop() error {
+	if this.isUpdated {
+		return nil
+	}
+
+	rpcClient, err := rpc.SharedRPC()
+	if err != nil {
+		return err
+	}
+	resp, err := rpcClient.RegionCityRPC().FindAllEnabledRegionCities(rpcClient.Context(), &pb.FindAllEnabledRegionCitiesRequest{})
+	if err != nil {
+		return err
+	}
+
+	m := map[string]int64{}
+	for _, city := range resp.RegionCities {
+		for _, code := range city.Codes {
+			m[types.String(city.RegionProvinceId)+"_"+code] = city.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.cityMap = m
+	this.isUpdated = true
+	this.locker.Unlock()
+
+	// 保存到本地缓存
+
+	err = ioutil.WriteFile(this.cacheFile, data, 0666)
+	return err
+}