1.4.5.2

EdgeReporter/internal/nodes/api_stream.go

@@ -0,0 +1,177 @@
+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})
+}