nezha/service/rpc/nezha.go

package rpc

import (
	"context"
	"log"

	"github.com/p14yground/nezha/model"
	pb "github.com/p14yground/nezha/proto"
	"github.com/p14yground/nezha/service/dao"
)

// NezhaHandler ..
type NezhaHandler struct {
	Auth *AuthHandler
}

// ReportState ..
func (s *NezhaHandler) ReportState(c context.Context, r *pb.State) (*pb.Receipt, error) {
	var clientID string
	var err error
	if clientID, err = s.Auth.Check(c); err != nil {
		return nil, err
	}
	dao.ServerLock.Lock()
	defer dao.ServerLock.Unlock()
	dao.ServerList[clientID].State = model.PB2State(r)
	return &pb.Receipt{Proced: true}, nil
}

// Heartbeat ..
func (s *NezhaHandler) Heartbeat(r *pb.Beat, stream pb.NezhaService_HeartbeatServer) error {
	var clientID string
	var err error
	defer log.Printf("Heartbeat exit server:%v err:%v", clientID, err)
	if clientID, err = s.Auth.Check(stream.Context()); err != nil {
		return err
	}
	err = stream.Send(&pb.Command{
		Type: model.MTReportState,
	})
	if err != nil {
		log.Printf("Heartbeat stream.Send err:%v", err)
	}
	select {}
}

// Register ..
func (s *NezhaHandler) Register(c context.Context, r *pb.Host) (*pb.Receipt, error) {
	var clientID string
	var err error
	if clientID, err = s.Auth.Check(c); err != nil {
		return nil, err
	}
	dao.ServerLock.Lock()
	defer dao.ServerLock.Unlock()
	dao.ServerList[clientID].Host = model.PB2Host(r)
	return &pb.Receipt{Proced: true}, nil
}