waf-platform/EdgeAPI/internal/rpc/services/httpdns/service_httpdns_node.go

package httpdns

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"path/filepath"

	"github.com/TeaOSLab/EdgeAPI/internal/db/models"
	"github.com/TeaOSLab/EdgeAPI/internal/goman"
	"github.com/TeaOSLab/EdgeAPI/internal/installers"
	"github.com/TeaOSLab/EdgeAPI/internal/rpc/services"
	"github.com/TeaOSLab/EdgeAPI/internal/setup"
	rpcutils "github.com/TeaOSLab/EdgeAPI/internal/rpc/utils"
	"github.com/TeaOSLab/EdgeCommon/pkg/nodeconfigs"
	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
	"github.com/iwind/TeaGo/dbs"
	"github.com/iwind/TeaGo/logs"
	stringutil "github.com/iwind/TeaGo/utils/string"
)

// HTTPDNSNodeService HTTPDNS节点服务
type HTTPDNSNodeService struct {
	services.BaseService
	pb.UnimplementedHTTPDNSNodeServiceServer
}

func (this *HTTPDNSNodeService) CreateHTTPDNSNode(ctx context.Context, req *pb.CreateHTTPDNSNodeRequest) (*pb.CreateHTTPDNSNodeResponse, error) {
	_, err := this.ValidateAdmin(ctx)
	if err != nil {
		return nil, err
	}
	if req.ClusterId <= 0 {
		return nil, errors.New("required 'clusterId'")
	}
	var nodeId int64
	err = this.RunTx(func(tx *dbs.Tx) error {
		nodeId, err = models.SharedHTTPDNSNodeDAO.CreateNode(tx, req.ClusterId, req.Name, req.InstallDir, req.IsOn)
		if err != nil {
			return err
		}
		return notifyHTTPDNSClusterTask(tx, req.ClusterId, models.HTTPDNSNodeTaskTypeConfigChanged)
	})
	if err != nil {
		return nil, err
	}
	return &pb.CreateHTTPDNSNodeResponse{NodeId: nodeId}, nil
}

func (this *HTTPDNSNodeService) UpdateHTTPDNSNode(ctx context.Context, req *pb.UpdateHTTPDNSNodeRequest) (*pb.RPCSuccess, error) {
	_, err := this.ValidateAdmin(ctx)
	if err != nil {
		return nil, err
	}
	err = this.RunTx(func(tx *dbs.Tx) error {
		node, err := models.SharedHTTPDNSNodeDAO.FindEnabledNode(tx, req.NodeId)
		if err != nil {
			return err
		}
		if node == nil {
			return errors.New("node not found")
		}

		err = models.SharedHTTPDNSNodeDAO.UpdateNode(tx, req.NodeId, req.Name, req.InstallDir, req.IsOn)
		if err != nil {
			return err
		}
		return notifyHTTPDNSClusterTask(tx, int64(node.ClusterId), models.HTTPDNSNodeTaskTypeConfigChanged)
	})
	if err != nil {
		return nil, err
	}
	return this.Success()
}

func (this *HTTPDNSNodeService) DeleteHTTPDNSNode(ctx context.Context, req *pb.DeleteHTTPDNSNodeRequest) (*pb.RPCSuccess, error) {
	_, err := this.ValidateAdmin(ctx)
	if err != nil {
		return nil, err
	}
	err = this.RunTx(func(tx *dbs.Tx) error {
		node, err := models.SharedHTTPDNSNodeDAO.FindEnabledNode(tx, req.NodeId)
		if err != nil {
			return err
		}
		if node == nil {
			return nil
		}

		err = models.SharedHTTPDNSNodeDAO.DisableNode(tx, req.NodeId)
		if err != nil {
			return err
		}
		return notifyHTTPDNSClusterTask(tx, int64(node.ClusterId), models.HTTPDNSNodeTaskTypeConfigChanged)
	})
	if err != nil {
		return nil, err
	}
	return this.Success()
}

