Skip to content

Commit

Permalink
remove powershell fromwindows registry interactions
Browse files Browse the repository at this point in the history
Signed-off-by: Evan Baker <[email protected]>
  • Loading branch information
rbtr authored Sep 11, 2024
1 parent 47b4329 commit bea86f2
Show file tree
Hide file tree
Showing 4 changed files with 85 additions and 80 deletions.
3 changes: 1 addition & 2 deletions cns/service/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion platform/os_linux.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand Down
108 changes: 65 additions & 43 deletions platform/os_windows.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -257,43 +248,73 @@ 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 {
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)
if !errors.Is(err, registry.ErrNotExist) {
return nil
}
return errors.Wrap(err, "could not open registry key")
}
if strings.EqualFold(exists, "false") {
log.Printf("hns state path does not exist, skip setting SdnRemoteArpMacAddress")
return nil
defer k.Close()
// set the key value
if err = setSDNRemoteARPMACAddress(k); err != nil {
return errors.Wrap(err, "could not set registry key")
}
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
}
// connect to the service manager
m, err := mgr.Connect()
if err != nil {
return errors.Wrap(err, "could not connect to service manager")
}
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")
}
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)
Expand Down Expand Up @@ -364,6 +385,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)
Expand Down
52 changes: 18 additions & 34 deletions platform/os_windows_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@ import (
"context"
"errors"
"os/exec"
"strings"
"testing"
"time"

Expand All @@ -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)
Expand Down Expand Up @@ -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) {
Expand Down

0 comments on commit bea86f2

Please sign in to comment.