From 5d4899cb783c035a5869b81dbdf9019f46b288fd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Zolt=C3=A1n=20Papp?= <zoltan.pmail@gmail.com>
Date: Thu, 28 Nov 2024 17:34:09 +0100
Subject: [PATCH] Fix logging and error handling

---
 relay/client/client.go             |  2 +-
 relay/client/dialer/quic/quic.go   | 12 ++++++--
 relay/client/dialer/race_dialer.go | 44 ++++++++++++++++++------------
 relay/client/dialer/ws/ws.go       | 12 ++++++--
 4 files changed, 46 insertions(+), 24 deletions(-)

diff --git a/relay/client/client.go b/relay/client/client.go
index 256d725e8..45ecef6ca 100644
--- a/relay/client/client.go
+++ b/relay/client/client.go
@@ -261,7 +261,7 @@ func (c *Client) Close() error {
 }
 
 func (c *Client) connect() error {
-	conn, err := dialer.RaceDial(c.connectionURL, quic.Dial, ws.Dial)
+	conn, err := dialer.RaceDial(c.log, c.connectionURL, quic.Dialer{}, ws.Dialer{})
 	if err != nil {
 		return err
 	}
diff --git a/relay/client/dialer/quic/quic.go b/relay/client/dialer/quic/quic.go
index b81a78a08..66000171c 100644
--- a/relay/client/dialer/quic/quic.go
+++ b/relay/client/dialer/quic/quic.go
@@ -15,7 +15,14 @@ import (
 	nbnet "github.com/netbirdio/netbird/util/net"
 )
 
-func Dial(ctx context.Context, address string) (net.Conn, error) {
+type Dialer struct {
+}
+
+func (d Dialer) Protocol() string {
+	return "QUIC"
+}
+
+func (d Dialer) Dial(ctx context.Context, address string) (net.Conn, error) {
 	quicURL, err := prepareURL(address)
 	if err != nil {
 		return nil, err
@@ -41,8 +48,7 @@ func Dial(ctx context.Context, address string) (net.Conn, error) {
 
 	session, err := quic.Dial(ctx, udpConn, udpAddr, quictls.ClientQUICTLSConfig(), quicConfig)
 	if err != nil {
-		if errors.Is(ctx.Err(), context.Canceled) {
-			log.Infof("QUIC connection aborted to: %s", quicURL)
+		if errors.Is(err, context.Canceled) {
 			return nil, err
 		}
 		log.Errorf("failed to dial to Relay server via QUIC '%s': %s", quicURL, err)
diff --git a/relay/client/dialer/race_dialer.go b/relay/client/dialer/race_dialer.go
index f58a705ec..4536c7374 100644
--- a/relay/client/dialer/race_dialer.go
+++ b/relay/client/dialer/race_dialer.go
@@ -9,46 +9,56 @@ import (
 	log "github.com/sirupsen/logrus"
 )
 
-type DialFn func(ctx context.Context, address string) (net.Conn, error)
-
-type dialResult struct {
-	Conn net.Conn
-	Err  error
+type DialerFn interface {
+	Dial(ctx context.Context, address string) (net.Conn, error)
+	Protocol() string
 }
 
-func RaceDial(serverURL string, DialFns ...DialFn) (net.Conn, error) {
-	connChan := make(chan dialResult, len(DialFns))
+type dialResult struct {
+	Conn     net.Conn
+	Protocol string
+	Err      error
+}
+
+func RaceDial(log *log.Entry, serverURL string, dialerFns ...DialerFn) (net.Conn, error) {
+	connChan := make(chan dialResult, len(dialerFns))
 	winnerConn := make(chan net.Conn, 1)
 	abortCtx, abort := context.WithCancel(context.Background())
 	defer abort()
 
-	for _, dfn := range DialFns {
+	for _, d := range dialerFns {
 		go func() {
 			ctx, cancel := context.WithTimeout(abortCtx, 30*time.Second)
 			defer cancel()
 
-			conn, err := dfn(ctx, serverURL)
-			if err != nil {
-				log.Errorf("failed to dial: %s", err)
-			}
-
-			connChan <- dialResult{Conn: conn, Err: err}
+			log.Infof("dialing Relay server via %s", d.Protocol())
+			conn, err := d.Dial(ctx, serverURL)
+			connChan <- dialResult{Conn: conn, Protocol: d.Protocol(), Err: err}
 		}()
 	}
 
 	go func() {
 		var hasWinner bool
-		for i := 0; i < len(DialFns); i++ {
+		for i := 0; i < len(dialerFns); i++ {
 			dr := <-connChan
 			if dr.Err != nil {
+				if errors.Is(dr.Err, context.Canceled) {
+					log.Infof("connection attempt aborted via: %s", dr.Protocol)
+				} else {
+					log.Errorf("failed to dial via: %s, %s", dr.Protocol, dr.Err)
+				}
 				continue
 			}
 
 			if hasWinner {
-				_ = dr.Conn.Close()
+				if cerr := dr.Conn.Close(); cerr != nil {
+					log.Warnf("failed to close connection via: %s, %s", dr.Protocol, cerr)
+				}
 				continue
 			}
 
+			log.Infof("successfully dialed via: %s", dr.Protocol)
+
 			hasWinner = true
 			winnerConn <- dr.Conn
 		}
@@ -57,7 +67,7 @@ func RaceDial(serverURL string, DialFns ...DialFn) (net.Conn, error) {
 
 	conn, ok := <-winnerConn
 	if !ok {
-		return nil, errors.New("failed to dial to Relay server")
+		return nil, errors.New("failed to dial to Relay server on any protocol")
 	}
 	return conn, nil
 }
diff --git a/relay/client/dialer/ws/ws.go b/relay/client/dialer/ws/ws.go
index 338e2aa6d..df91a66d4 100644
--- a/relay/client/dialer/ws/ws.go
+++ b/relay/client/dialer/ws/ws.go
@@ -16,7 +16,14 @@ import (
 	nbnet "github.com/netbirdio/netbird/util/net"
 )
 
-func Dial(ctx context.Context, address string) (net.Conn, error) {
+type Dialer struct {
+}
+
+func (d Dialer) Protocol() string {
+	return "WS"
+}
+
+func (d Dialer) Dial(ctx context.Context, address string) (net.Conn, error) {
 	wsURL, err := prepareURL(address)
 	if err != nil {
 		return nil, err
@@ -34,8 +41,7 @@ func Dial(ctx context.Context, address string) (net.Conn, error) {
 
 	wsConn, resp, err := websocket.Dial(ctx, parsedURL.String(), opts)
 	if err != nil {
-		if errors.Is(ctx.Err(), context.Canceled) {
-			log.Infof("WebSocket connection aborted to: %s", wsURL)
+		if errors.Is(err, context.Canceled) {
 			return nil, err
 		}
 		log.Errorf("failed to dial to Relay server '%s': %s", wsURL, err)