From 2f0b652dadf786b77bb6698b5745ee2b75baab49 Mon Sep 17 00:00:00 2001
From: Zoltan Papp <zoltan.pmail@gmail.com>
Date: Tue, 21 Feb 2023 10:46:58 +0100
Subject: [PATCH] Fix error handling in Stop/Start functions (#699)

Properly close all resources in case of any error
during the start or stop procedure
---
 client/internal/engine.go | 102 ++++++++++++++++++++------------------
 1 file changed, 55 insertions(+), 47 deletions(-)

diff --git a/client/internal/engine.go b/client/internal/engine.go
index d629e6b18..29bf7eccb 100644
--- a/client/internal/engine.go
+++ b/client/internal/engine.go
@@ -157,54 +157,8 @@ func (e *Engine) Stop() error {
 	// Removing peers happens in the conn.CLose() asynchronously
 	time.Sleep(500 * time.Millisecond)
 
-	log.Debugf("removing Netbird interface %s", e.config.WgIfaceName)
-	err = e.wgInterface.Close()
-	if err != nil {
-		log.Errorf("failed closing Netbird interface %s %v", e.config.WgIfaceName, err)
-		return err
-	}
-
-	if e.udpMux != nil {
-		if err := e.udpMux.Close(); err != nil {
-			log.Debugf("close udp mux: %v", err)
-		}
-	}
-
-	if e.udpMuxSrflx != nil {
-		if err := e.udpMuxSrflx.Close(); err != nil {
-			log.Debugf("close server reflexive udp mux: %v", err)
-		}
-	}
-
-	if e.udpMuxConn != nil {
-		if err := e.udpMuxConn.Close(); err != nil {
-			log.Debugf("close udp mux connection: %v", err)
-		}
-	}
-
-	if e.udpMuxConnSrflx != nil {
-		if err := e.udpMuxConnSrflx.Close(); err != nil {
-			log.Debugf("close server reflexive udp mux connection: %v", err)
-		}
-	}
-
-	if !isNil(e.sshServer) {
-		err := e.sshServer.Stop()
-		if err != nil {
-			log.Warnf("failed stopping the SSH server: %v", err)
-		}
-	}
-
-	if e.routeManager != nil {
-		e.routeManager.Stop()
-	}
-
-	if e.dnsServer != nil {
-		e.dnsServer.Stop()
-	}
-
+	e.close()
 	log.Infof("stopped Netbird Engine")
-
 	return nil
 }
 
@@ -234,12 +188,14 @@ func (e *Engine) Start() error {
 	e.udpMuxConn, err = net.ListenUDP(networkName, &net.UDPAddr{Port: e.config.UDPMuxPort})
 	if err != nil {
 		log.Errorf("failed listening on UDP port %d: [%s]", e.config.UDPMuxPort, err.Error())
+		e.close()
 		return err
 	}
 
 	e.udpMuxConnSrflx, err = net.ListenUDP(networkName, &net.UDPAddr{Port: e.config.UDPMuxSrflxPort})
 	if err != nil {
 		log.Errorf("failed listening on UDP port %d: [%s]", e.config.UDPMuxSrflxPort, err.Error())
+		e.close()
 		return err
 	}
 
@@ -249,12 +205,14 @@ func (e *Engine) Start() error {
 	err = e.wgInterface.Create()
 	if err != nil {
 		log.Errorf("failed creating tunnel interface %s: [%s]", wgIfaceName, err.Error())
+		e.close()
 		return err
 	}
 
 	err = e.wgInterface.Configure(myPrivateKey.String(), e.config.WgPort)
 	if err != nil {
 		log.Errorf("failed configuring Wireguard interface [%s]: %s", wgIfaceName, err.Error())
+		e.close()
 		return err
 	}
 
@@ -264,6 +222,7 @@ func (e *Engine) Start() error {
 		// todo fix custom address
 		dnsServer, err := dns.NewDefaultServer(e.ctx, e.wgInterface, e.config.CustomDNSAddress)
 		if err != nil {
+			e.close()
 			return err
 		}
 		e.dnsServer = dnsServer
@@ -1009,6 +968,55 @@ func (e *Engine) parseNATExternalIPMappings() []string {
 	return mappedIPs
 }
 
+func (e *Engine) close() {
+	log.Debugf("removing Netbird interface %s", e.config.WgIfaceName)
+	if e.wgInterface != nil {
+		if err := e.wgInterface.Close(); err != nil {
+			log.Errorf("failed closing Netbird interface %s %v", e.config.WgIfaceName, err)
+		}
+	}
+
+	if e.udpMux != nil {
+		if err := e.udpMux.Close(); err != nil {
+			log.Debugf("close udp mux: %v", err)
+		}
+	}
+
+	if e.udpMuxSrflx != nil {
+		if err := e.udpMuxSrflx.Close(); err != nil {
+			log.Debugf("close server reflexive udp mux: %v", err)
+		}
+	}
+
+	if e.udpMuxConn != nil {
+		if err := e.udpMuxConn.Close(); err != nil {
+			log.Debugf("close udp mux connection: %v", err)
+		}
+	}
+
+	if e.udpMuxConnSrflx != nil {
+		if err := e.udpMuxConnSrflx.Close(); err != nil {
+			log.Debugf("close server reflexive udp mux connection: %v", err)
+		}
+	}
+
+	if !isNil(e.sshServer) {
+		err := e.sshServer.Stop()
+		if err != nil {
+			log.Warnf("failed stopping the SSH server: %v", err)
+		}
+	}
+
+	if e.routeManager != nil {
+		e.routeManager.Stop()
+	}
+
+	if e.dnsServer != nil {
+		e.dnsServer.Stop()
+	}
+
+}
+
 func findIPFromInterfaceName(ifaceName string) (net.IP, error) {
 	iface, err := net.InterfaceByName(ifaceName)
 	if err != nil {