From a7760bf0a7aa18ead65807eca043c2ef1df184ec Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Zolt=C3=A1n=20Papp?= <zoltan.pmail@gmail.com>
Date: Fri, 14 Jun 2024 15:43:18 +0200
Subject: [PATCH] Configurable relay address with env variable

---
 client/internal/connect.go | 23 +++++++++++++++++++----
 relay/client/manager.go    |  6 ++++--
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/client/internal/connect.go b/client/internal/connect.go
index 44e85be53..2c0c2657b 100644
--- a/client/internal/connect.go
+++ b/client/internal/connect.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"net"
+	"os"
 	"runtime"
 	"runtime/debug"
 	"strings"
@@ -245,10 +246,13 @@ func (c *ConnectClient) run(
 
 		c.statusRecorder.MarkSignalConnected()
 
-		relayManager := relayClient.NewManager(engineCtx, loginResp.GetWiretrusteeConfig().GetRelayAddress(), myPrivateKey.PublicKey().String())
-		if err = relayManager.Serve(); err != nil {
-			log.Error(err)
-			return wrapErr(err)
+		relayAddress := relayAddress(loginResp)
+		relayManager := relayClient.NewManager(engineCtx, relayAddress, myPrivateKey.PublicKey().String())
+		if relayAddress != "" {
+			if err = relayManager.Serve(); err != nil {
+				log.Error(err)
+				return wrapErr(err)
+			}
 		}
 
 		peerConfig := loginResp.GetPeerConfig()
@@ -304,6 +308,17 @@ func (c *ConnectClient) run(
 	return nil
 }
 
+func relayAddress(resp *mgmProto.LoginResponse) string {
+	if envRelay := os.Getenv("NB_RELAY_ADDRESS"); envRelay != "" {
+		return envRelay
+	}
+
+	if resp.GetWiretrusteeConfig().GetRelayAddress() != "" {
+		return resp.GetWiretrusteeConfig().GetRelayAddress()
+	}
+	return ""
+}
+
 func (c *ConnectClient) Engine() *Engine {
 	var e *Engine
 	c.engineMutex.Lock()
diff --git a/relay/client/manager.go b/relay/client/manager.go
index 0223c98ab..6b816c4fb 100644
--- a/relay/client/manager.go
+++ b/relay/client/manager.go
@@ -12,6 +12,8 @@ import (
 
 var (
 	relayCleanupInterval = 60 * time.Second
+
+	errRelayClientNotConnected = fmt.Errorf("relay client not connected")
 )
 
 // RelayTrack hold the relay clients for the foregin relay servers.
@@ -80,7 +82,7 @@ func (m *Manager) Serve() error {
 // connection to the relay server.
 func (m *Manager) OpenConn(serverAddress, peerKey string) (net.Conn, error) {
 	if m.relayClient == nil {
-		return nil, fmt.Errorf("relay client not connected")
+		return nil, errRelayClientNotConnected
 	}
 
 	foreign, err := m.isForeignServer(serverAddress)
@@ -101,7 +103,7 @@ func (m *Manager) OpenConn(serverAddress, peerKey string) (net.Conn, error) {
 // This address will be sent to the target peer to choose the common relay server for the communication.
 func (m *Manager) RelayAddress() (net.Addr, error) {
 	if m.relayClient == nil {
-		return nil, fmt.Errorf("relay client not connected")
+		return nil, errRelayClientNotConnected
 	}
 	return m.relayClient.RelayRemoteAddress()
 }