diff --git a/internal/configs/api_config.go b/internal/configs/api_config.go index de78a1c1..6a9c141e 100644 --- a/internal/configs/api_config.go +++ b/internal/configs/api_config.go @@ -155,3 +155,15 @@ func (this *APIConfig) WriteFile(path string) error { return nil } + +// Clone 克隆当前配置 +func (this *APIConfig) Clone() *APIConfig { + return &APIConfig{ + RPC: struct { + Endpoints []string `yaml:"endpoints"` + DisableUpdate bool `yaml:"disableUpdate"` + }{}, + NodeId: this.NodeId, + Secret: this.Secret, + } +} diff --git a/internal/utils/apinodeutils/upgrader.go b/internal/utils/apinodeutils/upgrader.go new file mode 100644 index 00000000..ce3582e8 --- /dev/null +++ b/internal/utils/apinodeutils/upgrader.go @@ -0,0 +1,241 @@ +// Copyright 2023 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn . + +package apinodeutils + +import ( + "bytes" + "compress/gzip" + "crypto/md5" + "errors" + "fmt" + "github.com/TeaOSLab/EdgeAdmin/internal/configs" + "github.com/TeaOSLab/EdgeAdmin/internal/rpc" + "github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb" + "github.com/iwind/TeaGo/Tea" + "github.com/iwind/TeaGo/types" + stringutil "github.com/iwind/TeaGo/utils/string" + "io" + "os" + "os/exec" + "path/filepath" + "regexp" + "runtime" +) + +type Progress struct { + Percent float64 +} + +type Upgrader struct { + progress *Progress + apiExe string +} + +func NewUpgrader() *Upgrader { + return &Upgrader{ + apiExe: Tea.Root + "/edge-api/bin/edge-api", + progress: &Progress{Percent: 0}, + } +} + +func (this *Upgrader) CanUpgrade(apiVersion string) (canUpgrade bool, reason string) { + stat, err := os.Stat(this.apiExe) + if err != nil { + return false, "stat error: " + err.Error() + } + if stat.IsDir() { + return false, "is directory" + } + + localVersion, err := this.localVersion() + if err != nil { + return false, "lookup version failed: " + err.Error() + } + if stringutil.VersionCompare(localVersion, apiVersion) <= 0 { + return false, "need not upgrade, local '" + localVersion + "' vs remote '" + apiVersion + "'" + } + + return true, "" +} + +func (this *Upgrader) Upgrade(apiNodeId int64) error { + sharedClient, err := rpc.SharedRPC() + if err != nil { + return err + } + apiNodeResp, err := sharedClient.APINodeRPC().FindEnabledAPINode(sharedClient.Context(0), &pb.FindEnabledAPINodeRequest{ApiNodeId: apiNodeId}) + if err != nil { + return err + } + var apiNode = apiNodeResp.ApiNode + if apiNode == nil { + return errors.New("could not find api node with id '" + types.String(apiNodeId) + "'") + } + + apiConfig, err := configs.LoadAPIConfig() + if err != nil { + return err + } + var newAPIConfig = apiConfig.Clone() + newAPIConfig.RPC.Endpoints = apiNode.AccessAddrs + + rpcClient, err := rpc.NewRPCClient(newAPIConfig, false) + if err != nil { + return err + } + versionResp, err := rpcClient.APINodeRPC().FindCurrentAPINodeVersion(sharedClient.Context(0), &pb.FindCurrentAPINodeVersionRequest{}) + if err != nil { + return err + } + if !Tea.IsTesting() /** 开发环境下允许突破此限制方便测试 **/ && + (stringutil.VersionCompare(versionResp.Version, "0.6.4" /** 从0.6.4开始支持 **/) <= 0 || versionResp.Os != runtime.GOOS || versionResp.Arch != runtime.GOARCH) { + return errors.New("could not upgrade api node v" + versionResp.Version + "/" + versionResp.Os + "/" + versionResp.Arch) + } + + // 检查本地文件版本 + canUpgrade, reason := this.CanUpgrade(versionResp.Version) + if !canUpgrade { + return errors.New(reason) + } + + localVersion, err := this.localVersion() + if err != nil { + return errors.New("lookup version failed: " + err.Error()) + } + + // 检查要升级的文件 + var gzFile = this.apiExe + "." + localVersion + ".gz" + + gzReader, err := os.Open(gzFile) + if err != nil { + if !os.IsNotExist(err) { + return err + } + err = func() error { + // 压缩文件 + exeReader, err := os.Open(this.apiExe) + if err != nil { + return err + } + defer func() { + _ = exeReader.Close() + }() + var tmpGzFile = gzFile + ".tmp" + gzFileWriter, err := os.OpenFile(tmpGzFile, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0666) + if err != nil { + return err + } + var gzWriter = gzip.NewWriter(gzFileWriter) + defer func() { + _ = gzWriter.Close() + _ = gzFileWriter.Close() + + _ = os.Rename(tmpGzFile, gzFile) + }() + _, err = io.Copy(gzWriter, exeReader) + if err != nil { + return err + } + + return nil + }() + if err != nil { + return err + } + gzReader, err = os.Open(gzFile) + if err != nil { + return err + } + } + + defer func() { + _ = gzReader.Close() + }() + + // 开始上传 + var hash = md5.New() + var buf = make([]byte, 128*4096) + var isFirst = true + stat, err := gzReader.Stat() + if err != nil { + return err + } + var totalSize = stat.Size() + if totalSize == 0 { + _ = gzReader.Close() + _ = os.Remove(gzFile) + return errors.New("invalid gz file") + } + + var uploadedSize int64 = 0 + for { + n, err := gzReader.Read(buf) + if n > 0 { + // 计算Hash + hash.Write(buf[:n]) + + // 上传 + _, uploadErr := rpcClient.APINodeRPC().UploadAPINodeFile(rpcClient.Context(0), &pb.UploadAPINodeFileRequest{ + Filename: filepath.Base(this.apiExe), + Sum: "", + ChunkData: buf[:n], + IsFirstChunk: isFirst, + IsLastChunk: false, + }) + if uploadErr != nil { + return uploadErr + } + + // 进度 + uploadedSize += int64(n) + this.progress = &Progress{Percent: float64(uploadedSize*100) / float64(totalSize)} + } + if isFirst { + isFirst = false + } + if err != nil { + if err != io.EOF { + return err + } + if err == io.EOF { + _, uploadErr := rpcClient.APINodeRPC().UploadAPINodeFile(rpcClient.Context(0), &pb.UploadAPINodeFileRequest{ + Filename: filepath.Base(this.apiExe), + Sum: fmt.Sprintf("%x", hash.Sum(nil)), + ChunkData: buf[:n], + IsFirstChunk: isFirst, + IsLastChunk: true, + }) + if uploadErr != nil { + return uploadErr + } + + break + } + } + } + + return nil +} + +func (this *Upgrader) Progress() *Progress { + return this.progress +} + +func (this *Upgrader) localVersion() (string, error) { + var cmd = exec.Command(this.apiExe, "-V") + var output = &bytes.Buffer{} + cmd.Stdout = output + err := cmd.Run() + if err != nil { + return "", err + } + var localVersion = output.String() + + // 检查版本号 + var reg = regexp.MustCompile(`^[\d.]+$`) + if !reg.MatchString(localVersion) { + return "", errors.New("lookup version failed: " + localVersion) + } + + return localVersion, nil +} diff --git a/internal/utils/apinodeutils/upgrader_test.go b/internal/utils/apinodeutils/upgrader_test.go new file mode 100644 index 00000000..2f92b2fe --- /dev/null +++ b/internal/utils/apinodeutils/upgrader_test.go @@ -0,0 +1,22 @@ +// Copyright 2023 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn . + +package apinodeutils_test + +import ( + "github.com/TeaOSLab/EdgeAdmin/internal/utils/apinodeutils" + _ "github.com/iwind/TeaGo/bootstrap" + "testing" +) + +func TestUpgrader_CanUpgrade(t *testing.T) { + var upgrader = apinodeutils.NewUpgrader() + t.Log(upgrader.CanUpgrade("0.6.3")) +} + +func TestUpgrader_Upgrade(t *testing.T) { + var upgrader = apinodeutils.NewUpgrader() + err := upgrader.Upgrade(1) + if err != nil { + t.Fatal(err) + } +}