From a6e01c4963f4054945053fa27105e3dc6cab61c1 Mon Sep 17 00:00:00 2001
From: Abhishek Kondur <abhi281342@gmail.com>
Date: Mon, 14 Nov 2022 23:25:19 +0530
Subject: [PATCH] extclients track watch with context and exiting durig cleanup

---
 nm-proxy/common/common.go   |  2 ++
 nm-proxy/manager/manager.go | 42 +++++++++++++++++++++++--------------
 nm-proxy/wg/wg.go           |  2 +-
 3 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/nm-proxy/common/common.go b/nm-proxy/common/common.go
index a7ac6031..2cd2ced1 100644
--- a/nm-proxy/common/common.go
+++ b/nm-proxy/common/common.go
@@ -78,6 +78,8 @@ var WgIfaceKeyMap = make(map[string]struct{})
 
 var RelayPeerMap = make(map[string]map[string]RemotePeer)
 
+var ExtClientsWaitTh = make(map[string][]context.CancelFunc)
+
 // RunCmd - runs a local command
 func RunCmd(command string, printerr bool) (string, error) {
 	args := strings.Fields(command)
diff --git a/nm-proxy/manager/manager.go b/nm-proxy/manager/manager.go
index 3710959b..3b92fabb 100644
--- a/nm-proxy/manager/manager.go
+++ b/nm-proxy/manager/manager.go
@@ -1,6 +1,7 @@
 package manager
 
 import (
+	"context"
 	"crypto/md5"
 	"errors"
 	"fmt"
@@ -24,7 +25,6 @@ TODO:-
 			-> start remote conn after endpoint is updated
 		-->
 */
-var sent bool
 
 type ProxyAction string
 
@@ -73,13 +73,9 @@ func StartProxyManager(manageChan chan *ManagerAction) {
 
 		select {
 		case mI := <-manageChan:
-			if sent {
-				continue
-			}
 			log.Printf("-------> PROXY-MANAGER: %+v\n", mI)
 			switch mI.Action {
 			case AddInterface:
-				sent = true
 				common.IsRelay = mI.Payload.IsRelay
 				if mI.Payload.IsRelay {
 					mI.RelayPeers()
@@ -194,7 +190,13 @@ func cleanUp(iface string) {
 		}
 	}
 	delete(common.WgIFaceMap, iface)
-	time.Sleep(time.Second * 5)
+	if waitThs, ok := common.ExtClientsWaitTh[iface]; ok {
+		for _, cancelF := range waitThs {
+			cancelF()
+		}
+		delete(common.ExtClientsWaitTh, iface)
+	}
+
 	log.Println("CLEANED UP..........")
 }
 
@@ -277,6 +279,8 @@ func (m *ManagerAction) AddInterfaceToProxy() error {
 			go func(wgInterface *wg.WGIface, peer *wgtypes.PeerConfig,
 				isRelayed, isExtClient, isAttachedExtClient bool, relayTo *net.UDPAddr, peerConf PeerConf) {
 				addExtClient := false
+				ctx, cancel := context.WithCancel(context.Background())
+				common.ExtClientsWaitTh[wgInterface.Name] = append(common.ExtClientsWaitTh[wgInterface.Name], cancel)
 				defer func() {
 					if addExtClient {
 						log.Println("GOT ENDPOINT for Extclient adding peer...")
@@ -293,19 +297,25 @@ func (m *ManagerAction) AddInterfaceToProxy() error {
 					}
 				}()
 				for {
-					wgInterface, err := wg.NewWGIFace(ifaceName, "127.0.0.1/32", wg.DefaultMTU)
-					if err != nil {
-						log.Println("Failed init new interface: ", err)
-						continue
-					}
-					for _, devpeerI := range wgInterface.Device.Peers {
-						if devpeerI.PublicKey.String() == peer.PublicKey.String() && devpeerI.Endpoint != nil {
-							peer.Endpoint = devpeerI.Endpoint
-							addExtClient = true
+					select {
+					case <-ctx.Done():
+						log.Println("Exiting extclient watch Thread for: ", wgInterface.Device.PublicKey.String())
+						return
+					default:
+						wgInterface, err := wg.NewWGIFace(ifaceName, "127.0.0.1/32", wg.DefaultMTU)
+						if err != nil {
+							log.Println("Failed init new interface: ", err)
 							return
 						}
+						for _, devpeerI := range wgInterface.Device.Peers {
+							if devpeerI.PublicKey.String() == peer.PublicKey.String() && devpeerI.Endpoint != nil {
+								peer.Endpoint = devpeerI.Endpoint
+								addExtClient = true
+								return
+							}
+						}
+						time.Sleep(time.Second * 5)
 					}
-					time.Sleep(time.Second * 5)
 
 				}
 
diff --git a/nm-proxy/wg/wg.go b/nm-proxy/wg/wg.go
index b7fff0db..7ee5e148 100644
--- a/nm-proxy/wg/wg.go
+++ b/nm-proxy/wg/wg.go
@@ -74,7 +74,7 @@ func (w *WGIface) GetWgIface(iface string) error {
 		return err
 	}
 
-	log.Printf("----> DEVICE: %+v\n", dev)
+	//log.Printf("----> DEVICE: %+v\n", dev)
 	w.Device = dev
 	w.Port = dev.ListenPort
 	return nil