sendnrw/netbird
2
0
Fork 0
mirror of https://github.com/netbirdio/netbird.git synced 2026-04-21 01:36:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

add peer key rate limit

Browse Source
This commit is contained in:
Maycon Santos
2025-05-08 00:17:10 +02:00
parent 921b5606ce
commit 8bfab0d6dd
1 changed files with 17 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
management/server/grpcserver.go
View File

@@ -51,8 +51,10 @@ type GRPCServer struct {
peerLocks sync.Map
authManager auth.Manager
syncLimiter *rate.Limiter
loginLimiter *rate.Limiter
syncLimiter *rate.Limiter
loginLimiter *rate.Limiter
loginLimiterStore sync.Map
loginPeerLimit rate.Limit
}
// NewServer creates a new Management server
@@ -86,6 +88,10 @@ func NewServer(
if loginRatePerS == 0 || err != nil {
loginRatePerS = 200
}
loginPeerRatePerS, err := strconv.Atoi(os.Getenv("NB_LOGIN_PEER_RATE_PER_M"))
if loginPeerRatePerS == 0 || err != nil {
loginPeerRatePerS = 200
}
log.WithContext(ctx).Infof("login rate limit set to %d/min", loginRatePerS)
syncRatePerS, err := strconv.Atoi(os.Getenv("NB_SYNC_RATE_PER_M"))
@@ -107,6 +113,7 @@ func NewServer(
ephemeralManager: ephemeralManager,
syncLimiter: rate.NewLimiter(rate.Every(time.Minute/time.Duration(syncRatePerS)), 1),
loginLimiter: rate.NewLimiter(rate.Every(time.Minute/time.Duration(loginRatePerS)), 1),
loginPeerLimit: rate.Every(time.Minute / time.Duration(loginPeerRatePerS)),
}, nil
}
@@ -440,10 +447,17 @@ func (s *GRPCServer) parseRequest(ctx context.Context, req *proto.EncryptedMessa
// In case it isn't, the endpoint checks whether setup key is provided within the request and tries to register a peer.
// In case of the successful registration login is also successful
func (s *GRPCServer) Login(ctx context.Context, req *proto.EncryptedMessage) (*proto.EncryptedMessage, error) {
if !s.loginLimiter.Allow() {
limiter, _ := s.loginLimiterStore.LoadOrStore(req.WgPubKey, rate.NewLimiter(s.loginPeerLimit, 1))
if !limiter.(*rate.Limiter).Allow() {
return nil, status.Errorf(codes.Internal, "temp rate limit reached")
}
if os.Getenv("ENABLE_LOGIN_RATE_LIMIT") == "true" {
if !s.loginLimiter.Allow() {
return nil, status.Errorf(codes.Internal, "temp rate limit reached")
}
}
reqStart := time.Now()
defer func() {
if s.appMetrics != nil {