EdgeAPI/internal/rpc/services/service_api_node.go

package services

import (
	"compress/gzip"
	"context"
	"crypto/md5"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"runtime"

	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"
	"github.com/iwind/TeaGo/Tea"
	"github.com/iwind/TeaGo/dbs"
	stringutil "github.com/iwind/TeaGo/utils/string"
)

type APINodeService struct {
	BaseService
}

// CreateAPINode 创建API节点
func (this *APINodeService) CreateAPINode(ctx context.Context, req *pb.CreateAPINodeRequest) (*pb.CreateAPINodeResponse, error) {
	_, err := this.ValidateAdmin(ctx)
	if err != nil {
		return nil, err
	}

	var tx = this.NullTx()

	nodeId, err := models.SharedAPINodeDAO.CreateAPINode(tx, req.Name, req.Description, req.HttpJSON, req.HttpsJSON, req.RestIsOn, req.RestHTTPJSON, req.RestHTTPSJSON, req.AccessAddrsJSON, req.IsOn)
	if err != nil {
		return nil, err
	}

	return &pb.CreateAPINodeResponse{ApiNodeId: nodeId}, nil
}

// UpdateAPINode 修改API节点
func (this *APINodeService) UpdateAPINode(ctx context.Context, req *pb.UpdateAPINodeRequest) (*pb.RPCSuccess, error) {
	_, err := this.ValidateAdmin(ctx)
	if err != nil {
		return nil, err
	}

	var tx = this.NullTx()

	err = models.SharedAPINodeDAO.UpdateAPINode(tx, req.ApiNodeId, req.Name, req.Description, req.HttpJSON, req.HttpsJSON, req.RestIsOn, req.RestHTTPJSON, req.RestHTTPSJSON, req.AccessAddrsJSON, req.IsOn, req.IsPrimary)
	if err != nil {
		return nil, err
	}

	return this.Success()
}

// DeleteAPINode 删除API节点
func (this *APINodeService) DeleteAPINode(ctx context.Context, req *pb.DeleteAPINodeRequest) (*pb.RPCSuccess, error) {
	_, err := this.ValidateAdmin(ctx)
	if err != nil {
		return nil, err
	}

	var tx = this.NullTx()

	err = models.SharedAPINodeDAO.DisableAPINode(tx, req.ApiNodeId)
	if err != nil {
		return nil, err
	}

	return this.Success()
}

// FindAllEnabledAPINodes 列出所有可用API节点
func (this *APINodeService) FindAllEnabledAPINodes(ctx context.Context, req *pb.FindAllEnabledAPINodesRequest) (*pb.FindAllEnabledAPINodesResponse, error) {
	_, _, _, err := rpcutils.ValidateRequest(ctx, rpcutils.UserTypeAdmin, rpcutils.UserTypeUser, rpcutils.UserTypeNode, rpcutils.UserTypeDNS, rpcutils.UserTypeAuthority)
	if err != nil {
		return nil, err
	}

	var tx = this.NullTx()

	nodes, err := models.SharedAPINodeDAO.FindAllEnabledAPINodes(tx)
	if err != nil {
		return nil, err
	}

	result := []*pb.APINode{}
	for _, node := range nodes {
		accessAddrs, err := node.DecodeAccessAddrStrings()
		if err != nil {
			return nil, err
		}

		result = append(result, &pb.APINode{
			Id:              int64(node.Id),
			IsOn:            node.IsOn,
			NodeClusterId:   int64(node.ClusterId),
			UniqueId:        node.UniqueId,
			Secret:          node.Secret,
			Name:            node.Name,
			Description:     node.Description,
			HttpJSON:        node.Http,
			HttpsJSON:       node.Https,
			AccessAddrsJSON: node.AccessAddrs,
			AccessAddrs:     accessAddrs,
			IsPrimary:       node.IsPrimary,
		})
	}

	return &pb.FindAllEnabledAPINodesResponse{ApiNodes: result}, nil
}

// CountAllEnabledAPINodes 计算API节点数量
func (this *APINodeService) CountAllEnabledAPINodes(ctx context.Context, req *pb.CountAllEnabledAPINodesRequest) (*pb.RPCCountResponse, error) {
	_, err := this.ValidateAdmin(ctx)
	if err != nil {
		return nil, err
	}

	var tx = this.NullTx()

	count, err := models.SharedAPINodeDAO.CountAllEnabledAPINodes(tx)
	if err != nil {
		return nil, err
	}

	return this.SuccessCount(count)
}

// CountAllEnabledAndOnAPINodes 计算API节点数量
func (this *APINodeService) CountAllEnabledAndOnAPINodes(ctx context.Context, req *pb.CountAllEnabledAndOnAPINodesRequest) (*pb.RPCCountResponse, error) {
	_, err := this.ValidateAdmin(ctx)
	if err != nil {
		return nil, err
	}

	var tx = this.NullTx()

	count, err := models.SharedAPINodeDAO.CountAllEnabledAndOnAPINodes(tx)
	if err != nil {
		return nil, err
	}

	return this.SuccessCount(count)
}