func (this *HTTPDNSNodeService) FindHTTPDNSNode(ctx context.Context, req *pb.FindHTTPDNSNodeRequest) (*pb.FindHTTPDNSNodeResponse, error) {
	nodeId := req.NodeId
	if nodeId <= 0 {
		parsedNodeId, nodeErr := this.ValidateHTTPDNSNode(ctx)
		if nodeErr != nil {
			return nil, errors.New("invalid 'nodeId'")
		}
		nodeId = parsedNodeId
	} else {
		_, _, validateErr := this.ValidateAdminAndUser(ctx, true)
		if validateErr != nil {
			if _, nodeErr := this.ValidateHTTPDNSNode(ctx); nodeErr != nil {
				return nil, validateErr
			}
		}
	}

	node, err := models.SharedHTTPDNSNodeDAO.FindEnabledNode(this.NullTx(), nodeId)
	if err != nil {
		return nil, err
	}

	pbNode := toPBNode(node)

	// 认证信息
	if pbNode != nil {
		login, loginErr := models.SharedNodeLoginDAO.FindEnabledNodeLoginWithNodeId(this.NullTx(), nodeconfigs.NodeRoleHTTPDNS, nodeId)
		if loginErr != nil {
			return nil, loginErr
		}
		if login != nil {
			pbNode.NodeLogin = &pb.NodeLogin{
				Id:     int64(login.Id),
				Name:   login.Name,
				Type:   login.Type,
				Params: login.Params,
			}
		}
	}

	return &pb.FindHTTPDNSNodeResponse{Node: pbNode}, nil
}

func (this *HTTPDNSNodeService) ListHTTPDNSNodes(ctx context.Context, req *pb.ListHTTPDNSNodesRequest) (*pb.ListHTTPDNSNodesResponse, error) {
	_, _, err := this.ValidateAdminAndUser(ctx, true)
	if err != nil {
		return nil, err
	}
	nodes, err := models.SharedHTTPDNSNodeDAO.ListEnabledNodes(this.NullTx(), req.ClusterId)
	if err != nil {
		return nil, err
	}
	var pbNodes []*pb.HTTPDNSNode
	for _, node := range nodes {
		pbNodes = append(pbNodes, toPBNode(node))
	}
	return &pb.ListHTTPDNSNodesResponse{Nodes: pbNodes}, nil
}

func (this *HTTPDNSNodeService) UpdateHTTPDNSNodeStatus(ctx context.Context, req *pb.UpdateHTTPDNSNodeStatusRequest) (*pb.RPCSuccess, error) {
	nodeId := req.GetNodeId()
	isAdminCaller := false
	if nodeId > 0 {
		if _, adminErr := this.ValidateAdmin(ctx); adminErr == nil {
			isAdminCaller = true
		}
	}
	if !isAdminCaller {
		if nodeId <= 0 {
			parsedNodeId, err := this.ValidateHTTPDNSNode(ctx)
			if err != nil {
				return nil, err
			}
			nodeId = parsedNodeId
		}
	}
	if nodeId <= 0 {
		return nil, errors.New("invalid 'nodeId'")
	}

	err := models.SharedHTTPDNSNodeDAO.UpdateNodeStatus(this.NullTx(), nodeId, req.GetIsUp(), req.GetIsInstalled(), req.GetIsActive(), req.GetStatusJSON(), req.GetInstallStatusJSON())
	if err != nil {
		return nil, err
	}

	if isAdminCaller && shouldTriggerHTTPDNSInstall(req.GetInstallStatusJSON()) {
		goman.New(func() {
			installErr := installers.SharedHTTPDNSNodeQueue().InstallNodeProcess(nodeId, false)
			if installErr != nil {
				logs.Println("[RPC][HTTPDNS]install node failed:", installErr.Error())
			}
		})
	}

	return this.Success()
}

// UpdateHTTPDNSNodeLogin 修改HTTPDNS节点登录信息
func (this *HTTPDNSNodeService) UpdateHTTPDNSNodeLogin(ctx context.Context, req *pb.UpdateHTTPDNSNodeLoginRequest) (*pb.RPCSuccess, error) {
	_, err := this.ValidateAdmin(ctx)
	if err != nil {
		return nil, err
	}

	var tx = this.NullTx()

	if req.NodeLogin.Id <= 0 {
		loginId, createErr := models.SharedNodeLoginDAO.CreateNodeLogin(tx, nodeconfigs.NodeRoleHTTPDNS, req.NodeId, req.NodeLogin.Name, req.NodeLogin.Type, req.NodeLogin.Params)
		if createErr != nil {
			return nil, createErr
		}
		req.NodeLogin.Id = loginId
	}

	err = models.SharedNodeLoginDAO.UpdateNodeLogin(tx, req.NodeLogin.Id, req.NodeLogin.Name, req.NodeLogin.Type, req.NodeLogin.Params)
	if err != nil {
		return nil, err
	}

	return this.Success()
}

