EdgeNode/internal/iplibrary/updater.go

package iplibrary

import (
	"fmt"
	"github.com/TeaOSLab/EdgeCommon/pkg/nodeconfigs"
	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
	"github.com/TeaOSLab/EdgeNode/internal/errors"
	"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/iwind/TeaGo/Tea"
	"os"
	"time"
)

var SharedUpdater = NewUpdater()

func init() {
	events.On(events.EventStart, func() {
		goman.New(func() {
			SharedUpdater.Start()
		})
	})
	events.On(events.EventQuit, func() {
		SharedUpdater.Stop()
	})
}

// Updater IP库更新程序
type Updater struct {
	ticker *time.Ticker
}

// NewUpdater 获取新对象
func NewUpdater() *Updater {
	return &Updater{}
}

// Start 开始更新
func (this *Updater) Start() {
	// 这里不需要太频繁检查更新，因为通常不需要更新IP库
	this.ticker = time.NewTicker(1 * time.Hour)
	for range this.ticker.C {
		err := this.loop()
		if err != nil {
			remotelogs.ErrorObject("IP_LIBRARY", err)
		}
	}
}

func (this *Updater) Stop() {
	if this.ticker != nil {
		this.ticker.Stop()
	}
}

// 单次任务
func (this *Updater) loop() error {
	nodeConfig, err := nodeconfigs.SharedNodeConfig()
	if err != nil {
		return err
	}
	if nodeConfig.GlobalConfig == nil {
		return nil
	}
	code := nodeConfig.GlobalConfig.IPLibrary.Code
	if len(code) == 0 {
		code = serverconfigs.DefaultIPLibraryType
	}

	rpcClient, err := rpc.SharedRPC()
	if err != nil {
		return err
	}
	libraryResp, err := rpcClient.IPLibraryRPC().FindLatestIPLibraryWithType(rpcClient.Context(), &pb.FindLatestIPLibraryWithTypeRequest{Type: code})
	if err != nil {
		return err
	}
	lib := libraryResp.IpLibrary
	if lib == nil || lib.File == nil {
		return nil
	}

	typeInfo := serverconfigs.FindIPLibraryWithType(code)
	if typeInfo == nil {
		return errors.New("invalid ip library code '" + code + "'")
	}

	path := Tea.Root + "/resources/ipdata/" + code + "/" + code + "." + fmt.Sprintf("%d", lib.CreatedAt) + typeInfo.GetString("ext")

	// 是否已经存在
	_, err = os.Stat(path)
	if err == nil {
		return nil
	}

	// 开始下载
	fileChunkIdsResp, err := rpcClient.FileChunkRPC().FindAllFileChunkIds(rpcClient.Context(), &pb.FindAllFileChunkIdsRequest{FileId: lib.File.Id})
	if err != nil {
		return err
	}
	chunkIds := fileChunkIdsResp.FileChunkIds
	if len(chunkIds) == 0 {
		return nil
	}
	isOk := false

	fp, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY, 0666)
	if err != nil {
		return err
	}

	defer func() {
		// 如果保存不成功就直接删除
		if !isOk {
			_ = fp.Close()
			_ = os.Remove(path)
		}
	}()
	for _, chunkId := range chunkIds {
		chunkResp, err := rpcClient.FileChunkRPC().DownloadFileChunk(rpcClient.Context(), &pb.DownloadFileChunkRequest{FileChunkId: chunkId})
		if err != nil {
			return err
		}
		chunk := chunkResp.FileChunk

		if chunk == nil {
			continue
		}
		_, err = fp.Write(chunk.Data)
		if err != nil {
			return err
		}
	}

	err = fp.Close()
	if err != nil {
		return err
	}

	// 重新加载
	library, err := SharedManager.Load()
	if err != nil {
		return err
	}
	SharedLibrary = library

	isOk = true

	return nil
}