Skip to content

Commit cc99288

Browse files
authored
fix: remove PowerShell from Windows registry interactions (#2993)
* remove powershell fromwindows registry interactions Signed-off-by: Evan Baker <[email protected]> * address feedback Signed-off-by: Evan Baker <[email protected]> --------- Signed-off-by: Evan Baker <[email protected]>
1 parent 5011d9a commit cc99288

File tree

4 files changed

+93
-79
lines changed

4 files changed

+93
-79
lines changed

cns/service/main.go

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

789789
// Setting the remote ARP MAC address to 12-34-56-78-9a-bc on windows for external traffic if HNS is enabled
790-
execClient := platform.NewExecClient(nil)
791-
err = platform.SetSdnRemoteArpMacAddress(execClient)
790+
err = platform.SetSdnRemoteArpMacAddress(rootCtx)
792791
if err != nil {
793792
logger.Errorf("Failed to set remote ARP MAC address: %v", err)
794793
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(context.Context) error {
183183
return nil
184184
}
185185

platform/os_windows.go

Lines changed: 73 additions & 42 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,10 @@ 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 other utils from "golang.org/x/sys/windows"?
7769
GetMacAddressVFPPnpIDMapping = "Get-NetAdapter | Select-Object MacAddress, PnpDeviceID| Format-Table -HideTableHeaders"
7870

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

@@ -257,40 +246,82 @@ func (p *execClient) ExecutePowershellCommandWithContext(ctx context.Context, co
257246
}
258247

259248
// 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)
249+
func SetSdnRemoteArpMacAddress(ctx context.Context) error {
250+
log.Printf("Setting SDNRemoteArpMacAddress regKey")
251+
// open the registry key
252+
k, err := registry.OpenKey(registry.LOCAL_MACHINE, "SYSTEM\\CurrentControlSet\\Services\\hns\\State", registry.READ|registry.SET_VALUE)
262253
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)
254+
if errors.Is(err, registry.ErrNotExist) {
255+
return nil
256+
}
257+
return errors.Wrap(err, "could not open registry key")
266258
}
267-
if strings.EqualFold(exists, "false") {
268-
log.Printf("hns state path does not exist, skip setting SdnRemoteArpMacAddress")
269-
return nil
259+
defer k.Close()
260+
// set the key value
261+
if err = setSDNRemoteARPMACAddress(k); err != nil {
262+
return errors.Wrap(err, "could not set registry key")
270263
}
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-
}
264+
log.Printf("SDNRemoteArpMacAddress regKey set successfully")
265+
log.Printf("Restarting HNS service")
266+
// connect to the service manager
267+
m, err := mgr.Connect()
268+
if err != nil {
269+
return errors.Wrap(err, "could not connect to service manager")
270+
}
271+
defer m.Disconnect() //nolint:errcheck // ignore error
272+
// open the HNS service
273+
service, err := m.OpenService("hns")
274+
if err != nil {
275+
return errors.Wrap(err, "could not access service")
276+
}
277+
defer service.Close()
278+
if err := restartService(ctx, service); err != nil {
279+
return errors.Wrap(err, "could not restart service")
280+
}
281+
log.Printf("HNS service restarted successfully")
282+
return nil
283+
}
283284

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-
}
285+
type key interface {
286+
SetStringValue(name, value string) error
287+
}
290288

291-
sdnRemoteArpMacAddressSet = true
289+
func setSDNRemoteARPMACAddress(k key) error {
290+
// Set the reg key
291+
// was "Set-ItemProperty -Path HKLM:\\SYSTEM\\CurrentControlSet\\Services\\hns\\State -Name SDNRemoteArpMacAddress -Value \"12-34-56-78-9a-bc\""
292+
if err := k.SetStringValue("SDNRemoteArpMacAddress", SDNRemoteArpMacAddress); err != nil {
293+
return errors.Wrap(err, "could not set registry key")
292294
}
295+
return nil
296+
}
293297

298+
func restartService(ctx context.Context, s *mgr.Service) error {
299+
// Stop the service
300+
_, err := s.Control(svc.Stop)
301+
if err != nil {
302+
return errors.Wrap(err, "could not stop service")
303+
}
304+
// Wait for the service to stop
305+
ticker := time.NewTicker(500 * time.Millisecond) //nolint:gomnd // 500ms
306+
defer ticker.Stop()
307+
for { // hacky cancellable do-while
308+
status, err := s.Query()
309+
if err != nil {
310+
return errors.Wrap(err, "could not query service status")
311+
}
312+
if status.State == svc.Stopped {
313+
break
314+
}
315+
select {
316+
case <-ctx.Done():
317+
return errors.New("context cancelled")
318+
case <-ticker.C:
319+
}
320+
}
321+
// Start the service again
322+
if err := s.Start(); err != nil {
323+
return errors.Wrap(err, "could not start service")
324+
}
294325
return nil
295326
}
296327

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)