From 4a3c782a3180feba35608132a1d9dbf23b65d5e6 Mon Sep 17 00:00:00 2001
From: Pascal Fischer <pascal@netbird.io>
Date: Wed, 7 May 2025 20:35:41 +0200
Subject: [PATCH] add rate limiter to login and sync on grpc

---
 management/server/grpcserver.go | 14 ++++++++++++++
 management/server/peer.go       |  8 --------
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/management/server/grpcserver.go b/management/server/grpcserver.go
index 43d35f643..d25684996 100644
--- a/management/server/grpcserver.go
+++ b/management/server/grpcserver.go
@@ -2,6 +2,7 @@ package server
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"net"
 	"net/netip"
@@ -13,6 +14,7 @@ import (
 	"github.com/golang/protobuf/ptypes/timestamp"
 	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/realip"
 	log "github.com/sirupsen/logrus"
+	"golang.org/x/time/rate"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/peer"
@@ -125,9 +127,15 @@ func getRealIP(ctx context.Context) net.IP {
 	return nil
 }
 
+var syncLimiter = rate.NewLimiter(rate.Every(time.Minute/300), 1)
+
 // Sync validates the existence of a connecting peer, sends an initial state (all available for the connecting peers) and
 // notifies the connected peer of any updates (e.g. new peers under the same account)
 func (s *GRPCServer) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_SyncServer) error {
+	if !syncLimiter.Allow() {
+		return errors.New("rate limit exceeded")
+	}
+
 	reqStart := time.Now()
 	if s.appMetrics != nil {
 		s.appMetrics.GRPCMetrics().CountSyncRequest()
@@ -411,11 +419,17 @@ func (s *GRPCServer) parseRequest(ctx context.Context, req *proto.EncryptedMessa
 	return peerKey, nil
 }
 
+var loginLimiter = rate.NewLimiter(rate.Every(time.Minute/300), 1)
+
 // Login endpoint first checks whether peer is registered under any account
 // In case it is, the login is successful
 // 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 !loginLimiter.Allow() {
+		return nil, errors.New("rate limit exceeded")
+	}
+
 	reqStart := time.Now()
 	defer func() {
 		if s.appMetrics != nil {
diff --git a/management/server/peer.go b/management/server/peer.go
index c07aefaa2..9ff80442e 100644
--- a/management/server/peer.go
+++ b/management/server/peer.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"crypto/sha256"
 	b64 "encoding/base64"
-	"errors"
 	"fmt"
 	"net"
 	"slices"
@@ -15,7 +14,6 @@ import (
 	"github.com/rs/xid"
 	log "github.com/sirupsen/logrus"
 	"golang.org/x/exp/maps"
-	"golang.org/x/time/rate"
 
 	"github.com/netbirdio/netbird/management/domain"
 	"github.com/netbirdio/netbird/management/server/geolocation"
@@ -782,17 +780,11 @@ func (am *DefaultAccountManager) handlePeerLoginNotFound(ctx context.Context, lo
 	return nil, nil, nil, status.Errorf(status.Internal, "failed while logging in peer")
 }
 
-var loginLimiter = rate.NewLimiter(rate.Every(time.Minute/500), 1)
-
 // LoginPeer logs in or registers a peer.
 // If peer doesn't exist the function checks whether a setup key or a user is present and registers a new peer if so.
 func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login types.PeerLogin) (*nbpeer.Peer, *types.NetworkMap, []*posture.Checks, error) {
 	accountID, err := am.Store.GetAccountIDByPeerPubKey(ctx, login.WireGuardPubKey)
 	if err != nil {
-		if !loginLimiter.Allow() {
-			return nil, nil, nil, errors.New("rate limit exceeded")
-		}
-
 		return am.handlePeerLoginNotFound(ctx, login, err)
 	}