Skip to content

Commit ee7e885

Browse files
authored
remove powershell fromwindows registry interactions
Signed-off-by: Evan Baker <[email protected]>
1 parent 47b4329 commit ee7e885

File tree

4 files changed

+85
-80
lines changed

4 files changed

+85
-80
lines changed

cns/service/main.go

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -789,8 +789,7 @@ func main() {
789789
}
790790

791791
// Setting the remote ARP MAC address to 12-34-56-78-9a-bc on windows for external traffic if HNS is enabled
792-
execClient := platform.NewExecClient(nil)
793-
err = platform.SetSdnRemoteArpMacAddress(execClient)
792+
err = platform.SetSdnRemoteArpMacAddress()
794793
if err != nil {
795794
logger.Errorf("Failed to set remote ARP MAC address: %v", err)
796795
return

platform/os_linux.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -179,7 +179,7 @@ func (p *execClient) KillProcessByName(processName string) error {
179179

180180
// SetSdnRemoteArpMacAddress sets the regkey for SDNRemoteArpMacAddress needed for multitenancy
181181
// This operation is specific to windows OS
182-
func SetSdnRemoteArpMacAddress(_ ExecClient) error {
182+
func SetSdnRemoteArpMacAddress() error {
183183
return nil
184184
}
185185

platform/os_windows.go

Lines changed: 65 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,9 @@ import (
2020
"github.com/pkg/errors"
2121
"go.uber.org/zap"
2222
"golang.org/x/sys/windows"
23+
"golang.org/x/sys/windows/registry"
24+
"golang.org/x/sys/windows/svc"
25+
"golang.org/x/sys/windows/svc/mgr"
2326
)
2427

2528
const (
@@ -61,24 +64,12 @@ const (
6164
// for vlan tagged arp requests
6265
SDNRemoteArpMacAddress = "12-34-56-78-9a-bc"
6366

64-
// Command to get SDNRemoteArpMacAddress registry key
65-
GetSdnRemoteArpMacAddressCommand = "(Get-ItemProperty " +
66-
"-Path HKLM:\\SYSTEM\\CurrentControlSet\\Services\\hns\\State -Name SDNRemoteArpMacAddress).SDNRemoteArpMacAddress"
67-
68-
// Command to set SDNRemoteArpMacAddress registry key
69-
SetSdnRemoteArpMacAddressCommand = "Set-ItemProperty " +
70-
"-Path HKLM:\\SYSTEM\\CurrentControlSet\\Services\\hns\\State -Name SDNRemoteArpMacAddress -Value \"12-34-56-78-9a-bc\""
71-
72-
// Command to check if system has hns state path or not
73-
CheckIfHNSStatePathExistsCommand = "Test-Path " +
74-
"-Path HKLM:\\SYSTEM\\CurrentControlSet\\Services\\hns\\State"
75-
7667
// Command to fetch netadapter and pnp id
68+
// TODO can we replace this (and things in endpoint_windows) with "golang.org/x/sys/windows"
69+
// var adapterInfo windows.IpAdapterInfo
70+
// var bufferSize uint32 = uint32(unsafe.Sizeof(adapterInfo))
7771
GetMacAddressVFPPnpIDMapping = "Get-NetAdapter | Select-Object MacAddress, PnpDeviceID| Format-Table -HideTableHeaders"
7872

79-
// Command to restart HNS service
80-
RestartHnsServiceCommand = "Restart-Service -Name hns"
81-
8273
// Interval between successive checks for mellanox adapter's PriorityVLANTag value
8374
defaultMellanoxMonitorInterval = 30 * time.Second
8475

@@ -257,43 +248,73 @@ func (p *execClient) ExecutePowershellCommandWithContext(ctx context.Context, co
257248
}
258249

259250
// SetSdnRemoteArpMacAddress sets the regkey for SDNRemoteArpMacAddress needed for multitenancy if hns is enabled
260-
func SetSdnRemoteArpMacAddress(execClient ExecClient) error {
261-
exists, err := execClient.ExecutePowershellCommand(CheckIfHNSStatePathExistsCommand)
251+
func SetSdnRemoteArpMacAddress() error {
252+
// open the registry key
253+
k, err := registry.OpenKey(registry.LOCAL_MACHINE, "SYSTEM\\CurrentControlSet\\Services\\hns\\State", registry.READ|registry.SET_VALUE)
262254
if err != nil {
263-
errMsg := fmt.Sprintf("Failed to check the existent of hns state path due to error %s", err.Error())
264-
log.Printf(errMsg)
265-
return errors.Errorf(errMsg)
255+
if errors.Is(err, registry.ErrNotExist) {
256+
return nil
257+
}
258+
return errors.Wrap(err, "could not open registry key")
266259
}
267-
if strings.EqualFold(exists, "false") {
268-
log.Printf("hns state path does not exist, skip setting SdnRemoteArpMacAddress")
269-
return nil
260+
defer k.Close()
261+
// set the key value
262+
if err = setSDNRemoteARPMACAddress(k); err != nil {
263+
return errors.Wrap(err, "could not set registry key")
270264
}
271-
if sdnRemoteArpMacAddressSet == false {
272-
result, err := execClient.ExecutePowershellCommand(GetSdnRemoteArpMacAddressCommand)
273-
if err != nil {
274-
return err
275-
}
276-
277-
// Set the reg key if not already set or has incorrect value
278-
if result != SDNRemoteArpMacAddress {
279-
if _, err = execClient.ExecutePowershellCommand(SetSdnRemoteArpMacAddressCommand); err != nil {
280-
log.Printf("Failed to set SDNRemoteArpMacAddress due to error %s", err.Error())
281-
return err
282-
}
265+
// connect to the service manager
266+
m, err := mgr.Connect()
267+
if err != nil {
268+
return errors.Wrap(err, "could not connect to service manager")
269+
}
270+
defer m.Disconnect() //nolint:errcheck // ignore error
271+
// open the HNS service
272+
service, err := m.OpenService("hns")
273+
if err != nil {
274+
return errors.Wrap(err, "could not access service")
275+
}
276+
defer service.Close()
277+
return errors.Wrap(restartService(service), "could not restart service")
278+
}
283279

284-
log.Printf("[Azure CNS] SDNRemoteArpMacAddress regKey set successfully. Restarting hns service.")
285-
if _, err := execClient.ExecutePowershellCommand(RestartHnsServiceCommand); err != nil {
286-
log.Printf("Failed to Restart HNS Service due to error %s", err.Error())
287-
return err
288-
}
289-
}
280+
type key interface {
281+
SetStringValue(name, value string) error
282+
}
290283

291-
sdnRemoteArpMacAddressSet = true
284+
func setSDNRemoteARPMACAddress(k key) error {
285+
// Set the reg key
286+
// was "Set-ItemProperty -Path HKLM:\\SYSTEM\\CurrentControlSet\\Services\\hns\\State -Name SDNRemoteArpMacAddress -Value \"12-34-56-78-9a-bc\""
287+
if err := k.SetStringValue("SDNRemoteArpMacAddress", SDNRemoteArpMacAddress); err != nil {
288+
return errors.Wrap(err, "could not set registry key")
292289
}
293-
290+
log.Printf("[Azure CNS] SDNRemoteArpMacAddress regKey set successfully. Restarting hns service.")
294291
return nil
295292
}
296293

294+
type serv interface {
295+
Control(code svc.Cmd) (svc.Status, error)
296+
Query() (svc.Status, error)
297+
Start(...string) error
298+
}
299+
300+
// straight out of chat gpt
301+
func restartService(s serv) error {
302+
// Stop the service
303+
_, err := s.Control(svc.Stop)
304+
if err != nil {
305+
return errors.Wrap(err, "could not stop service")
306+
}
307+
// Wait for the service to stop
308+
for status, err := s.Query(); status.State != svc.Stopped; status, err = s.Query() {
309+
if err != nil {
310+
return errors.Wrap(err, "could not query service status")
311+
}
312+
time.Sleep(500 * time.Millisecond) //nolint:gomnd // 500ms
313+
}
314+
// Start the service again
315+
return errors.Wrap(s.Start(), "could not start service")
316+
}
317+
297318
func HasMellanoxAdapter() bool {
298319
m := &mellanox.Mellanox{}
299320
return hasNetworkAdapter(m)
@@ -364,6 +385,7 @@ func GetProcessNameByID(pidstr string) (string, error) {
364385
pidstr = strings.Trim(pidstr, "\r\n")
365386
cmd := fmt.Sprintf("Get-Process -Id %s|Format-List", pidstr)
366387
p := NewExecClient(nil)
388+
// TODO not riemovign this because it seems to only be called in test?
367389
out, err := p.ExecutePowershellCommand(cmd)
368390
if err != nil {
369391
log.Printf("Process is not running. Output:%v, Error %v", out, err)

platform/os_windows_test.go

Lines changed: 18 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ import (
44
"context"
55
"errors"
66
"os/exec"
7-
"strings"
87
"testing"
98
"time"
109

@@ -16,6 +15,19 @@ import (
1615

1716
var errTestFailure = errors.New("test failure")
1817

18+
type mockKey struct {
19+
values map[string]string
20+
}
21+
22+
func (k *mockKey) SetStringValue(name, value string) error {
23+
k.values[name] = value
24+
return nil
25+
}
26+
27+
func (k *mockKey) Close() error {
28+
return nil
29+
}
30+
1931
// Test if hasNetworkAdapter returns false on actual error or empty adapter name(an error)
2032
func TestHasNetworkAdapterReturnsError(t *testing.T) {
2133
ctrl := gomock.NewController(t)
@@ -115,39 +127,11 @@ func TestExecuteCommandError(t *testing.T) {
115127
require.ErrorIs(t, err, exec.ErrNotFound)
116128
}
117129

118-
func TestSetSdnRemoteArpMacAddress_hnsNotEnabled(t *testing.T) {
119-
mockExecClient := NewMockExecClient(false)
120-
// testing skip setting SdnRemoteArpMacAddress when hns not enabled
121-
mockExecClient.SetPowershellCommandResponder(func(_ string) (string, error) {
122-
return "False", nil
123-
})
124-
err := SetSdnRemoteArpMacAddress(mockExecClient)
125-
assert.NoError(t, err)
126-
assert.Equal(t, false, sdnRemoteArpMacAddressSet)
127-
128-
// testing the scenario when there is an error in checking if hns is enabled or not
129-
mockExecClient.SetPowershellCommandResponder(func(_ string) (string, error) {
130-
return "", errTestFailure
131-
})
132-
err = SetSdnRemoteArpMacAddress(mockExecClient)
133-
assert.ErrorAs(t, err, &errTestFailure)
134-
assert.Equal(t, false, sdnRemoteArpMacAddressSet)
135-
}
136-
137-
func TestSetSdnRemoteArpMacAddress_hnsEnabled(t *testing.T) {
138-
mockExecClient := NewMockExecClient(false)
139-
// happy path
140-
mockExecClient.SetPowershellCommandResponder(func(cmd string) (string, error) {
141-
if strings.Contains(cmd, "Test-Path") {
142-
return "True", nil
143-
}
144-
return "", nil
145-
})
146-
err := SetSdnRemoteArpMacAddress(mockExecClient)
147-
assert.NoError(t, err)
148-
assert.Equal(t, true, sdnRemoteArpMacAddressSet)
149-
// reset sdnRemoteArpMacAddressSet
150-
sdnRemoteArpMacAddressSet = false
130+
func TestSetSDNRemoteARPMACAddress(t *testing.T) {
131+
k := &mockKey{values: make(map[string]string)}
132+
err := setSDNRemoteARPMACAddress(k)
133+
require.NoError(t, err)
134+
assert.Equal(t, SDNRemoteArpMacAddress, k.values["SDNRemoteArpMacAddress"])
151135
}
152136

153137
func TestFetchPnpIDMapping(t *testing.T) {

0 commit comments

Comments
 (0)