@@ -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
2528const (
@@ -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+
297318func 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 )
0 commit comments