remove powershell fromwindows registry interactions

Signed-off-by: Evan Baker <[email protected]>

Author: rbtr

cns/service/main.go

@@ -789,8 +789,7 @@ func main() {
 	}
 
 	// Setting the remote ARP MAC address to 12-34-56-78-9a-bc on windows for external traffic if HNS is enabled
-	execClient := platform.NewExecClient(nil)
-	err = platform.SetSdnRemoteArpMacAddress(execClient)
+	err = platform.SetSdnRemoteArpMacAddress()
 	if err != nil {
 		logger.Errorf("Failed to set remote ARP MAC address: %v", err)
 		return

platform/os_linux.go

@@ -179,7 +179,7 @@ func (p *execClient) KillProcessByName(processName string) error {
 
 // SetSdnRemoteArpMacAddress sets the regkey for SDNRemoteArpMacAddress needed for multitenancy
 // This operation is specific to windows OS
-func SetSdnRemoteArpMacAddress(_ ExecClient) error {
+func SetSdnRemoteArpMacAddress() error {
 	return nil
 }
 

platform/os_windows.go

@@ -20,6 +20,9 @@ import (
 	"github.com/pkg/errors"
 	"go.uber.org/zap"
 	"golang.org/x/sys/windows"
+	"golang.org/x/sys/windows/registry"
+	"golang.org/x/sys/windows/svc"
+	"golang.org/x/sys/windows/svc/mgr"
 )
 
 const (
@@ -61,24 +64,12 @@ const (
 	// for vlan tagged arp requests
 	SDNRemoteArpMacAddress = "12-34-56-78-9a-bc"
 
-	// Command to get SDNRemoteArpMacAddress registry key
-	GetSdnRemoteArpMacAddressCommand = "(Get-ItemProperty " +
-		"-Path HKLM:\\SYSTEM\\CurrentControlSet\\Services\\hns\\State -Name SDNRemoteArpMacAddress).SDNRemoteArpMacAddress"
-
-	// Command to set SDNRemoteArpMacAddress registry key
-	SetSdnRemoteArpMacAddressCommand = "Set-ItemProperty " +
-		"-Path HKLM:\\SYSTEM\\CurrentControlSet\\Services\\hns\\State -Name SDNRemoteArpMacAddress -Value \"12-34-56-78-9a-bc\""
-
-	// Command to check if system has hns state path or not
-	CheckIfHNSStatePathExistsCommand = "Test-Path " +
-		"-Path HKLM:\\SYSTEM\\CurrentControlSet\\Services\\hns\\State"
-
 	// Command to fetch netadapter and pnp id
+	// TODO can we replace this (and things in endpoint_windows) with "golang.org/x/sys/windows"
+	// var adapterInfo windows.IpAdapterInfo
+	// var bufferSize uint32 = uint32(unsafe.Sizeof(adapterInfo))
 	GetMacAddressVFPPnpIDMapping = "Get-NetAdapter | Select-Object MacAddress, PnpDeviceID| Format-Table -HideTableHeaders"
 
-	// Command to restart HNS service
-	RestartHnsServiceCommand = "Restart-Service -Name hns"
-
 	// Interval between successive checks for mellanox adapter's PriorityVLANTag value
 	defaultMellanoxMonitorInterval = 30 * time.Second
 
@@ -257,43 +248,70 @@ func (p *execClient) ExecutePowershellCommandWithContext(ctx context.Context, co
 }
 
 // SetSdnRemoteArpMacAddress sets the regkey for SDNRemoteArpMacAddress needed for multitenancy if hns is enabled
-func SetSdnRemoteArpMacAddress(execClient ExecClient) error {
-	exists, err := execClient.ExecutePowershellCommand(CheckIfHNSStatePathExistsCommand)
+func SetSdnRemoteArpMacAddress() error {
+	// open the registry key
+	k, err := registry.OpenKey(registry.LOCAL_MACHINE, "SYSTEM\\CurrentControlSet\\Services\\hns\\State", registry.READ|registry.SET_VALUE)
+	if err != nil && !errors.Is(err, registry.ErrNotExist) {
+		return errors.Wrap(err, "could not open registry key")
+	}
+	defer k.Close()
+	// set the key value
+	if err = setSDNRemoteARPMACAddress(k); err != nil {
+		return errors.Wrap(err, "could not set registry key")
+	}
+	// connect to the service manager
+	m, err := mgr.Connect()
 	if err != nil {
-		errMsg := fmt.Sprintf("Failed to check the existent of hns state path due to error %s", err.Error())
-		log.Printf(errMsg)
-		return errors.Errorf(errMsg)
+		return errors.Wrap(err, "could not connect to service manager")
 	}
-	if strings.EqualFold(exists, "false") {
-		log.Printf("hns state path does not exist, skip setting SdnRemoteArpMacAddress")
-		return nil
+	defer m.Disconnect() //nolint:errcheck // ignore error
+	// open the HNS service
+	service, err := m.OpenService("hns")
+	if err != nil {
+		return errors.Wrap(err, "could not access service")
 	}
-	if sdnRemoteArpMacAddressSet == false {
-		result, err := execClient.ExecutePowershellCommand(GetSdnRemoteArpMacAddressCommand)
-		if err != nil {
-			return err
-		}
-
-		// Set the reg key if not already set or has incorrect value
-		if result != SDNRemoteArpMacAddress {
-			if _, err = execClient.ExecutePowershellCommand(SetSdnRemoteArpMacAddressCommand); err != nil {
-				log.Printf("Failed to set SDNRemoteArpMacAddress due to error %s", err.Error())
-				return err
-			}
+	defer service.Close()
+	return errors.Wrap(restartService(service), "could not restart service")
+}
 
-			log.Printf("[Azure CNS] SDNRemoteArpMacAddress regKey set successfully. Restarting hns service.")
-			if _, err := execClient.ExecutePowershellCommand(RestartHnsServiceCommand); err != nil {
-				log.Printf("Failed to Restart HNS Service due to error %s", err.Error())
-				return err
-			}
-		}
+type key interface {
+	SetStringValue(name, value string) error
+}
 
-		sdnRemoteArpMacAddressSet = true
+func setSDNRemoteARPMACAddress(k key) error {
+	// Set the reg key
+	// was "Set-ItemProperty -Path HKLM:\\SYSTEM\\CurrentControlSet\\Services\\hns\\State -Name SDNRemoteArpMacAddress -Value \"12-34-56-78-9a-bc\""
+	if err := k.SetStringValue("SDNRemoteArpMacAddress", SDNRemoteArpMacAddress); err != nil {
+		return errors.Wrap(err, "could not set registry key")
 	}
-
+	log.Printf("[Azure CNS] SDNRemoteArpMacAddress regKey set successfully. Restarting hns service.")
 	return nil
 }
 
+type serv interface {
+	Control(code svc.Cmd) (svc.Status, error)
+	Query() (svc.Status, error)
+	Start(...string) error
+}
+
+// straight out of chat gpt
+func restartService(s serv) error {
+	// Stop the service
+	_, err := s.Control(svc.Stop)
+	if err != nil {
+		return errors.Wrap(err, "could not stop service")
+	}
+	// Wait for the service to stop
+	for status, err := s.Query(); status.State != svc.Stopped; status, err = s.Query() {
+		if err != nil {
+			return errors.Wrap(err, "could not query service status")
+		}
+		time.Sleep(500 * time.Millisecond) //nolint:gomnd // 500ms
+	}
+	// Start the service again
+	return errors.Wrap(s.Start(), "could not start service")
+}
+
 func HasMellanoxAdapter() bool {
 	m := &mellanox.Mellanox{}
 	return hasNetworkAdapter(m)
@@ -364,6 +382,7 @@ func GetProcessNameByID(pidstr string) (string, error) {
 	pidstr = strings.Trim(pidstr, "\r\n")
 	cmd := fmt.Sprintf("Get-Process -Id %s|Format-List", pidstr)
 	p := NewExecClient(nil)
+	// TODO not riemovign this because it seems to only be called in test?
 	out, err := p.ExecutePowershellCommand(cmd)
 	if err != nil {
 		log.Printf("Process is not running. Output:%v, Error %v", out, err)

platform/os_windows_test.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"errors"
 	"os/exec"
-	"strings"
 	"testing"
 	"time"
 
@@ -16,6 +15,19 @@ import (
 
 var errTestFailure = errors.New("test failure")
 
+type mockKey struct {
+	values map[string]string
+}
+
+func (k *mockKey) SetStringValue(name, value string) error {
+	k.values[name] = value
+	return nil
+}
+
+func (k *mockKey) Close() error {
+	return nil
+}
+
 // Test if hasNetworkAdapter returns false on actual error or empty adapter name(an error)
 func TestHasNetworkAdapterReturnsError(t *testing.T) {
 	ctrl := gomock.NewController(t)
@@ -115,39 +127,11 @@ func TestExecuteCommandError(t *testing.T) {
 	require.ErrorIs(t, err, exec.ErrNotFound)
 }
 
-func TestSetSdnRemoteArpMacAddress_hnsNotEnabled(t *testing.T) {
-	mockExecClient := NewMockExecClient(false)
-	// testing skip setting SdnRemoteArpMacAddress when hns not enabled
-	mockExecClient.SetPowershellCommandResponder(func(_ string) (string, error) {
-		return "False", nil
-	})
-	err := SetSdnRemoteArpMacAddress(mockExecClient)
-	assert.NoError(t, err)
-	assert.Equal(t, false, sdnRemoteArpMacAddressSet)
-
-	// testing the scenario when there is an error in checking if hns is enabled or not
-	mockExecClient.SetPowershellCommandResponder(func(_ string) (string, error) {
-		return "", errTestFailure
-	})
-	err = SetSdnRemoteArpMacAddress(mockExecClient)
-	assert.ErrorAs(t, err, &errTestFailure)
-	assert.Equal(t, false, sdnRemoteArpMacAddressSet)
-}
-
-func TestSetSdnRemoteArpMacAddress_hnsEnabled(t *testing.T) {
-	mockExecClient := NewMockExecClient(false)
-	// happy path
-	mockExecClient.SetPowershellCommandResponder(func(cmd string) (string, error) {
-		if strings.Contains(cmd, "Test-Path") {
-			return "True", nil
-		}
-		return "", nil
-	})
-	err := SetSdnRemoteArpMacAddress(mockExecClient)
-	assert.NoError(t, err)
-	assert.Equal(t, true, sdnRemoteArpMacAddressSet)
-	// reset sdnRemoteArpMacAddressSet
-	sdnRemoteArpMacAddressSet = false
+func TestSetSDNRemoteARPMACAddress(t *testing.T) {
+	k := &mockKey{values: make(map[string]string)}
+	err := setSDNRemoteARPMACAddress(k)
+	require.NoError(t, err)
+	assert.Equal(t, SDNRemoteArpMacAddress, k.values["SDNRemoteArpMacAddress"])
 }
 
 func TestFetchPnpIDMapping(t *testing.T) {