[management] Send relay credentials with turn updates (#3164)

send relay credentials when sending turn credentials update to avoid removing servers from clients
netbirdio · Jan 10, 2025 · 649bfb2 · 649bfb2
1 parent 409003b
commit 649bfb2
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 10 deletions.
diff --git 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
@@ -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
@@ -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
+				}
 			}
 		}
 	}