// ListEnabledAPINodes 列出单页的API节点
func (this *APINodeService) ListEnabledAPINodes(ctx context.Context, req *pb.ListEnabledAPINodesRequest) (*pb.ListEnabledAPINodesResponse, error) {
	_, err := this.ValidateAdmin(ctx)
	if err != nil {
		return nil, err
	}

	var tx = this.NullTx()

	nodes, err := models.SharedAPINodeDAO.ListEnabledAPINodes(tx, req.Offset, req.Size)
	if err != nil {
		return nil, err
	}

	result := []*pb.APINode{}
	for _, node := range nodes {
		accessAddrs, err := node.DecodeAccessAddrStrings()
		if err != nil {
			return nil, err
		}

		result = append(result, &pb.APINode{
			Id:              int64(node.Id),
			IsOn:            node.IsOn,
			NodeClusterId:   int64(node.ClusterId),
			UniqueId:        node.UniqueId,
			Secret:          node.Secret,
			Name:            node.Name,
			Description:     node.Description,
			HttpJSON:        node.Http,
			HttpsJSON:       node.Https,
			RestIsOn:        node.RestIsOn == 1,
			RestHTTPJSON:    node.RestHTTP,
			RestHTTPSJSON:   node.RestHTTPS,
			AccessAddrsJSON: node.AccessAddrs,
			AccessAddrs:     accessAddrs,
			StatusJSON:      node.Status,
			IsPrimary:       node.IsPrimary,
		})
	}

	return &pb.ListEnabledAPINodesResponse{ApiNodes: result}, nil
}

// FindEnabledAPINode 根据ID查找节点
func (this *APINodeService) FindEnabledAPINode(ctx context.Context, req *pb.FindEnabledAPINodeRequest) (*pb.FindEnabledAPINodeResponse, error) {
	_, _, err := this.ValidateAdminAndUser(ctx, false)
	if err != nil {
		return nil, err
	}

	var tx = this.NullTx()

	node, err := models.SharedAPINodeDAO.FindEnabledAPINode(tx, req.ApiNodeId, nil)
	if err != nil {
		return nil, err
	}

	if node == nil {
		return &pb.FindEnabledAPINodeResponse{ApiNode: nil}, nil
	}

	accessAddrs, err := node.DecodeAccessAddrStrings()
	if err != nil {
		return nil, err
	}

	result := &pb.APINode{
		Id:              int64(node.Id),
		IsOn:            node.IsOn,
		NodeClusterId:   int64(node.ClusterId),
		UniqueId:        node.UniqueId,
		Secret:          node.Secret,
		Name:            node.Name,
		Description:     node.Description,
		HttpJSON:        node.Http,
		HttpsJSON:       node.Https,
		RestIsOn:        node.RestIsOn == 1,
		RestHTTPJSON:    node.RestHTTP,
		RestHTTPSJSON:   node.RestHTTPS,
		AccessAddrsJSON: node.AccessAddrs,
		AccessAddrs:     accessAddrs,
		IsPrimary:       node.IsPrimary,
		StatusJSON:      node.Status,
	}
	return &pb.FindEnabledAPINodeResponse{ApiNode: result}, nil
}

// FindCurrentAPINodeVersion 获取当前API节点的版本
func (this *APINodeService) FindCurrentAPINodeVersion(ctx context.Context, req *pb.FindCurrentAPINodeVersionRequest) (*pb.FindCurrentAPINodeVersionResponse, error) {
	role, _, _, err := rpcutils.ValidateRequest(ctx)
	if err != nil {
		return nil, err
	}

	return &pb.FindCurrentAPINodeVersionResponse{
		Version: teaconst.Version,
		Os:      runtime.GOOS,
		Arch:    runtime.GOARCH,
		Role:    role,
	}, nil
}

// FindCurrentAPINode 获取当前API节点的信息
func (this *APINodeService) FindCurrentAPINode(ctx context.Context, req *pb.FindCurrentAPINodeRequest) (*pb.FindCurrentAPINodeResponse, error) {
	_, err := this.ValidateAdmin(ctx)
	if err != nil {
		return nil, err
	}

	var nodeId = teaconst.NodeId
	var tx *dbs.Tx
	node, err := models.SharedAPINodeDAO.FindEnabledAPINode(tx, nodeId, nil)
	if err != nil {
		return nil, err
	}
	if node == nil {
		return &pb.FindCurrentAPINodeResponse{ApiNode: nil}, nil
	}

	accessAddrs, err := node.DecodeAccessAddrStrings()
	if err != nil {
		return nil, err
	}

	return &pb.FindCurrentAPINodeResponse{ApiNode: &pb.APINode{
		Id:              int64(node.Id),
		IsOn:            node.IsOn,
		NodeClusterId:   0,
		UniqueId:        "",
		Secret:          "",
		Name:            "",
		Description:     "",
		HttpJSON:        nil,
		HttpsJSON:       nil,
		RestIsOn:        false,
		RestHTTPJSON:    nil,
		RestHTTPSJSON:   nil,
		AccessAddrsJSON: node.AccessAddrs,
		AccessAddrs:     accessAddrs,
		StatusJSON:      node.Status,
		IsPrimary:       node.IsPrimary,
		InstanceCode:    teaconst.InstanceCode,
	}}, nil
}

