Buriburizaemon/nezha_domains
1
0
Fork 0
mirror of https://github.com/Buriburizaem0n/nezha_domains.git synced 2026-02-04 12:40:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

展示服务器列表

Browse Source
This commit is contained in:
奶爸
2019-12-09 18:14:31 +08:00
parent 58277ba0b6
commit bfe6d48927
18 changed files with 173 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
service/rpc/auth.go
View File

@@ -2,9 +2,7 @@ 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"
@@ -19,7 +17,7 @@ type AuthHandler struct {
// GetRequestMetadata ..
func (a *AuthHandler) GetRequestMetadata(ctx context.Context, uri ...string) (map[string]string, error) {
return map[string]string{"app_key": a.ClientID, "app_secret": a.ClientSecret}, nil
return map[string]string{"client_id": a.ClientID, "client_secret": a.ClientSecret}, nil
}
// RequireTransportSecurity ..
@@ -28,26 +26,27 @@ func (a *AuthHandler) RequireTransportSecurity() bool {
}
// Check ..
func (a *AuthHandler) Check(ctx context.Context) error {
func (a *AuthHandler) Check(ctx context.Context) (clientID string, err error) {
md, ok := metadata.FromIncomingContext(ctx)
if !ok {
return status.Errorf(codes.Unauthenticated, "获取 metaData 失败")
err = status.Errorf(codes.Unauthenticated, "获取 metaData 失败")
return
}
var (
ClientID string
ClientSecret string
clientSecret string
)
if value, ok := md["app_key"]; ok {
ClientID = value[0]
if value, ok := md["client_id"]; ok {
clientID = value[0]
}
if value, ok := md["app_secret"]; ok {
ClientSecret = value[0]
if value, ok := md["client_secret"]; ok {
clientSecret = value[0]
}
if _, ok := dao.Cache.Get(fmt.Sprintf("%s%s%s", model.CtxKeyServer, ClientID, ClientSecret)); !ok {
return status.Errorf(codes.Unauthenticated, "客户端认证失败")
dao.ServerLock.RLock()
defer dao.ServerLock.RUnlock()
if server, has := dao.ServerList[clientID]; !has || server.Secret != clientSecret {
err = status.Errorf(codes.Unauthenticated, "客户端认证失败")
}
return nil
return
}

26
service/rpc/nezha.go
View File

@@ -2,11 +2,11 @@ package rpc
import (
"context"
"fmt"
"log"
"github.com/p14yground/nezha/model"
pb "github.com/p14yground/nezha/proto"
"github.com/p14yground/nezha/service/dao"
)
// NezhaHandler ..
@@ -16,20 +16,26 @@ type NezhaHandler struct {
// ReportState ..
func (s *NezhaHandler) ReportState(c context.Context, r *pb.State) (*pb.Receipt, error) {
if err := s.Auth.Check(c); err != nil {
var clientID string
var err error
if clientID, err = s.Auth.Check(c); err != nil {
return nil, err
}
fmt.Printf("ReportState receive: %s\n", r)
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 {
defer log.Println("Heartbeat exit")
if err := s.Auth.Check(stream.Context()); err != nil {
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{
err = stream.Send(&pb.Command{
Type: model.MTReportState,
})
if err != nil {
@@ -40,9 +46,13 @@ func (s *NezhaHandler) Heartbeat(r *pb.Beat, stream pb.NezhaService_HeartbeatSer
// Register ..
func (s *NezhaHandler) Register(c context.Context, r *pb.Host) (*pb.Receipt, error) {
if err := s.Auth.Check(c); err != nil {
var clientID string
var err error
if clientID, err = s.Auth.Check(c); err != nil {
return nil, err
}
fmt.Printf("Register receive: %s\n", r)
dao.ServerLock.Lock()
defer dao.ServerLock.Unlock()
dao.ServerList[clientID].Host = model.PB2Host(r)
return &pb.Receipt{Proced: true}, nil
}