EdgeAdmin/internal/utils/apinodeutils/upgrader.go

// Copyright 2023 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .

package apinodeutils

import (
	"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"
	"path/filepath"
	"runtime"
)

type Progress struct {
	Percent float64
}

type Upgrader struct {
	progress  *Progress
	apiExe    string
	apiNodeId int64
}

func NewUpgrader(apiNodeId int64) *Upgrader {
	return &Upgrader{
		apiExe:    apiExe(),
		progress:  &Progress{Percent: 0},
		apiNodeId: apiNodeId,
	}
}

func (this *Upgrader) Upgrade() error {
	sharedClient, err := rpc.SharedRPC()
	if err != nil {
		return err
	}
	apiNodeResp, err := sharedClient.APINodeRPC().FindEnabledAPINode(sharedClient.Context(0), &pb.FindEnabledAPINodeRequest{ApiNodeId: this.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(this.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 := CanUpgrade(versionResp.Version, versionResp.Os, versionResp.Arch)
	if !canUpgrade {
		return errors.New(reason)
	}

	localVersion, err := 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
}