客户端授权验证

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
+	clientID     string
-	appSecret string
+	clientSecret string
+	debug        bool
 )
 
 func main() {
-	rootCmd.PersistentFlags().StringVarP(&appKey, "id", "i", "", "客户端ID")
+	rootCmd.PersistentFlags().StringVarP(&clientID, "id", "i", "", "客户端ID")
-	rootCmd.PersistentFlags().StringVarP(&appSecret, "secret", "p", "", "客户端Secret")
+	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,
+		ClientID:     clientID,
-		AppSecret: appSecret,
+		ClientSecret: clientSecret,
 	}
-	conn, err := grpc.Dial(":5555", grpc.WithInsecure(), grpc.WithPerRPCCredentials(&auth))
+	go reportState()
-	if err != nil {
+	var err error
-		panic(err)
+	var conn *grpc.ClientConn
-	}
+	var hc pb.NezhaService_HeartbeatClient
-	defer conn.Close()
+	for {
-	client := pb.NewNezhaServiceClient(conn)
+		log.Println("Try to reconnect ...")
-	ctx := context.Background()
+		time.Sleep(time.Second * 5)
-
+		conn, err = grpc.Dial(":5555", grpc.WithInsecure(), grpc.WithPerRPCCredentials(&auth))
-	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())
 		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",
+			ClientID:     "naiba",
-			AppSecret: "123456",
+			ClientSecret: "123456",
 		},
 	})
 	listen, err := net.Listen("tcp", ":5555")

model/common.go

@@ -7,6 +7,8 @@ const CtxKeyIsUserLogin = "ckiul"
 
 // CtxKeyOauth2State ..
 const CtxKeyOauth2State = "cko2s"
+// CtxKeyServer ..
+const CtxKeyServer = "cks"
 
 // Common ..
 type Common struct {

model/monitor.go

@@ -2,6 +2,12 @@ package model
 
 import pb "github.com/p14yground/nezha/proto"
 
+const (
+	_ = iota
+	// MTReportState ..
+	MTReportState
+)
+
 // State ..
 type State struct {
 	CPU       float64
@@ -39,6 +45,8 @@ type Host struct {
 	Virtualization  string
 	Uptime          string
 	BootTime        string
+	IP              string
+	CountryCode     string
 	Version         string
 }
 
@@ -52,6 +60,8 @@ func (h *Host) PB() *pb.Host {
 		Virtualization:  h.Virtualization,
 		Uptime:          h.Uptime,
 		BootTime:        h.BootTime,
+		Ip:              h.IP,
+		CountryCode:     h.CountryCode,
 		Version:         h.Version,
 	}
 }

proto/nezha.proto

@@ -16,7 +16,9 @@ message Host {
     string virtualization = 5;
     string uptime = 6;
     string boot_time = 7;
-    string version = 8;
+    string ip = 8;
+    string country_code = 9;
+    string version = 10;
 }
 
 message State {

resource/static/main.css

@@ -4,8 +4,19 @@
     }
 }
 
+.nb-container {
+    padding-top: 75px;
+    min-height: 80%;
+}
+
+.footer.segment {
+    margin-top: 40px !important;
+    padding-top: 40px;
+    padding-bottom: 40px;
+}
+
 .login-form {
-    height: 100%;
+    height: 82%;
 }
 
 .login-form .column {
@@ -16,12 +27,29 @@
     height: 100%;
 }
 
