waf-platform/EdgeReporter/internal/nodes/api_stream.go

package nodes

import (
	"context"
	"encoding/json"
	"github.com/TeaOSLab/EdgeCommon/pkg/errors"
	"github.com/TeaOSLab/EdgeCommon/pkg/messageconfigs"
	"github.com/TeaOSLab/EdgeCommon/pkg/reporterconfigs"
	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
	teaconst "github.com/TeaOSLab/EdgeReporter/internal/const"
	"github.com/TeaOSLab/EdgeReporter/internal/events"
	"github.com/TeaOSLab/EdgeReporter/internal/remotelogs"
	"github.com/TeaOSLab/EdgeReporter/internal/rpc"
	"github.com/TeaOSLab/EdgeReporter/internal/utils"
	"github.com/iwind/TeaGo/logs"
	"os/exec"
	"strconv"
	"time"
)

type APIStream struct {
	stream pb.ReportNodeService_ReportNodeStreamClient

	isQuiting  bool
	cancelFunc context.CancelFunc
}

func NewAPIStream() *APIStream {
	return &APIStream{}
}

func (this *APIStream) Start() {
	events.On(events.EventQuit, func() {
		this.isQuiting = true
		if this.cancelFunc != nil {
			this.cancelFunc()
		}
	})
	for {
		if this.isQuiting {
			return
		}
		err := this.loop()
		if err != nil {
			remotelogs.Error("API_STREAM", err.Error())
			time.Sleep(10 * time.Second)
			continue
		}
		time.Sleep(1 * time.Second)
	}
}

func (this *APIStream) loop() error {
	rpcClient, err := rpc.SharedRPC()
	if err != nil {
		return errors.Wrap(err)
	}

	ctx, cancelFunc := context.WithCancel(rpcClient.Context())
	this.cancelFunc = cancelFunc
	defer func() {
		cancelFunc()
	}()

	nodeStream, err := rpcClient.ReportNodeRPC().ReportNodeStream(ctx)
	if err != nil {
		if this.isQuiting {
			return nil
		}
		return errors.Wrap(err)
	}
	this.stream = nodeStream

	for {
		if this.isQuiting {
			logs.Println("API_STREAM", "quit")
			break
		}

		message, err := nodeStream.Recv()
		if err != nil {
			if this.isQuiting {
				remotelogs.Println("API_STREAM", "quit")
				return nil
			}
			return errors.Wrap(err)
		}

		// 处理消息
		switch message.Code {
		case reporterconfigs.MessageCodeConnectedAPINode: // 连接API节点成功
			err = this.handleConnectedAPINode(message)
		case reporterconfigs.MessageCodeNewNodeTask: // 有新的任务
			err = this.handleNewNodeTask(message)
		case reporterconfigs.MessageCodeCheckSystemdService: // 检查Systemd服务
			err = this.handleCheckSystemdService(message)
		default:
			err = this.handleUnknownMessage(message)
		}
		if err != nil {
			remotelogs.Error("API_STREAM", "handle message failed: "+err.Error())
		}
	}

	return nil
}

// 连接API节点成功
func (this *APIStream) handleConnectedAPINode(message *pb.ReportNodeStreamMessage) error {
	// 更改连接的APINode信息
	if len(message.DataJSON) == 0 {
		return nil
	}
	msg := &messageconfigs.ConnectedAPINodeMessage{}
	err := json.Unmarshal(message.DataJSON, msg)
	if err != nil {
		return errors.Wrap(err)
	}

	remotelogs.Println("API_STREAM", "connected to api node '"+strconv.FormatInt(msg.APINodeId, 10)+"'")
	return nil
}

// 处理配置变化
func (this *APIStream) handleNewNodeTask(message *pb.ReportNodeStreamMessage) error {
	select {
	case nodeTaskNotify <- true:
	default:

	}
	this.replyOk(message.RequestId, "ok")
	return nil
}

// 检查Systemd服务
func (this *APIStream) handleCheckSystemdService(message *pb.ReportNodeStreamMessage) error {
	systemctl, err := exec.LookPath("systemctl")
	if err != nil {
		this.replyFail(message.RequestId, "'systemctl' not found")
		return nil
	}
	if len(systemctl) == 0 {
		this.replyFail(message.RequestId, "'systemctl' not found")
		return nil
	}

	cmd := utils.NewCommandExecutor()
	shortName := teaconst.SystemdServiceName
	cmd.Add(systemctl, "is-enabled", shortName)
	output, err := cmd.Run()
	if err != nil {
		this.replyFail(message.RequestId, "'systemctl' command error: "+err.Error())
		return nil
	}
	if output == "enabled" {
		this.replyOk(message.RequestId, "ok")
	} else {
		this.replyFail(message.RequestId, "not installed")
	}
	return nil
}

// 处理未知消息
func (this *APIStream) handleUnknownMessage(message *pb.ReportNodeStreamMessage) error {
	this.replyFail(message.RequestId, "unknown message code '"+message.Code+"'")
	return nil
}

// 回复失败
func (this *APIStream) replyFail(requestId int64, message string) {
	_ = this.stream.Send(&pb.ReportNodeStreamMessage{RequestId: requestId, IsOk: false, Message: message})
}

// 回复成功
func (this *APIStream) replyOk(requestId int64, message string) {
	_ = this.stream.Send(&pb.ReportNodeStreamMessage{RequestId: requestId, IsOk: true, Message: message})
}