fc98c0919f
* new geoip method * report geoip separately, fix server creation & deletion bugs * fix struct tag * fix write name * remove deleteion list * remove rpc realip header * Revert "remove rpc realip header" This reverts commit 8a5f86cf2d7df87f28cfa2a3b3430f449dd6ed73.
70 lines
1.7 KiB
Go
70 lines
1.7 KiB
Go
package rpc
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
|
|
petname "github.com/dustinkirkland/golang-petname"
|
|
"github.com/hashicorp/go-uuid"
|
|
"google.golang.org/grpc/codes"
|
|
"google.golang.org/grpc/metadata"
|
|
"google.golang.org/grpc/status"
|
|
|
|
"github.com/naiba/nezha/model"
|
|
"github.com/naiba/nezha/service/singleton"
|
|
)
|
|
|
|
type authHandler struct {
|
|
ClientSecret string
|
|
ClientUUID string
|
|
}
|
|
|
|
func (a *authHandler) Check(ctx context.Context) (uint64, error) {
|
|
md, ok := metadata.FromIncomingContext(ctx)
|
|
if !ok {
|
|
return 0, status.Errorf(codes.Unauthenticated, "获取 metaData 失败")
|
|
}
|
|
|
|
//realIp := ctx.Value(model.CtxKeyRealIP{})
|
|
//log.Printf("bingo rpc realIp: %s, metadata: %v", realIp, md)
|
|
|
|
var clientSecret string
|
|
if value, ok := md["client_secret"]; ok {
|
|
clientSecret = value[0]
|
|
}
|
|
|
|
if clientSecret != singleton.Conf.AgentSecretKey {
|
|
return 0, status.Error(codes.Unauthenticated, "客户端认证失败")
|
|
}
|
|
|
|
var clientUUID string
|
|
if value, ok := md["client_uuid"]; ok {
|
|
clientUUID = value[0]
|
|
}
|
|
|
|
if _, err := uuid.ParseUUID(clientUUID); err != nil {
|
|
return 0, status.Error(codes.Unauthenticated, "客户端 UUID 不合法")
|
|
}
|
|
|
|
singleton.ServerLock.RLock()
|
|
defer singleton.ServerLock.RUnlock()
|
|
|
|
clientID, hasID := singleton.ServerUUIDToID[clientUUID]
|
|
if !hasID {
|
|
s := model.Server{UUID: clientUUID, Name: petname.Generate(2, "-")}
|
|
if err := singleton.DB.Create(&s).Error; err != nil {
|
|
return 0, status.Error(codes.Unauthenticated, err.Error())
|
|
}
|
|
s.Host = &model.Host{}
|
|
s.State = &model.HostState{}
|
|
s.TaskCloseLock = new(sync.Mutex)
|
|
// generate a random silly server name
|
|
singleton.ServerList[s.ID] = &s
|
|
singleton.ServerUUIDToID[clientUUID] = s.ID
|
|
singleton.ReSortServer()
|
|
clientID = s.ID
|
|
}
|
|
|
|
return clientID, nil
|
|
}
|