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}) }