From 649bfb236bb820a2c67940029a0f677d155b6cfb Mon Sep 17 00:00:00 2001
From: Maycon Santos <mlsmaycon@gmail.com>
Date: Fri, 10 Jan 2025 09:44:02 +0100
Subject: [PATCH] [management] Send relay credentials with turn updates (#3164)

send relay credentials when sending turn credentials update to avoid removing servers
from clients
---
 management/server/peer_test.go      |  5 +++--
 management/server/token_mgr.go      | 17 ++++++++++++++---
 management/server/token_mgr_test.go | 13 ++++++++-----
 3 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/management/server/peer_test.go b/management/server/peer_test.go
index 5f500c2267c..0c751e6c409 100644
--- a/management/server/peer_test.go
+++ b/management/server/peer_test.go
@@ -13,13 +13,14 @@ import (
 	"testing"
 	"time"
 
-	"github.com/netbirdio/netbird/management/server/util"
 	"github.com/rs/xid"
 	log "github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 
+	"github.com/netbirdio/netbird/management/server/util"
+
 	resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
 	routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
 	networkTypes "github.com/netbirdio/netbird/management/server/networks/types"
@@ -937,7 +938,7 @@ func BenchmarkUpdateAccountPeers(b *testing.B) {
 		{"Small single", 50, 10, 90, 120, 90, 120},
 		{"Medium single", 500, 10, 110, 170, 120, 200},
 		{"Large 5", 5000, 15, 1300, 2100, 4900, 7000},
-		{"Extra Large", 2000, 2000, 1300, 2400, 4000, 6400},
+		{"Extra Large", 2000, 2000, 1300, 2400, 3900, 6400},
 	}
 
 	log.SetOutput(io.Discard)
diff --git a/management/server/token_mgr.go b/management/server/token_mgr.go
index ef8276b5924..fd67fa3e328 100644
--- a/management/server/token_mgr.go
+++ b/management/server/token_mgr.go
@@ -158,7 +158,7 @@ func (m *TimeBasedAuthSecretsManager) refreshTURNTokens(ctx context.Context, pee
 			log.WithContext(ctx).Debugf("stopping TURN refresh for %s", peerID)
 			return
 		case <-ticker.C:
-			m.pushNewTURNTokens(ctx, peerID)
+			m.pushNewTURNAndRelayTokens(ctx, peerID)
 		}
 	}
 }
@@ -178,7 +178,7 @@ func (m *TimeBasedAuthSecretsManager) refreshRelayTokens(ctx context.Context, pe
 	}
 }
 
-func (m *TimeBasedAuthSecretsManager) pushNewTURNTokens(ctx context.Context, peerID string) {
+func (m *TimeBasedAuthSecretsManager) pushNewTURNAndRelayTokens(ctx context.Context, peerID string) {
 	turnToken, err := m.turnHmacToken.GenerateToken(sha1.New)
 	if err != nil {
 		log.Errorf("failed to generate token for peer '%s': %s", peerID, err)
@@ -201,10 +201,21 @@ func (m *TimeBasedAuthSecretsManager) pushNewTURNTokens(ctx context.Context, pee
 	update := &proto.SyncResponse{
 		WiretrusteeConfig: &proto.WiretrusteeConfig{
 			Turns: turns,
-			// omit Relay to avoid updates there
 		},
 	}
 
+	// workaround for the case when client is unable to handle turn and relay updates at different time
+	if m.relayCfg != nil {
+		token, err := m.GenerateRelayToken()
+		if err == nil {
+			update.WiretrusteeConfig.Relay = &proto.RelayConfig{
+				Urls:           m.relayCfg.Addresses,
+				TokenPayload:   token.Payload,
+				TokenSignature: token.Signature,
+			}
+		}
+	}
+
 	log.WithContext(ctx).Debugf("sending new TURN credentials to peer %s", peerID)
 	m.updateManager.SendUpdate(ctx, peerID, &UpdateMessage{Update: update})
 }
diff --git a/management/server/token_mgr_test.go b/management/server/token_mgr_test.go
index 3e63346c2d0..2aafb9f6856 100644
--- a/management/server/token_mgr_test.go
+++ b/management/server/token_mgr_test.go
@@ -133,11 +133,14 @@ loop:
 			}
 		}
 		if relay := update.Update.GetWiretrusteeConfig().GetRelay(); relay != nil {
-			relayUpdates++
-			if relayUpdates == 1 {
-				firstRelayUpdate = relay
-			} else {
-				secondRelayUpdate = relay
+			// avoid updating on turn updates since they also send relay credentials
+			if update.Update.GetWiretrusteeConfig().GetTurns() == nil {
+				relayUpdates++
+				if relayUpdates == 1 {
+					firstRelayUpdate = relay
+				} else {
+					secondRelayUpdate = relay
+				}
 			}
 		}
 	}