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

improve: use stream reduce auth check time

Browse Source
This commit is contained in:
naiba
2024-11-23 12:43:02 +08:00
parent cd42b1b9d5
commit 22738b6244
10 changed files with 249 additions and 141 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
cmd/dashboard/rpc/rpc.go
View File

@@ -20,17 +20,26 @@ import (
)
func ServeRPC() *grpc.Server {
server := grpc.NewServer(grpc.UnaryInterceptor(getRealIp))
server := grpc.NewServer(grpc.ChainUnaryInterceptor(getRealIp, waf))
rpcService.NezhaHandlerSingleton = rpcService.NewNezhaHandler()
proto.RegisterNezhaServiceServer(server, rpcService.NezhaHandlerSingleton)
return server
}
func waf(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
if err := model.CheckIP(singleton.DB, ctx.Value(model.CtxKeyRealIP{}).(string)); err != nil {
return nil, err
}
return handler(ctx, req)
}
func getRealIp(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
if singleton.Conf.RealIPHeader == "" {
return handler(ctx, req)
}
var ip string
if singleton.Conf.RealIPHeader == model.ConfigUsePeerIP {
p, ok := peer.FromContext(ctx)
if !ok {
@@ -40,18 +49,19 @@ func getRealIp(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo,
if err != nil {
return nil, err
}
ctx = context.WithValue(ctx, model.CtxKeyRealIP{}, addrPort.Addr().String())
return handler(ctx, req)
ip = addrPort.Addr().String()
} else {
vals := metadata.ValueFromIncomingContext(ctx, singleton.Conf.RealIPHeader)
if len(vals) == 0 {
return nil, fmt.Errorf("real ip header not found")
}
var err error
ip, err = utils.GetIPFromHeader(vals[0])
if err != nil {
return nil, err
}
}
vals := metadata.ValueFromIncomingContext(ctx, singleton.Conf.RealIPHeader)
if len(vals) == 0 {
return nil, fmt.Errorf("real ip header not found")
}
ip, err := utils.GetIPFromHeader(vals[0])
if err != nil {
return nil, err
}
ctx = context.WithValue(ctx, model.CtxKeyRealIP{}, ip)
return handler(ctx, req)
}