go mod tidy

go.mod

@@ -106,6 +106,7 @@ require (
 	golang.org/x/oauth2 v0.24.0
 	golang.org/x/sync v0.13.0
 	golang.org/x/term v0.31.0
+	golang.org/x/time v0.5.0
 	google.golang.org/api v0.177.0
 	gopkg.in/yaml.v3 v3.0.1
 	gorm.io/driver/mysql v1.5.7
@@ -240,7 +241,6 @@ require (
 	golang.org/x/image v0.18.0 // indirect
 	golang.org/x/mod v0.17.0 // indirect
 	golang.org/x/text v0.24.0 // indirect
-	golang.org/x/time v0.5.0 // indirect
 	golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d // indirect
 	golang.zx2c4.com/wintun v0.0.0-20230126152724-0fa3db229ce2 // indirect
 	google.golang.org/genproto/googleapis/api v0.0.0-20240509183442-62759503f434 // indirect

management/server/account.go

@@ -1522,8 +1522,8 @@ func (am *DefaultAccountManager) SyncPeerMeta(ctx context.Context, peerPubKey st
 		return err
 	}
 
-	unlock := am.Store.AcquireReadLockByUID(ctx, accountID)
+	// unlock := am.Store.AcquireReadLockByUID(ctx, accountID)
-	defer unlock()
+	// defer unlock()
 
 	unlockPeer := am.Store.AcquireWriteLockByUID(ctx, peerPubKey)
 	defer unlockPeer()

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"
@@ -411,11 +413,17 @@ func (s *GRPCServer) parseRequest(ctx context.Context, req *proto.EncryptedMessa
 	return peerKey, nil
 }
 
+var loginLimiter = rate.NewLimiter(rate.Every(time.Minute/100), 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 {
@@ -434,6 +442,16 @@ func (s *GRPCServer) Login(ctx context.Context, req *proto.EncryptedMessage) (*p
 		return nil, err
 	}
 
+	err = TryAcquire(loginReq.GetPeerKeys().String())
+	if err != nil {
+		log.WithContext(ctx).Debugf("error while acquiring grpc semaphore for %s: %v", loginReq.GetPeerKeys().String(), err)
+		return nil, err
+	}
+
+	defer func() {
+		Release(loginReq.GetPeerKeys().String())
+	}()
+
 	//nolint
 	ctx = context.WithValue(ctx, nbContext.PeerIDKey, peerKey.String())
 	accountID, err := s.accountManager.GetAccountIDForPeerKey(ctx, peerKey.String())

management/server/peer.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/sha256"
 	b64 "encoding/base64"
+	"errors"
 	"fmt"
 	"net"
 	"slices"
@@ -445,6 +446,30 @@ func (am *DefaultAccountManager) GetPeerNetwork(ctx context.Context, peerID stri
 	return account.Network.Copy(), err
 }
 
+var (
+	mu      sync.Mutex
+	busyMap = make(map[string]bool)
+)
+
+func TryAcquire(key string) error {
+	mu.Lock()
+	defer mu.Unlock()
+
+	if busyMap[key] {
+		return errors.New("keep your calm")
+	}
+
+	busyMap[key] = true
+	return nil
+}
+
+func Release(key string) {
+	mu.Lock()
+	defer mu.Unlock()
+
+	delete(busyMap, key)
+}
+
 // AddPeer adds a new peer to the Store.
 // Each Account has a list of pre-authorized SetupKey and if no Account has a given key err with a code status.PermissionDenied
 // will be returned, meaning the setup key is invalid or not found.
@@ -474,11 +499,13 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s
 		return nil, nil, nil, status.Errorf(status.NotFound, "failed adding new peer: account not found")
 	}
 
-	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
+	if err := TryAcquire(accountID); err != nil {
-	defer func() {
+		log.Debugf("failed to acquire semaphore: %v", err)
-		if unlock != nil {
+		return nil, nil, nil, status.Errorf(status.PreconditionFailed, "failed to acquire semaphore: %v", err)
-			unlock()
+	}
-		}
+
+	go func() {
+		defer Release(accountID)
 	}()
 
 	// This is a handling for the case when the same machine (with the same WireGuard pub key) tries to register twice.
@@ -657,9 +684,6 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s
 
 	am.StoreEvent(ctx, opEvent.InitiatorID, opEvent.TargetID, opEvent.AccountID, opEvent.Activity, opEvent.Meta)
 
-	unlock()
-	unlock = nil
-
 	if updateAccountPeers {
 		am.BufferUpdateAccountPeers(ctx, accountID)
 	}