-.nb-container {
+.status.cards .header>.icon {
-    margin-top: 75px;
+    float: right;
+    margin-right: 0;
 }
 
-.footer.segment {
+.status.cards .wide.column {
-    margin-top: 40px !important;
+    padding-top: 0 !important;
-    padding-top: 40px;
+    padding-bottom: 0 !important;
-    padding-bottom: 40px;
+    height: 2rem !important;
+}
+
+.status.cards .three.wide.column {
+    padding-right: 0 !important;
+}
+
+.status.cards .wide.column:nth-child(1) {
+    margin-top: 1rem !important;
+}
+
+.status.cards .wide.column:nth-child(2) {
+    margin-top: 1rem !important;
+}
+
+.status.cards .description {
+    padding-bottom: 1rem !important;
 }

resource/static/main.js

@@ -1,11 +1,12 @@
 $('.ui.checkbox').checkbox();
+$('.yellow.info.circle.icon').popup();
 
 const confirmBtn = $('.mini.confirm.modal .positive.button')
 function showConfirm(title, content, callFn, extData) {
     const modal = $('.mini.confirm.modal')
     modal.children('.header').text(title)
     modal.children('.content').text(content)
-    if (confirmBtn.hasClass('loading')) {
+    if (confirmBtn.hasCslass('loading')) {
         return false
     }
     modal.modal({

resource/template/page/home.html

@@ -3,17 +3,40 @@
 {{template "common/menu" .}}
 <div class="nb-container">
     <div class="ui container">
-        <div class="ui four cards">
+        <div class="ui four status cards">
             {{range $server := .Servers}}
             <div class="card">
                 <div class="content">
-                    <div class="header">{{.Name}}</div>
+                    <div class="header"><i class="ae flag"></i>{{.Name}}
+                        <i data-html="<div class='header'>User Rating</div><div class='content'><div class='ui star rating'><i class='active icon'></i><i class='active icon'></i><i class='active icon'></i><i class='active icon'></i><i class='active icon'></i></div></div>"
+                            class="yellow info circle icon"></i>
+                    </div>
                     <div class="description">
-                        <div class="ui active progress">
+                        <div class="ui grid">
-                            <div class="bar">
+                            <div class="three wide column">CPU</div>
-                                <div class="progress"></div>
+                            <div class="thirteen wide column">
+                                <div class="ui active progress">
+                                    <div class="bar">
+                                        <div class="progress"></div>
+                                    </div>
+                                </div>
+                            </div>
+                            <div class="three wide column">硬盘</div>
+                            <div class="thirteen wide column">
+                                <div class="ui active progress">
+                                    <div class="bar">
+                                        <div class="progress"></div>
+                                    </div>
+                                </div>
+                            </div>
+                            <div class="three wide column">内存</div>
+                            <div class="thirteen wide column">
+                                <div class="ui active progress">
+                                    <div class="bar">
+                                        <div class="progress"></div>
+                                    </div>
+                                </div>
                             </div>
-                            <div class="label">CPU</div>
                         </div>
                     </div>
                 </div>

service/monitor/monitor.go

@@ -1,7 +1,11 @@
 package monitor
 
 import (
+	"encoding/json"
 	"fmt"
+	"io/ioutil"
+	"net/http"
+	"strings"
 	"time"
 
 	"github.com/shirou/gopsutil/cpu"
@@ -13,6 +17,11 @@ import (
 	"github.com/p14yground/nezha/model"
 )
 
+type ipDotSbGeoIP struct {
+	CountryCode string
+	IP          string
+}
+
 // GetHost ..
 func GetHost() *model.Host {
 	hi, _ := host.Info()
@@ -21,6 +30,16 @@ func GetHost() *model.Host {
 	for i := 0; i < len(ci); i++ {
 		cpus = append(cpus, fmt.Sprintf("%v-%vC%vT", ci[i].ModelName, ci[i].Cores, ci[i].Stepping))
 	}
+	ip := ipDotSbGeoIP{
+		IP:          "127.0.0.1",
+		CountryCode: "cn",
+	}
+	resp, err := http.Get("https://api.ip.sb/geoip")
+	if err == nil {
+		defer resp.Body.Close()
+		body, _ := ioutil.ReadAll(resp.Body)
+		json.Unmarshal(body, &ip)
+	}
 	return &model.Host{
 		Platform:        hi.OS,
 		PlatformVersion: hi.PlatformVersion,
@@ -29,6 +48,8 @@ func GetHost() *model.Host {
 		Virtualization:  hi.VirtualizationSystem,
 		Uptime:          fmt.Sprintf("%v", hi.Uptime),
 		BootTime:        fmt.Sprintf("%v", hi.BootTime),
+		IP:              ip.IP,
+		CountryCode:     strings.ToLower(ip.CountryCode),
 	}
 }
 

service/rpc/auth.go

@@ -2,7 +2,10 @@ package rpc
 
 import (
 	"context"
+	"fmt"
 
+	"github.com/p14yground/nezha/model"
+	"github.com/p14yground/nezha/service/dao"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
@@ -10,51 +13,41 @@ import (
 
 // AuthHandler ..
 type AuthHandler struct {
-	AppKey    string
+	ClientID     string
-	AppSecret string
+	ClientSecret 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
+	return map[string]string{"app_key": a.ClientID, "app_secret": a.ClientSecret}, nil
 }
 
 // RequireTransportSecurity ..
 func (a *AuthHandler) RequireTransportSecurity() bool {
-	return false
+	return !dao.Conf.Debug
 }
 
 // Check ..
 func (a *AuthHandler) Check(ctx context.Context) error {
 	md, ok := metadata.FromIncomingContext(ctx)
 	if !ok {
-		return status.Errorf(codes.Unauthenticated, "metadata.FromIncomingContext err")
+		return status.Errorf(codes.Unauthenticated, "获取 metaData 失败")
 	}
 
 	var (
-		AppKey    string
+		ClientID     string
-		AppSecret string
+		ClientSecret string
 	)
 	if value, ok := md["app_key"]; ok {
-		AppKey = value[0]
+		ClientID = value[0]
 	}
 	if value, ok := md["app_secret"]; ok {
-		AppSecret = value[0]
+		ClientSecret = value[0]
 	}
 
-	if AppKey != a.GetAppKey() || AppSecret != a.GetAppSecret() {
+	if _, ok := dao.Cache.Get(fmt.Sprintf("%s%s%s", model.CtxKeyServer, ClientID, ClientSecret)); !ok {
-		return status.Errorf(codes.Unauthenticated, "invalid token")
+		return status.Errorf(codes.Unauthenticated, "客户端认证失败")
 	}
 
 	return nil
 }
-
-// GetAppKey ..
-func (a *AuthHandler) GetAppKey() string {
-	return a.AppKey
-}
-
-// GetAppSecret ..
-func (a *AuthHandler) GetAppSecret() string {
-	return a.AppSecret
-}

service/rpc/nezha.go

@@ -3,7 +3,9 @@ package rpc
 import (
 	"context"
 	"fmt"
+	"log"
 
+	"github.com/p14yground/nezha/model"
 	pb "github.com/p14yground/nezha/proto"
 )
 
@@ -23,11 +25,17 @@ func (s *NezhaHandler) ReportState(c context.Context, r *pb.State) (*pb.Receipt,
 
 // Heartbeat ..
 func (s *NezhaHandler) Heartbeat(r *pb.Beat, stream pb.NezhaService_HeartbeatServer) error {
+	defer log.Println("Heartbeat exit")
 	if err := s.Auth.Check(stream.Context()); err != nil {
 		return err
 	}
-	fmt.Printf("ReportState receive: %s\n", r)
+	err := stream.Send(&pb.Command{
-	return nil
+		Type: model.MTReportState,
+	})
+	if err != nil {
+		log.Printf("Heartbeat stream.Send err:%v", err)
+	}
+	select {}
 }
 
 // Register ..