From 86fa1eaa1619c3e8f2e870c0fcaa3fd5070c8d3b Mon Sep 17 00:00:00 2001
From: Maycon Santos <mlsmaycon@gmail.com>
Date: Thu, 8 May 2025 10:43:51 +0200
Subject: [PATCH] check global limiter before create peer limiter

---
 management/server/grpcserver.go | 44 ++++++++++++++++++++++++++-------
 1 file changed, 35 insertions(+), 9 deletions(-)

diff --git a/management/server/grpcserver.go b/management/server/grpcserver.go
index 0a685e64e..39dd97981 100644
--- a/management/server/grpcserver.go
+++ b/management/server/grpcserver.go
@@ -454,19 +454,45 @@ 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) {
-	limiter, _ := s.loginLimiterStore.LoadOrStore(req.WgPubKey, rate.NewLimiter(s.loginPeerLimit, 1))
-	if !limiter.(*rate.Limiter).Allow() {
-		time.Sleep(time.Millisecond * time.Duration(rand.IntN(10)*100))
-		log.WithContext(ctx).Warnf("rate limit exceeded for %s", req.WgPubKey)
-		return nil, status.Errorf(codes.Internal, "temp rate limit reached")
-	}
-
-	if os.Getenv("ENABLE_LOGIN_RATE_LIMIT") == "true" {
+	limiterIface, ok := s.loginLimiterStore.Load(req.WgPubKey)
+	if !ok {
+		// Check global limiter before allowing a new peer limiter
 		if !s.loginLimiter.Allow() {
-			return nil, status.Errorf(codes.Internal, "temp rate limit reached")
+			return nil, fmt.Errorf("temp rate limit reached (global limit)")
+		}
+
+		// Create new limiter for this peer
+		newLimiter := rate.NewLimiter(s.loginPeerLimit, 1)
+		s.loginLimiterStore.Store(req.WgPubKey, newLimiter)
+
+		if !newLimiter.Allow() {
+			time.Sleep(time.Second + (time.Millisecond * time.Duration(rand.IntN(20)*100)))
+			log.WithContext(ctx).Warnf("rate limit exceeded for peer %s", req.WgPubKey)
+			return nil, fmt.Errorf("temp rate limit reached (new peer limit)")
+		}
+	} else {
+		// Use existing limiter for this peer
+		limiter := limiterIface.(*rate.Limiter)
+		if !limiter.Allow() {
+			time.Sleep(time.Second + (time.Millisecond * time.Duration(rand.IntN(20)*100)))
+			log.WithContext(ctx).Warnf("rate limit exceeded for peer %s", req.WgPubKey)
+			return nil, fmt.Errorf("temp rate limit reached (peer limit)")
 		}
 	}
 
+	//limiter, _ := s.loginLimiterStore.LoadOrStore(req.WgPubKey, rate.NewLimiter(s.loginPeerLimit, 1))
+	//if !limiter.(*rate.Limiter).Allow() {
+	//	time.Sleep(time.Millisecond * time.Duration(rand.IntN(10)*100))
+	//	log.WithContext(ctx).Warnf("rate limit exceeded for %s", req.WgPubKey)
+	//	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 {