@@ -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,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
0 commit comments