// CheckHTTPDNSNodeLatestVersion 检查HTTPDNS节点新版本
func (this *HTTPDNSNodeService) CheckHTTPDNSNodeLatestVersion(ctx context.Context, req *pb.CheckHTTPDNSNodeLatestVersionRequest) (*pb.CheckHTTPDNSNodeLatestVersionResponse, error) {
	_, _, _, err := rpcutils.ValidateRequest(ctx, rpcutils.UserTypeAdmin, rpcutils.UserTypeHTTPDNS)
	if err != nil {
		return nil, err
	}

	deployFiles := installers.SharedDeployManager.LoadHTTPDNSNodeFiles()
	for _, file := range deployFiles {
		if file.OS == req.Os && file.Arch == req.Arch && stringutil.VersionCompare(file.Version, req.CurrentVersion) > 0 {
			return &pb.CheckHTTPDNSNodeLatestVersionResponse{
				HasNewVersion: true,
				NewVersion:    file.Version,
			}, nil
		}
	}
	return &pb.CheckHTTPDNSNodeLatestVersionResponse{HasNewVersion: false}, nil
}

// DownloadHTTPDNSNodeInstallationFile 下载最新HTTPDNS节点安装文件
func (this *HTTPDNSNodeService) DownloadHTTPDNSNodeInstallationFile(ctx context.Context, req *pb.DownloadHTTPDNSNodeInstallationFileRequest) (*pb.DownloadHTTPDNSNodeInstallationFileResponse, error) {
	nodeId, err := this.ValidateHTTPDNSNode(ctx)
	if err != nil {
		return nil, err
	}

	// 检查自动升级开关
	upgradeConfig, _ := setup.LoadUpgradeConfig()
	if upgradeConfig != nil && !upgradeConfig.AutoUpgrade {
		return &pb.DownloadHTTPDNSNodeInstallationFileResponse{}, nil
	}

	var file = installers.SharedDeployManager.FindHTTPDNSNodeFile(req.Os, req.Arch)
	if file == nil {
		return &pb.DownloadHTTPDNSNodeInstallationFileResponse{}, nil
	}

	sum, err := file.Sum()
	if err != nil {
		return nil, err
	}

	data, offset, err := file.Read(req.ChunkOffset)
	if err != nil && err != io.EOF {
		return nil, err
	}

	// 增加下载速度监控
	installers.SharedUpgradeLimiter.UpdateNodeBytes(nodeconfigs.NodeRoleHTTPDNS, nodeId, int64(len(data)))

	return &pb.DownloadHTTPDNSNodeInstallationFileResponse{
		Sum:       sum,
		Offset:    offset,
		ChunkData: data,
		Version:   file.Version,
		Filename:  filepath.Base(file.Path),
	}, nil
}

// CountAllUpgradeHTTPDNSNodesWithClusterId 计算需要升级的HTTPDNS节点数量
func (this *HTTPDNSNodeService) CountAllUpgradeHTTPDNSNodesWithClusterId(ctx context.Context, req *pb.CountAllUpgradeHTTPDNSNodesWithClusterIdRequest) (*pb.RPCCountResponse, error) {
	_, err := this.ValidateAdmin(ctx)
	if err != nil {
		return nil, err
	}

	var tx = this.NullTx()
	deployFiles := installers.SharedDeployManager.LoadHTTPDNSNodeFiles()
	total := int64(0)
	for _, deployFile := range deployFiles {
		count, err := models.SharedHTTPDNSNodeDAO.CountAllLowerVersionNodesWithClusterId(tx, req.ClusterId, deployFile.OS, deployFile.Arch, deployFile.Version)
		if err != nil {
			return nil, err
		}
		total += count
	}
	return this.SuccessCount(total)
}

