远程升级API节点时自动上传边缘节点安装文件

internal/installers/deploy_manager.go

@@ -5,6 +5,7 @@ import (
 	"github.com/iwind/TeaGo/files"
 	stringutil "github.com/iwind/TeaGo/utils/string"
 	"regexp"
+	"sync"
 )
 
 var SharedDeployManager = NewDeployManager()
@@ -12,9 +13,12 @@ var SharedDeployManager = NewDeployManager()
 // DeployManager 节点部署文件管理器
 // 如果节点部署文件有变化，需要重启API节点以便于生效
 type DeployManager struct {
-	dir         string
+	dir string
+
 	nodeFiles   []*DeployFile
 	nsNodeFiles []*DeployFile
+
+	locker sync.Mutex
 }
 
 // NewDeployManager 获取新节点部署文件管理器
@@ -29,24 +33,27 @@ func NewDeployManager() *DeployManager {
 
 // LoadNodeFiles 加载所有边缘节点文件
 func (this *DeployManager) LoadNodeFiles() []*DeployFile {
+	this.locker.Lock()
+	defer this.locker.Unlock()
+
 	if len(this.nodeFiles) > 0 {
 		return this.nodeFiles
 	}
 
-	keyMap := map[string]*DeployFile{} // key => File
+	var keyMap = map[string]*DeployFile{} // key => File
 
-	reg := regexp.MustCompile(`^edge-node-(\w+)-(\w+)-v([0-9.]+)\.zip$`)
+	var reg = regexp.MustCompile(`^edge-node-(\w+)-(\w+)-v([0-9.]+)\.zip$`)
 	for _, file := range files.NewFile(this.dir).List() {
-		name := file.Name()
+		var name = file.Name()
 		if !reg.MatchString(name) {
 			continue
 		}
-		matches := reg.FindStringSubmatch(name)
-		osName := matches[1]
-		arch := matches[2]
-		version := matches[3]
+		var matches = reg.FindStringSubmatch(name)
+		var osName = matches[1]
+		var arch = matches[2]
+		var version = matches[3]
 
-		key := osName + "_" + arch
+		var key = osName + "_" + arch
 		oldFile, ok := keyMap[key]
 		if ok && stringutil.VersionCompare(oldFile.Version, version) > 0 {
 			continue
@@ -59,7 +66,7 @@ func (this *DeployManager) LoadNodeFiles() []*DeployFile {
 		}
 	}
 
-	result := []*DeployFile{}
+	var result = []*DeployFile{}
 	for _, v := range keyMap {
 		result = append(result, v)
 	}
@@ -81,24 +88,27 @@ func (this *DeployManager) FindNodeFile(os string, arch string) *DeployFile {
 
 // LoadNSNodeFiles 加载所有NS节点安装文件
 func (this *DeployManager) LoadNSNodeFiles() []*DeployFile {
+	this.locker.Lock()
+	defer this.locker.Unlock()
+
 	if len(this.nsNodeFiles) > 0 {
 		return this.nsNodeFiles
 	}
 
-	keyMap := map[string]*DeployFile{} // key => File
+	var keyMap = map[string]*DeployFile{} // key => File
 
-	reg := regexp.MustCompile(`^edge-dns-(\w+)-(\w+)-v([0-9.]+)\.zip$`)
+	var reg = regexp.MustCompile(`^edge-dns-(\w+)-(\w+)-v([0-9.]+)\.zip$`)
 	for _, file := range files.NewFile(this.dir).List() {
-		name := file.Name()
+		var name = file.Name()
 		if !reg.MatchString(name) {
 			continue
 		}
-		matches := reg.FindStringSubmatch(name)
-		osName := matches[1]
-		arch := matches[2]
-		version := matches[3]
+		var matches = reg.FindStringSubmatch(name)
+		var osName = matches[1]
+		var arch = matches[2]
+		var version = matches[3]
 
-		key := osName + "_" + arch
+		var key = osName + "_" + arch
 		oldFile, ok := keyMap[key]
 		if ok && stringutil.VersionCompare(oldFile.Version, version) > 0 {
 			continue
@@ -111,7 +121,7 @@ func (this *DeployManager) LoadNSNodeFiles() []*DeployFile {
 		}
 	}
 
-	result := []*DeployFile{}
+	var result = []*DeployFile{}
 	for _, v := range keyMap {
 		result = append(result, v)
 	}
@@ -130,3 +140,12 @@ func (this *DeployManager) FindNSNodeFile(os string, arch string) *DeployFile {
 	}
 	return nil
 }
+
+// Reload 重置缓存
+func (this *DeployManager) Reload() {
+	this.locker.Lock()
+	defer this.locker.Unlock()
+
+	this.nodeFiles = nil
+	this.nsNodeFiles = nil
+}

internal/rpc/services/service_api_node.go

@@ -8,6 +8,7 @@ import (
 	"fmt"
 	teaconst "github.com/TeaOSLab/EdgeAPI/internal/const"
 	"github.com/TeaOSLab/EdgeAPI/internal/db/models"
+	"github.com/TeaOSLab/EdgeAPI/internal/installers"
 	rpcutils "github.com/TeaOSLab/EdgeAPI/internal/rpc/utils"
 	executils "github.com/TeaOSLab/EdgeAPI/internal/utils/exec"
 	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
@@ -481,3 +482,106 @@ func (this *APINodeService) UploadAPINodeFile(ctx context.Context, req *pb.Uploa
 
 	return &pb.UploadAPINodeFileResponse{}, nil
 }
+
+// UploadDeployFileToAPINode 上传节点安装文件
+func (this *APINodeService) UploadDeployFileToAPINode(ctx context.Context, req *pb.UploadDeployFileToAPINodeRequest) (*pb.RPCSuccess, error) {
+	_, err := this.ValidateAdmin(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	var targetDir = Tea.Root + "/deploy/"
+	var targetTmpFile = targetDir + "/" + req.Filename + ".tmp"
+	var targetFile = targetDir + "/" + req.Filename
+
+	if req.IsFirstChunk {
+		_ = os.Remove(targetTmpFile)
+	}
+
+	if len(req.ChunkData) > 0 {
+		err = func() error {
+			var flags = os.O_CREATE | os.O_WRONLY
+			if req.IsFirstChunk {
+				flags |= os.O_TRUNC
+			} else {
+				flags |= os.O_APPEND
+			}
+			fp, err := os.OpenFile(targetTmpFile, flags, 0666)
+			if err != nil {
+				return err
+			}
+			defer func() {
+				_ = fp.Close()
+			}()
+
+			_, err = fp.Write(req.ChunkData)
+			return err
+		}()
+		if err != nil {
+			return nil, errors.New("write file failed: " + err.Error())
+		}
+	}
+
+	if req.IsLastChunk {
+		// 检查SUM
+		fp, err := os.Open(targetTmpFile)
+		if err != nil {
+			return nil, err
+		}
+		var hash = md5.New()
+		_, err = io.Copy(hash, fp)
+		_ = fp.Close()
+		if err != nil {
+			return nil, err
+		}
+		var tmpSum = fmt.Sprintf("%x", hash.Sum(nil))
+		if tmpSum != req.Sum {
+			_ = os.Remove(targetTmpFile)
+			return nil, errors.New("check sum failed")
+		}
+
+		// 正式改名
+		err = os.Rename(targetTmpFile, targetFile)
+		if err != nil {
+			return nil, errors.New("rename failed: " + err.Error())
+		}
+
+		// 重载数据
+		installers.SharedDeployManager.Reload()
+	}
+
+	return this.Success()
+}
+
+// FindLatestDeployFiles 查找已有节点安装文件信息
+func (this *APINodeService) FindLatestDeployFiles(ctx context.Context, req *pb.FindLatestDeployFilesRequest) (*pb.FindLatestDeployFilesResponse, error) {
+	_, err := this.ValidateAdmin(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	var pbNodeFiles = []*pb.FindLatestDeployFilesResponse_DeployFile{}
+	var nodeFiles = installers.SharedDeployManager.LoadNodeFiles()
+	for _, nodeFile := range nodeFiles {
+		pbNodeFiles = append(pbNodeFiles, &pb.FindLatestDeployFilesResponse_DeployFile{
+			Os:      nodeFile.OS,
+			Arch:    nodeFile.Arch,
+			Version: nodeFile.Version,
+		})
+	}
+
+	var pbNSNodeFiles = []*pb.FindLatestDeployFilesResponse_DeployFile{}
+	var nsNodeFiles = installers.SharedDeployManager.LoadNSNodeFiles()
+	for _, nodeFile := range nsNodeFiles {
+		pbNSNodeFiles = append(pbNSNodeFiles, &pb.FindLatestDeployFilesResponse_DeployFile{
+			Os:      nodeFile.OS,
+			Arch:    nodeFile.Arch,
+			Version: nodeFile.Version,
+		})
+	}
+
+	return &pb.FindLatestDeployFilesResponse{
+		NodeDeployFiles:   pbNodeFiles,
+		NsNodeDeployFiles: pbNSNodeFiles,
+	}, nil
+}