review comments

management/internals/shared/grpc/proxy.go

@@ -183,7 +183,10 @@ func (s *ProxyServiceServer) GetMappingUpdate(req *proto.GetMappingUpdateRequest
 	if token != nil && token.AccountID != nil {
 		accountID = token.AccountID
 
-		existingProxy, _ := s.proxyManager.GetAccountProxy(ctx, *accountID)
+		existingProxy, err := s.proxyManager.GetAccountProxy(ctx, *accountID)
+		if err != nil {
+			log.WithContext(ctx).Debugf("failed to get account proxy for %s: %v", *accountID, err)
+		}
 		if existingProxy != nil && existingProxy.ID != proxyID {
 			if existingProxy.Status == proxy.StatusConnected {
 				return status.Errorf(codes.ResourceExhausted, "limit of 1 self-hosted proxy per account")
@@ -223,7 +226,7 @@ func (s *ProxyServiceServer) GetMappingUpdate(req *proto.GetMappingUpdateRequest
 	if err := s.proxyManager.Connect(ctx, proxyID, proxyAddress, peerInfo, accountID); err != nil {
 		if accountID != nil {
 			cancel()
-			if strings.Contains(err.Error(), "UNIQUE constraint") || strings.Contains(err.Error(), "duplicate key") || strings.Contains(err.Error(), "idx_proxy_account_id_unique") {
+			if errors.Is(err, proxy.ErrAccountProxyAlreadyExists) {
 				return status.Errorf(codes.ResourceExhausted, "limit of 1 self-hosted proxy per account")
 			}
 			return status.Errorf(codes.Internal, "failed to register BYOP proxy: %v", err)
@@ -453,14 +456,18 @@ func (s *ProxyServiceServer) SendAccessLog(ctx context.Context, req *proto.SendA
 // BYOP proxies only receive updates for their own account's services.
 func (s *ProxyServiceServer) SendServiceUpdate(update *proto.GetMappingUpdateResponse) {
 	log.Debugf("Broadcasting service update to all connected proxy servers")
-	var updateAccountID string
-	if len(update.Mapping) > 0 {
-		updateAccountID = update.Mapping[0].AccountId
+	updateAccountIDs := make(map[string]struct{})
+	for _, m := range update.Mapping {
+		if m.AccountId != "" {
+			updateAccountIDs[m.AccountId] = struct{}{}
+		}
 	}
 	s.connectedProxies.Range(func(key, value interface{}) bool {
 		conn := value.(*proxyConnection)
-		if conn.accountID != nil && updateAccountID != "" && *conn.accountID != updateAccountID {
-			return true
+		if conn.accountID != nil && len(updateAccountIDs) > 0 {
+			if _, ok := updateAccountIDs[*conn.accountID]; !ok {
+				return true
+			}
 		}
 		msg := s.perProxyMessage(update, conn.proxyID)
 		if msg == nil {

management/internals/shared/grpc/proxy_group_access_test.go

@@ -91,6 +91,9 @@ func (m *mockReverseProxyManager) StopServiceFromPeer(_ context.Context, _, _, _
 func (m *mockReverseProxyManager) StartExposeReaper(_ context.Context) {}
 
 func (m *mockReverseProxyManager) GetServiceByDomain(_ context.Context, domain string) (*service.Service, error) {
+	if m.err != nil {
+		return nil, m.err
+	}
 	for _, services := range m.proxiesByAccount {
 		for _, svc := range services {
 			if svc.Domain == domain {

management/internals/shared/grpc/proxy_test.go

@@ -11,8 +11,11 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc/codes"
+	grpcstatus "google.golang.org/grpc/status"
 
 	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
+	"github.com/netbirdio/netbird/management/server/types"
 	"github.com/netbirdio/netbird/shared/management/proto"
 )
 
@@ -304,6 +307,58 @@ func TestValidateState_RejectsOldTwoPartFormat(t *testing.T) {
 	assert.Contains(t, err.Error(), "invalid state format")
 }
 
+func scopedCtx(accountID string) context.Context {
+	token := &types.ProxyAccessToken{
+		ID:        "token-1",
+		AccountID: &accountID,
+	}
+	return context.WithValue(context.Background(), ProxyTokenContextKey, token)
+}
+
+func globalCtx() context.Context {
+	token := &types.ProxyAccessToken{
+		ID: "token-global",
+	}
+	return context.WithValue(context.Background(), ProxyTokenContextKey, token)
+}
+
+func TestEnforceAccountScope_AllowsMatchingAccount(t *testing.T) {
+	err := enforceAccountScope(scopedCtx("acc-1"), "acc-1")
+	assert.NoError(t, err)
+}
+
+func TestEnforceAccountScope_BlocksMismatchedAccount(t *testing.T) {
+	err := enforceAccountScope(scopedCtx("acc-1"), "acc-2")
+	require.Error(t, err)
+	st, ok := grpcstatus.FromError(err)
+	require.True(t, ok)
+	assert.Equal(t, codes.PermissionDenied, st.Code())
+}
+
+func TestEnforceAccountScope_BlocksEmptyRequestAccountID(t *testing.T) {
+	err := enforceAccountScope(scopedCtx("acc-1"), "")
+	require.Error(t, err)
+	st, ok := grpcstatus.FromError(err)
+	require.True(t, ok)
+	assert.Equal(t, codes.PermissionDenied, st.Code())
+}
+
+func TestEnforceAccountScope_AllowsGlobalToken(t *testing.T) {
+	err := enforceAccountScope(globalCtx(), "acc-1")
+	assert.NoError(t, err)
+
+	err = enforceAccountScope(globalCtx(), "acc-2")
+	assert.NoError(t, err)
+
+	err = enforceAccountScope(globalCtx(), "")
+	assert.NoError(t, err)
+}
+
+func TestEnforceAccountScope_AllowsNoTokenInContext(t *testing.T) {
+	err := enforceAccountScope(context.Background(), "acc-1")
+	assert.NoError(t, err)
+}
+
 func TestValidateState_RejectsInvalidHMAC(t *testing.T) {
 	ctx := context.Background()
 	pkceStore, err := NewPKCEVerifierStore(ctx, 10*time.Minute, 10*time.Minute, 100)