// FindAllUpgradeHTTPDNSNodesWithClusterId 列出所有需要升级的HTTPDNS节点
func (this *HTTPDNSNodeService) FindAllUpgradeHTTPDNSNodesWithClusterId(ctx context.Context, req *pb.FindAllUpgradeHTTPDNSNodesWithClusterIdRequest) (*pb.FindAllUpgradeHTTPDNSNodesWithClusterIdResponse, error) {
	_, err := this.ValidateAdmin(ctx)
	if err != nil {
		return nil, err
	}

	var tx = this.NullTx()
	deployFiles := installers.SharedDeployManager.LoadHTTPDNSNodeFiles()
	var result []*pb.FindAllUpgradeHTTPDNSNodesWithClusterIdResponse_HTTPDNSNodeUpgrade
	for _, deployFile := range deployFiles {
		nodes, err := models.SharedHTTPDNSNodeDAO.FindAllLowerVersionNodesWithClusterId(tx, req.ClusterId, deployFile.OS, deployFile.Arch, deployFile.Version)
		if err != nil {
			return nil, err
		}
		for _, node := range nodes {
			// 解析状态获取当前版本
			var oldVersion string
			if len(node.Status) > 0 {
				var statusMap map[string]interface{}
				if json.Unmarshal(node.Status, &statusMap) == nil {
					if v, ok := statusMap["buildVersion"]; ok {
						oldVersion, _ = v.(string)
					}
				}
			}

			pbNode := toPBNode(node)

			// 认证信息
			login, loginErr := models.SharedNodeLoginDAO.FindEnabledNodeLoginWithNodeId(tx, nodeconfigs.NodeRoleHTTPDNS, int64(node.Id))
			if loginErr != nil {
				return nil, loginErr
			}
			if login != nil && pbNode != nil {
				pbNode.NodeLogin = &pb.NodeLogin{
					Id:     int64(login.Id),
					Name:   login.Name,
					Type:   login.Type,
					Params: login.Params,
				}
			}

			result = append(result, &pb.FindAllUpgradeHTTPDNSNodesWithClusterIdResponse_HTTPDNSNodeUpgrade{
				Node:       pbNode,
				Os:         deployFile.OS,
				Arch:       deployFile.Arch,
				OldVersion: oldVersion,
				NewVersion: deployFile.Version,
			})
		}
	}
	return &pb.FindAllUpgradeHTTPDNSNodesWithClusterIdResponse{Nodes: result}, nil
}

// UpgradeHTTPDNSNode 升级单个HTTPDNS节点
func (this *HTTPDNSNodeService) UpgradeHTTPDNSNode(ctx context.Context, req *pb.UpgradeHTTPDNSNodeRequest) (*pb.RPCSuccess, error) {
	_, err := this.ValidateAdmin(ctx)
	if err != nil {
		return nil, err
	}

	var tx = this.NullTx()

	err = models.SharedHTTPDNSNodeDAO.UpdateNodeIsInstalled(tx, req.NodeId, false)
	if err != nil {
		return nil, err
	}

	// 重置安装状态
	installStatus, err := models.SharedHTTPDNSNodeDAO.FindNodeInstallStatus(tx, req.NodeId)
	if err != nil {
		return nil, err
	}
	if installStatus == nil {
		installStatus = &models.NodeInstallStatus{}
	}
	installStatus.IsOk = false
	installStatus.IsFinished = false
	err = models.SharedHTTPDNSNodeDAO.UpdateNodeInstallStatus(tx, req.NodeId, installStatus)
	if err != nil {
		return nil, err
	}

	goman.New(func() {
		installErr := installers.SharedHTTPDNSNodeQueue().InstallNodeProcess(req.NodeId, true)
		if installErr != nil {
			logs.Println("[RPC][HTTPDNS]upgrade node failed:", installErr.Error())
		}
	})

	return this.Success()
}

func shouldTriggerHTTPDNSInstall(installStatusJSON []byte) bool {
	if len(installStatusJSON) == 0 {
		return false
	}

	installStatus := &models.NodeInstallStatus{}
	err := json.Unmarshal(installStatusJSON, installStatus)
	if err != nil {
		return false
	}
	return installStatus.IsRunning && !installStatus.IsFinished
}