@@ -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,73 @@ 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+ // check the key value
261+ if v , _ , _ := k .GetStringValue ("SDNRemoteArpMacAddress" ); v == SDNRemoteArpMacAddress {
262+ log .Printf ("SDNRemoteArpMacAddress regKey already set" )
263+ return nil // already set
270264 }
271- if sdnRemoteArpMacAddressSet == false {
272- result , err := execClient .ExecutePowershellCommand (GetSdnRemoteArpMacAddressCommand )
265+ if err = k .SetStringValue ("SDNRemoteArpMacAddress" , SDNRemoteArpMacAddress ); err != nil {
266+ return errors .Wrap (err , "could not set registry key" )
267+ }
268+ log .Printf ("SDNRemoteArpMacAddress regKey set successfully" )
269+ log .Printf ("Restarting HNS service" )
270+ // connect to the service manager
271+ m , err := mgr .Connect ()
272+ if err != nil {
273+ return errors .Wrap (err , "could not connect to service manager" )
274+ }
275+ defer m .Disconnect () //nolint:errcheck // ignore error
276+ // open the HNS service
277+ service , err := m .OpenService ("hns" )
278+ if err != nil {
279+ return errors .Wrap (err , "could not access service" )
280+ }
281+ defer service .Close ()
282+ if err := restartService (ctx , service ); err != nil {
283+ return errors .Wrap (err , "could not restart service" )
284+ }
285+ log .Printf ("HNS service restarted successfully" )
286+ return nil
287+ }
288+
289+ func restartService (ctx context.Context , s * mgr.Service ) error {
290+ // Stop the service
291+ _ , err := s .Control (svc .Stop )
292+ if err != nil {
293+ return errors .Wrap (err , "could not stop service" )
294+ }
295+ // Wait for the service to stop
296+ ticker := time .NewTicker (500 * time .Millisecond ) //nolint:gomnd // 500ms
297+ defer ticker .Stop ()
298+ for { // hacky cancellable do-while
299+ status , err := s .Query ()
273300 if err != nil {
274- return err
301+ return errors . Wrap ( err , "could not query service status" )
275302 }
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- }
283-
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- }
303+ if status .State == svc .Stopped {
304+ break
305+ }
306+ select {
307+ case <- ctx .Done ():
308+ return errors .New ("context cancelled" )
309+ case <- ticker .C :
289310 }
290-
291- sdnRemoteArpMacAddressSet = true
292311 }
293-
312+ // Start the service again
313+ if err := s .Start (); err != nil {
314+ return errors .Wrap (err , "could not start service" )
315+ }
294316 return nil
295317}
296318
0 commit comments