From bc33abb9d4b0d7201adac5d71ace8cdb19f48b5b Mon Sep 17 00:00:00 2001
From: Evan Baker <rbtr@users.noreply.github.com>
Date: Mon, 13 Jan 2025 22:14:48 +0000
Subject: [PATCH] fix: prevent CNS hanging on HNS restart

Signed-off-by: Evan Baker <rbtr@users.noreply.github.com>
---
 cns/service/main.go    |  6 +++++-
 platform/os_windows.go | 31 +++++++++++++++++++------------
 2 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/cns/service/main.go b/cns/service/main.go
index f242c3da84..f8235bba9e 100644
--- a/cns/service/main.go
+++ b/cns/service/main.go
@@ -799,11 +799,15 @@ func main() {
 	}
 
 	// Setting the remote ARP MAC address to 12-34-56-78-9a-bc on windows for external traffic if HNS is enabled
-	err = platform.SetSdnRemoteArpMacAddress(rootCtx)
+	arpCtx, arpCtxCancel := context.WithTimeout(rootCtx, 30*time.Second)
+	err = platform.SetSdnRemoteArpMacAddress(arpCtx)
 	if err != nil {
 		logger.Errorf("Failed to set remote ARP MAC address: %v", err)
+		arpCtxCancel()
 		return
 	}
+	arpCtxCancel()
+
 	// We are only setting the PriorityVLANTag in 'cns.Direct' mode, because it neatly maps today, to 'isUsingMultitenancy'
 	// In the future, we would want to have a better CNS flag, to explicitly say, this CNS is using multitenancy
 	if cnsconfig.ChannelMode == cns.Direct {
diff --git a/platform/os_windows.go b/platform/os_windows.go
index c76b8a7042..1ce81b24d8 100644
--- a/platform/os_windows.go
+++ b/platform/os_windows.go
@@ -247,6 +247,18 @@ func (p *execClient) ExecutePowershellCommandWithContext(ctx context.Context, co
 
 // SetSdnRemoteArpMacAddress sets the regkey for SDNRemoteArpMacAddress needed for multitenancy if hns is enabled
 func SetSdnRemoteArpMacAddress(ctx context.Context) error {
+	if err := setSDNRemoteARPRegKey(); err != nil {
+		return err
+	}
+	log.Printf("SDNRemoteArpMacAddress regKey set successfully")
+	if err := restartHNS(ctx); err != nil {
+		return err
+	}
+	log.Printf("HNS service restarted successfully")
+	return nil
+}
+
+func setSDNRemoteARPRegKey() error {
 	log.Printf("Setting SDNRemoteArpMacAddress regKey")
 	// open the registry key
 	k, err := registry.OpenKey(registry.LOCAL_MACHINE, `SYSTEM\CurrentControlSet\Services\hns\State`, registry.READ|registry.SET_VALUE)
@@ -265,7 +277,10 @@ func SetSdnRemoteArpMacAddress(ctx context.Context) error {
 	if err = k.SetStringValue("SDNRemoteArpMacAddress", SDNRemoteArpMacAddress); err != nil {
 		return errors.Wrap(err, "could not set registry key")
 	}
-	log.Printf("SDNRemoteArpMacAddress regKey set successfully")
+	return nil
+}
+
+func restartHNS(ctx context.Context) error {
 	log.Printf("Restarting HNS service")
 	// connect to the service manager
 	m, err := mgr.Connect()
@@ -279,16 +294,8 @@ func SetSdnRemoteArpMacAddress(ctx context.Context) error {
 		return errors.Wrap(err, "could not access service")
 	}
 	defer service.Close()
-	if err := restartService(ctx, service); err != nil {
-		return errors.Wrap(err, "could not restart service")
-	}
-	log.Printf("HNS service restarted successfully")
-	return nil
-}
-
-func restartService(ctx context.Context, s *mgr.Service) error {
 	// Stop the service
-	_, err := s.Control(svc.Stop)
+	_, err = service.Control(svc.Stop)
 	if err != nil {
 		return errors.Wrap(err, "could not stop service")
 	}
@@ -296,7 +303,7 @@ func restartService(ctx context.Context, s *mgr.Service) error {
 	ticker := time.NewTicker(500 * time.Millisecond) //nolint:gomnd // 500ms
 	defer ticker.Stop()
 	for { // hacky cancellable do-while
-		status, err := s.Query()
+		status, err := service.Query()
 		if err != nil {
 			return errors.Wrap(err, "could not query service status")
 		}
@@ -310,7 +317,7 @@ func restartService(ctx context.Context, s *mgr.Service) error {
 		}
 	}
 	// Start the service again
-	if err := s.Start(); err != nil {
+	if err := service.Start(); err != nil {
 		return errors.Wrap(err, "could not start service")
 	}
 	return nil