客户端授权验证

cmd/agent/main.go

@@ -3,6 +3,7 @@ package main
 import (
 	"context"
 	"fmt"
+	"io"
 	"log"
 	"os"
 	"time"
@@ -10,7 +11,9 @@ import (
 	"github.com/spf13/cobra"
 	"google.golang.org/grpc"
 
+	"github.com/p14yground/nezha/model"
 	pb "github.com/p14yground/nezha/proto"
+	"github.com/p14yground/nezha/service/dao"
 	"github.com/p14yground/nezha/service/monitor"
 	"github.com/p14yground/nezha/service/rpc"
 )
@@ -25,51 +28,89 @@ var (
 啦啦啦，啦啦啦，我是 mjj 小行家`,
 		Run: run,
 	}
-	appKey    string
-	appSecret string
+	clientID     string
+	clientSecret string
+	debug        bool
 )
 
 func main() {
-	rootCmd.PersistentFlags().StringVarP(&appKey, "id", "i", "", "客户端ID")
-	rootCmd.PersistentFlags().StringVarP(&appSecret, "secret", "p", "", "客户端Secret")
+	rootCmd.PersistentFlags().StringVarP(&clientID, "id", "i", "", "客户端ID")
+	rootCmd.PersistentFlags().StringVarP(&clientSecret, "secret", "p", "", "客户端Secret")
+	rootCmd.PersistentFlags().BoolVarP(&debug, "debug", "d", false, "开启Debug")
 	if err := rootCmd.Execute(); err != nil {
 		fmt.Println(err)
 		os.Exit(1)
 	}
 }
 
+var endReport time.Time
+var reporting bool
+var client pb.NezhaServiceClient
+var ctx = context.Background()
+
 func run(cmd *cobra.Command, args []string) {
+	dao.Conf = &model.Config{
+		Debug: debug,
+	}
 	auth := rpc.AuthHandler{
-		AppKey:    appKey,
-		AppSecret: appSecret,
+		ClientID:     clientID,
+		ClientSecret: clientSecret,
 	}
-	conn, err := grpc.Dial(":5555", grpc.WithInsecure(), grpc.WithPerRPCCredentials(&auth))
-	if err != nil {
-		panic(err)
-	}
-	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(ctx, monitor.GetState(3).PB())
+	go reportState()
+	var err error
+	var conn *grpc.ClientConn
+	var hc pb.NezhaService_HeartbeatClient
+	for {
+		log.Println("Try to reconnect ...")
+		time.Sleep(time.Second * 5)
+		conn, err = grpc.Dial(":5555", grpc.WithInsecure(), grpc.WithPerRPCCredentials(&auth))
 		if err != nil {
-			log.Printf("client.ReportState err: %v", err)
+			log.Printf("grpc.Dial err: %v", err)
+			continue
 		}
-		log.Printf("ReportState resp: %s", resp)
+		client = pb.NewNezhaServiceClient(conn)
+		// 第一步注册
+		client.Register(ctx, monitor.GetHost().PB())
+		hc, err = client.Heartbeat(ctx, &pb.Beat{
+			Timestamp: fmt.Sprintf("%v", time.Now()),
+		})
+		if err != nil {
+			log.Printf("client.Register err: %v", err)
+			continue
+		}
+		err = receiveCommand(hc)
+		log.Printf("receiveCommand exit to main: %v", err)
+	}
+}
+
+func receiveCommand(hc pb.NezhaService_HeartbeatClient) error {
+	var err error
+	var action *pb.Command
+	defer log.Printf("receiveCommand exit %v %v => %v", time.Now(), action, err)
+	for {
+		action, err = hc.Recv()
+		if err == io.EOF {
+			return nil
+		}
+		if err != nil {
+			return err
+		}
+		switch action.GetType() {
+		case model.MTReportState:
+			endReport = time.Now().Add(time.Second * 10)
+		default:
+			log.Printf("Unknown action: %v", action)
+		}
+	}
+}
+
+func reportState() {
+	var err error
+	defer log.Printf("reportState exit %v %v => %v", endReport, time.Now(), err)
+	for {
+		if endReport.After(time.Now()) {
+			_, err = client.ReportState(ctx, monitor.GetState(0).PB())
+		}
+		time.Sleep(time.Second)
 	}
 }

cmd/dashboard/controller/member_api.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/gin-gonic/gin"
 	"github.com/naiba/com"
+	"github.com/patrickmn/go-cache"
 
 	"github.com/p14yground/nezha/model"
 	"github.com/p14yground/nezha/pkg/mygin"
@@ -37,9 +38,9 @@ type serverForm struct {
 
 func (ma *memberAPI) addServer(c *gin.Context) {
 	var sf serverForm
+	var s model.Server
 	err := c.ShouldBindJSON(&sf)
 	if err == nil {
-		var s model.Server
 		s.Name = sf.Name
 		s.Secret = com.MD5(fmt.Sprintf("%s%s%d", time.Now(), sf.Name, dao.Admin.ID))
 		s.Secret = s.Secret[:10]
@@ -52,6 +53,7 @@ func (ma *memberAPI) addServer(c *gin.Context) {
 		})
 		return
 	}
+	dao.Cache.Set(fmt.Sprintf("%s%d%s", model.CtxKeyServer, s.ID, s.Secret), s, cache.NoExpiration)
 	c.JSON(http.StatusOK, model.Response{
 		Code: http.StatusOK,
 	})

cmd/dashboard/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"fmt"
 	"time"
 
 	"github.com/jinzhu/gorm"
@@ -36,6 +37,12 @@ func init() {
 
 func initDB() {
 	dao.DB.AutoMigrate(model.Server{})
+	// load cache
+	var servers []model.Server
+	dao.DB.Find(&servers)
+	for _, s := range servers {
+		dao.Cache.Set(fmt.Sprintf("%s%d%s", model.CtxKeyServer, s.ID, s.Secret), s, cache.NoExpiration)
+	}
 }
 
 func main() {

cmd/dashboard/rpc/rpc.go

@@ -14,8 +14,8 @@ func ServeRPC() {
 	server := grpc.NewServer()
 	pb.RegisterNezhaServiceServer(server, &rpcService.NezhaHandler{
 		Auth: &rpcService.AuthHandler{
-			AppKey:    "naiba",
-			AppSecret: "123456",
+			ClientID:     "naiba",
+			ClientSecret: "123456",
 		},
 	})
 	listen, err := net.Listen("tcp", ":5555")