gRPC 上报系统信息

cmd/agent/main.go

@@ -2,33 +2,21 @@ package main
 
 import (
 	"context"
+	"fmt"
 	"log"
+	"time"
 
 	"google.golang.org/grpc"
 
 	pb "github.com/p14yground/nezha/proto"
+	"github.com/p14yground/nezha/service/handler"
+	"github.com/p14yground/nezha/service/monitor"
 )
 
-// Auth ..
-type Auth struct {
-	AppKey    string
-	AppSecret string
-}
-
-// GetRequestMetadata ..
-func (a *Auth) GetRequestMetadata(ctx context.Context, uri ...string) (map[string]string, error) {
-	return map[string]string{"app_key": a.AppKey, "app_secret": a.AppSecret}, nil
-}
-
-// RequireTransportSecurity ..
-func (a *Auth) RequireTransportSecurity() bool {
-	return false
-}
-
 func main() {
-	auth := Auth{
+	auth := handler.AuthHandler{
 		AppKey:    "naiba",
-		AppSecret: "nbsecret",
+		AppSecret: "123456",
 	}
 	conn, err := grpc.Dial(":5555", grpc.WithInsecure(), grpc.WithPerRPCCredentials(&auth))
 	if err != nil {
@@ -36,13 +24,27 @@ func main() {
 	}
 	defer conn.Close()
 	client := pb.NewNezhaServiceClient(conn)
+	ctx := context.Background()
+
+	resp, err := client.Register(ctx, monitor.GetHost().PB())
+	if err != nil {
+		log.Printf("client.Register err: %v", err)
+	}
+	log.Printf("Register resp: %s", resp)
+
+	hc, err := client.Heartbeat(ctx, &pb.Beat{
+		Timestamp: fmt.Sprintf("%v", time.Now()),
+	})
+	if err != nil {
+		log.Printf("client.Register err: %v", err)
+	}
+	log.Printf("Register resp: %s", hc)
 
 	for i := 0; i < 3; i++ {
-		resp, err := client.ReportState(context.Background(), &pb.State{})
+		resp, err := client.ReportState(ctx, monitor.GetState(3).PB())
 		if err != nil {
-			log.Fatalf("client.Search err: %v", err)
+			log.Printf("client.ReportState err: %v", err)
 		}
-
+		log.Printf("ReportState resp: %s", resp)
-		log.Printf("resp: %s", resp)
 	}
 }

cmd/dashboard/main.go

@@ -1,35 +1,22 @@
 package main
 
 import (
-	"context"
-	"fmt"
 	"net"
 
 	"google.golang.org/grpc"
-	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/metadata"
-	"google.golang.org/grpc/status"
 
 	pb "github.com/p14yground/nezha/proto"
+	"github.com/p14yground/nezha/service/handler"
 )
 
-// NezhaService ..
-type NezhaService struct {
-	auth *Auth
-}
-
-// ReportState ..
-func (s *NezhaService) ReportState(ctx context.Context, r *pb.State) (*pb.Receipt, error) {
-	if err := s.auth.Check(ctx); err != nil {
-		return nil, err
-	}
-	fmt.Printf("receive: %s\n", r)
-	return &pb.Receipt{}, nil
-}
-
 func main() {
 	server := grpc.NewServer()
-	pb.RegisterNezhaServiceServer(server, &NezhaService{})
+	pb.RegisterNezhaServiceServer(server, &handler.NezhaHandler{
+		Auth: &handler.AuthHandler{
+			AppKey:    "naiba",
+			AppSecret: "123456",
+		},
+	})
 
 	lis, err := net.Listen("tcp", ":5555")
 	if err != nil {
@@ -38,44 +25,3 @@ func main() {
 
 	server.Serve(lis)
 }
-
-// Auth ..
-type Auth struct {
-	appKey    string
-	appSecret string
-}
-
-// Check ..
-func (a *Auth) Check(ctx context.Context) error {
-	md, ok := metadata.FromIncomingContext(ctx)
-	if !ok {
-		return status.Errorf(codes.Unauthenticated, "metadata.FromIncomingContext err")
-	}
-
-	var (
-		appKey    string
-		appSecret string
-	)
-	if value, ok := md["app_key"]; ok {
-		appKey = value[0]
-	}
-	if value, ok := md["app_secret"]; ok {
-		appSecret = value[0]
-	}
-
-	if appKey != a.GetAppKey() || appSecret != a.GetAppSecret() {
-		return status.Errorf(codes.Unauthenticated, "invalid token")
-	}
-
-	return nil
-}
-
-// GetAppKey ..
-func (a *Auth) GetAppKey() string {
-	return "naiba"
-}
-
-// GetAppSecret ..
-func (a *Auth) GetAppSecret() string {
-	return "nbsecret"
-}

cmd/playground/main.go

@@ -1,53 +1,12 @@
 package main
 
 import (
-	"fmt"
 	"log"
 	"os/exec"
-	"strconv"
-	"time"
-
-	"github.com/shirou/gopsutil/cpu"
-	"github.com/shirou/gopsutil/disk"
-	"github.com/shirou/gopsutil/host"
-	"github.com/shirou/gopsutil/mem"
-	"github.com/shirou/gopsutil/net"
 )
 
 func main() {
-	// Host info
+
-	hi, _ := host.Info()
-	fmt.Printf("「HostInfo」 platform:%v platformVersion:%v kernelArch:%v virtualizationSystem:%v uptime:%v boottime:%v\n", hi.OS, hi.PlatformVersion, hi.KernelArch, hi.VirtualizationSystem, hi.Uptime, hi.BootTime)
-	// Memory
-	mv, _ := mem.VirtualMemory()
-	ms, _ := mem.SwapMemory()
-	fmt.Printf("「VirtualMemory」 Total: %v, Free:%v, UsedPercent:%f%%\n", mv.Total, mv.Free, mv.UsedPercent)
-	fmt.Printf("「SwapMemory」 Total: %v, Free:%v, UsedPercent:%f%%\n", ms.Total, ms.Free, ms.UsedPercent)
-	// Disk
-	dparts, _ := disk.Partitions(false)
-	for _, part := range dparts {
-		fmt.Printf("「Disk」 %v\n", part)
-		u, _ := disk.Usage(part.Mountpoint)
-		fmt.Println("\t" + u.Path + "\t" + strconv.FormatFloat(u.UsedPercent, 'f', 2, 64) + "% full.")
-		fmt.Println("\t\tTotal: " + strconv.FormatUint(u.Total/1024/1024/1024, 10) + " GiB")
-		fmt.Println("\t\tFree:  " + strconv.FormatUint(u.Free/1024/1024/1024, 10) + " GiB")
-		fmt.Println("\t\tUsed:  " + strconv.FormatUint(u.Used/1024/1024/1024, 10) + " GiB")
-	}
-	// CPU
-	go func() {
-		cp, _ := cpu.Percent(time.Second*2, false)
-		ci, _ := cpu.Info()
-		for i := 0; i < len(ci); i++ {
-			fmt.Printf("「CPU」 %v core:%v step:%v", ci[i].ModelName, ci[i].Cores, ci[i].Stepping)
-		}
-		fmt.Printf(" percentIn2sec:%v%%\n", cp[0])
-	}()
-	// Network
-	nc, _ := net.IOCounters(true)
-	for _, ni := range nc {
-		fmt.Printf("「Net」%v\n", ni)
-	}
-	select {}
 }
 
 func cmdExec() {

model/monitor.go

@@ -0,0 +1,57 @@
+package model
+
+import pb "github.com/p14yground/nezha/proto"
+
+// State ..
+type State struct {
+	CPU       float64
+	MEMTotal  uint64
+	MEMUsed   uint64
+	SwapTotal uint64
+	SwapUsed  uint64
+	DiskTotal uint64
+	DiskUsed  uint64
+	NetIn     uint64
+	NetOut    uint64
+}
+
+// PB ..
+func (s *State) PB() *pb.State {
+	return &pb.State{
+		Cpu:       s.CPU,
+		MemTotal:  s.MEMTotal,
+		MemUsed:   s.MEMUsed,
+		SwapTotal: s.SwapTotal,
+		SwapUsed:  s.SwapUsed,
+		DiskTotal: s.DiskTotal,
+		DiskUsed:  s.DiskUsed,
+		NetIn:     s.NetIn,
+		NetOut:    s.NetOut,
+	}
+}
+
+// Host ..
+type Host struct {
+	Platform        string
+	PlatformVersion string
+	CPU             []string
+	Arch            string
+	Virtualization  string
+	Uptime          string
+	BootTime        string
+	Version         string
+}
+
+// PB ..
+func (h *Host) PB() *pb.Host {
+	return &pb.Host{
+		Platform:        h.Platform,
+		PlatformVersion: h.PlatformVersion,
+		Cpu:             h.CPU,
+		Arch:            h.Arch,
+		Virtualization:  h.Virtualization,
+		Uptime:          h.Uptime,
+		BootTime:        h.BootTime,
+		Version:         h.Version,
+	}
+}

proto/nezha.proto

@@ -3,11 +3,43 @@ syntax = "proto3";
 package proto;
 
 service NezhaService {
+    rpc Heartbeat(Beat)returns(stream Command){}
     rpc ReportState(State)returns(Receipt){}
+    rpc Register(Host)returns(Receipt){}
+}
+
+message Host {
+    string platform = 1;
+    string platform_version = 2;
+    repeated string cpu = 3;
+    string arch = 4;
+    string virtualization = 5;
+    string uptime = 6;
+    string boot_time = 7;
+    string version = 8;
 }
 
 message State {
+    double cpu = 1;
+    uint64 mem_total = 2;
+    uint64 mem_used = 3;
+    uint64 swap_total = 4;
+    uint64 swap_used = 5;
+    uint64 disk_total = 6;
+    uint64 disk_used = 7;
+    uint64 net_in = 8;
+    uint64 net_out = 9;
 }
 
 message Receipt{
+    bool proced = 1;
+}
+
+message Beat {
+    string timestamp = 1;
+}
+
+message Command {
+    uint64 type = 1;
+    string data = 2;
 }

service/handler/auth.go

@@ -0,0 +1,60 @@
+package handler
+
+import (
+	"context"
+
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
+)
+
+// AuthHandler ..
+type AuthHandler struct {
+	AppKey    string
+	AppSecret string
+}
+
+// GetRequestMetadata ..
+func (a *AuthHandler) GetRequestMetadata(ctx context.Context, uri ...string) (map[string]string, error) {
+	return map[string]string{"app_key": a.AppKey, "app_secret": a.AppSecret}, nil
+}
+
+// RequireTransportSecurity ..
+func (a *AuthHandler) RequireTransportSecurity() bool {
+	return false
+}
+
+// Check ..
+func (a *AuthHandler) Check(ctx context.Context) error {
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		return status.Errorf(codes.Unauthenticated, "metadata.FromIncomingContext err")
+	}
+
+	var (
+		AppKey    string
+		AppSecret string
+	)
+	if value, ok := md["app_key"]; ok {
+		AppKey = value[0]
+	}
+	if value, ok := md["app_secret"]; ok {
+		AppSecret = value[0]
+	}
+
+	if AppKey != a.GetAppKey() || AppSecret != a.GetAppSecret() {
+		return status.Errorf(codes.Unauthenticated, "invalid token")
+	}
+
+	return nil
+}
+
+// GetAppKey ..
+func (a *AuthHandler) GetAppKey() string {
+	return a.AppKey
+}
+
+// GetAppSecret ..
+func (a *AuthHandler) GetAppSecret() string {
+	return a.AppSecret
+}

service/handler/nezha.go

@@ -0,0 +1,40 @@
+package handler
+
+import (
+	"context"
+	"fmt"
+
+	pb "github.com/p14yground/nezha/proto"
+)
+
+// NezhaHandler ..
+type NezhaHandler struct {
+	Auth *AuthHandler
+}
+
+// ReportState ..
+func (s *NezhaHandler) ReportState(c context.Context, r *pb.State) (*pb.Receipt, error) {
+	if err := s.Auth.Check(c); err != nil {
+		return nil, err
+	}
+	fmt.Printf("ReportState receive: %s\n", r)
+	return &pb.Receipt{Proced: true}, nil
+}
+
+// Heartbeat ..
+func (s *NezhaHandler) Heartbeat(r *pb.Beat, stream pb.NezhaService_HeartbeatServer) error {
+	if err := s.Auth.Check(stream.Context()); err != nil {
+		return err
+	}
+	fmt.Printf("ReportState receive: %s\n", r)
+	return nil
+}
+
+// Register ..
+func (s *NezhaHandler) Register(c context.Context, r *pb.Host) (*pb.Receipt, error) {
+	if err := s.Auth.Check(c); err != nil {
+		return nil, err
+	}
+	fmt.Printf("Register receive: %s\n", r)
+	return &pb.Receipt{Proced: true}, nil
+}

service/monitor/monitor.go

@@ -0,0 +1,73 @@
+package monitor
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/shirou/gopsutil/cpu"
+	"github.com/shirou/gopsutil/disk"
+	"github.com/shirou/gopsutil/host"
+	"github.com/shirou/gopsutil/mem"
+	"github.com/shirou/gopsutil/net"
+
+	"github.com/p14yground/nezha/model"
+)
+
+// GetHost ..
+func GetHost() *model.Host {
+	hi, _ := host.Info()
+	var cpus []string
+	ci, _ := cpu.Info()
+	for i := 0; i < len(ci); i++ {
+		cpus = append(cpus, fmt.Sprintf("%v-%vC%vT", ci[i].ModelName, ci[i].Cores, ci[i].Stepping))
+	}
+	return &model.Host{
+		Platform:        hi.OS,
+		PlatformVersion: hi.PlatformVersion,
+		CPU:             cpus,
+		Arch:            hi.KernelArch,
+		Virtualization:  hi.VirtualizationSystem,
+		Uptime:          fmt.Sprintf("%v", hi.Uptime),
+		BootTime:        fmt.Sprintf("%v", hi.BootTime),
+	}
+}
+
+// GetState ..
+func GetState(delay uint64) *model.State {
+	// Memory
+	mv, _ := mem.VirtualMemory()
+	ms, _ := mem.SwapMemory()
+	// Disk
+	var diskTotal, diskUsed uint64
+	dparts, _ := disk.Partitions(true)
+	for _, part := range dparts {
+		u, _ := disk.Usage(part.Mountpoint)
+		diskTotal += u.Total
+		diskUsed += u.Used
+	}
+	// CPU
+	var cpuPercent float64
+	cp, err := cpu.Percent(time.Second*time.Duration(delay), false)
+	if err == nil {
+		cpuPercent = cp[0]
+
+	}
+	// Network
+	var netIn, netOut uint64
+	nc, err := net.IOCounters(false)
+	if err == nil {
+		netIn = nc[0].BytesRecv
+		netOut = nc[0].BytesSent
+	}
+	return &model.State{
+		CPU:       cpuPercent,
+		MEMTotal:  mv.Total,
+		MEMUsed:   mv.Used,
+		SwapTotal: ms.Total,
+		SwapUsed:  ms.Used,
+		DiskTotal: diskTotal,
+		DiskUsed:  diskUsed,
+		NetIn:     netIn,
+		NetOut:    netOut,
+	}
+}