// CountAllEnabledAPINodesWithSSLCertId 计算使用某个SSL证书的API节点数量
func (this *APINodeService) CountAllEnabledAPINodesWithSSLCertId(ctx context.Context, req *pb.CountAllEnabledAPINodesWithSSLCertIdRequest) (*pb.RPCCountResponse, error) {
	_, err := this.ValidateAdmin(ctx)
	if err != nil {
		return nil, err
	}

	var tx = this.NullTx()
	policyIds, err := models.SharedSSLPolicyDAO.FindAllEnabledPolicyIdsWithCertId(tx, req.SslCertId)
	if err != nil {
		return nil, err
	}
	if len(policyIds) == 0 {
		return this.SuccessCount(0)
	}

	count, err := models.SharedAPINodeDAO.CountAllEnabledAPINodesWithSSLPolicyIds(tx, policyIds)
	if err != nil {
		return nil, err
	}
	return this.SuccessCount(count)
}

// DebugAPINode 修改调试模式状态
func (this *APINodeService) DebugAPINode(ctx context.Context, req *pb.DebugAPINodeRequest) (*pb.RPCSuccess, error) {
	_, err := this.ValidateAdmin(ctx)
	if err != nil {
		return nil, err
	}

	teaconst.Debug = req.Debug
	return this.Success()
}

// UploadAPINodeFile 上传新版API节点文件
func (this *APINodeService) UploadAPINodeFile(ctx context.Context, req *pb.UploadAPINodeFileRequest) (*pb.UploadAPINodeFileResponse, error) {
	_, err := this.ValidateAdmin(ctx)
	if err != nil {
		return nil, err
	}

	exe, err := os.Executable()
	if err != nil {
		return nil, errors.New("can not find executable file: " + err.Error())
	}

	var targetDir = filepath.Dir(exe)
	var targetFilename = teaconst.ProcessName // 这里不使用 filepath.Base() 是因为文件名可能变成修改后的临时文件名
	var targetCompressedFile = targetDir + "/." + targetFilename + ".gz"
	var targetFile = targetDir + "/." + targetFilename

	if req.IsFirstChunk {
		_ = os.Remove(targetCompressedFile)
		_ = os.Remove(targetFile)
	}

	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(targetCompressedFile, 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 {
		err = func() error {
			// 删除压缩文件
			defer func() {
				_ = os.Remove(targetCompressedFile)
			}()

			// 检查SUM
			fp, err := os.Open(targetCompressedFile)
			if err != nil {
				return err
			}
			defer func() {
				_ = fp.Close()
			}()

			var hash = md5.New()
			_, err = io.Copy(hash, fp)
			if err != nil {
				return err
			}

			var sum = fmt.Sprintf("%x", hash.Sum(nil))
			if sum != req.Sum {
				return errors.New("check sum failed: '" + sum + "' expected: '" + req.Sum + "'")
			}

			// 解压
			fp2, err := os.Open(targetCompressedFile)
			if err != nil {
				return err
			}

			defer func() {
				_ = fp2.Close()
			}()

			gzipReader, err := gzip.NewReader(fp2)
			if err != nil {
				return err
			}
			defer func() {
				_ = gzipReader.Close()
			}()
			targetWriter, err := os.OpenFile(targetFile, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0777)
			if err != nil {
				return err
			}
			defer func() {
				_ = targetWriter.Close()
			}()
			_, err = io.Copy(targetWriter, gzipReader)
			if err != nil {
				return err
			}

			return nil
		}()
		if err != nil {
			return nil, errors.New("extract file failed: " + err.Error())
		}

		// 检查文件是否可执行
		var versionCmd = executils.NewCmd(targetFile, "-V").WithStdout().WithStderr()
		err = versionCmd.Run()
		if err != nil {
			return nil, errors.New("test file failed: " + versionCmd.Stderr())
		}
		var newVersion = versionCmd.Stdout()
		if len(newVersion) == 0 {
			return nil, errors.New("test file failed, new version should not be empty")
		}

		// 检查版本
		if stringutil.VersionCompare(newVersion, teaconst.Version) <= 0 {
			return &pb.UploadAPINodeFileResponse{}, nil
		}

		// 替换文件
		err = os.Remove(exe)
		if err != nil {
			return nil, errors.New("remove old file failed: " + err.Error())
		}
		err = os.Rename(targetFile, exe)
		if err != nil {
			return nil, errors.New("rename file failed: " + err.Error())
		}

		// 执行升级
		if !Tea.IsTesting() { // 开发环境下防止破坏本地数据库
			var upgradeCmd = executils.NewCmd(exe, "upgrade").WithStderr()
			err = upgradeCmd.Run()
			if err != nil {
				return nil, errors.New("execute 'upgrade' command failed: " + upgradeCmd.Stderr())
			}
		}

		// 重启
		var restartCmd = executils.NewCmd(exe, "restart").WithStderr()
		err = restartCmd.Start()
		if err != nil {
			return nil, errors.New("start new process failed: " + restartCmd.Stderr())
		}
	}

	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
}