@@ -19,6 +19,9 @@ import (
1919 "github.com/pkg/errors"
2020 "go.uber.org/zap"
2121 "golang.org/x/sys/windows"
22+ "golang.org/x/sys/windows/registry"
23+ "golang.org/x/sys/windows/svc"
24+ "golang.org/x/sys/windows/svc/mgr"
2225)
2326
2427const (
@@ -60,21 +63,6 @@ const (
6063 // for vlan tagged arp requests
6164 SDNRemoteArpMacAddress = "12-34-56-78-9a-bc"
6265
63- // Command to get SDNRemoteArpMacAddress registry key
64- GetSdnRemoteArpMacAddressCommand = "(Get-ItemProperty " +
65- "-Path HKLM:\\ SYSTEM\\ CurrentControlSet\\ Services\\ hns\\ State -Name SDNRemoteArpMacAddress).SDNRemoteArpMacAddress"
66-
67- // Command to set SDNRemoteArpMacAddress registry key
68- SetSdnRemoteArpMacAddressCommand = "Set-ItemProperty " +
69- "-Path HKLM:\\ SYSTEM\\ CurrentControlSet\\ Services\\ hns\\ State -Name SDNRemoteArpMacAddress -Value \" 12-34-56-78-9a-bc\" "
70-
71- // Command to check if system has hns state path or not
72- CheckIfHNSStatePathExistsCommand = "Test-Path " +
73- "-Path HKLM:\\ SYSTEM\\ CurrentControlSet\\ Services\\ hns\\ State"
74-
75- // Command to restart HNS service
76- RestartHnsServiceCommand = "Restart-Service -Name hns"
77-
7866 // Interval between successive checks for mellanox adapter's PriorityVLANTag value
7967 defaultMellanoxMonitorInterval = 30 * time .Second
8068
@@ -195,40 +183,73 @@ func (p *execClient) ExecutePowershellCommand(command string) (string, error) {
195183}
196184
197185// SetSdnRemoteArpMacAddress sets the regkey for SDNRemoteArpMacAddress needed for multitenancy if hns is enabled
198- func SetSdnRemoteArpMacAddress (execClient ExecClient ) error {
199- exists , err := execClient .ExecutePowershellCommand (CheckIfHNSStatePathExistsCommand )
186+ func SetSdnRemoteArpMacAddress (ctx context.Context ) error {
187+ log .Printf ("Setting SDNRemoteArpMacAddress regKey" )
188+ // open the registry key
189+ k , err := registry .OpenKey (registry .LOCAL_MACHINE , `SYSTEM\CurrentControlSet\Services\hns\State` , registry .READ | registry .SET_VALUE )
200190 if err != nil {
201- errMsg := fmt .Sprintf ("Failed to check the existent of hns state path due to error %s" , err .Error ())
202- log .Printf (errMsg )
203- return errors .Errorf (errMsg )
191+ if errors .Is (err , registry .ErrNotExist ) {
192+ return nil
193+ }
194+ return errors .Wrap (err , "could not open registry key" )
204195 }
205- if strings .EqualFold (exists , "false" ) {
206- log .Printf ("hns state path does not exist, skip setting SdnRemoteArpMacAddress" )
207- return nil
196+ defer k .Close ()
197+ // check the key value
198+ if v , _ , _ := k .GetStringValue ("SDNRemoteArpMacAddress" ); v == SDNRemoteArpMacAddress {
199+ log .Printf ("SDNRemoteArpMacAddress regKey already set" )
200+ return nil // already set
208201 }
209- if sdnRemoteArpMacAddressSet == false {
210- result , err := execClient .ExecutePowershellCommand (GetSdnRemoteArpMacAddressCommand )
202+ if err = k .SetStringValue ("SDNRemoteArpMacAddress" , SDNRemoteArpMacAddress ); err != nil {
203+ return errors .Wrap (err , "could not set registry key" )
204+ }
205+ log .Printf ("SDNRemoteArpMacAddress regKey set successfully" )
206+ log .Printf ("Restarting HNS service" )
207+ // connect to the service manager
208+ m , err := mgr .Connect ()
209+ if err != nil {
210+ return errors .Wrap (err , "could not connect to service manager" )
211+ }
212+ defer m .Disconnect () //nolint:errcheck // ignore error
213+ // open the HNS service
214+ service , err := m .OpenService ("hns" )
215+ if err != nil {
216+ return errors .Wrap (err , "could not access service" )
217+ }
218+ defer service .Close ()
219+ if err := restartService (ctx , service ); err != nil {
220+ return errors .Wrap (err , "could not restart service" )
221+ }
222+ log .Printf ("HNS service restarted successfully" )
223+ return nil
224+ }
225+
226+ func restartService (ctx context.Context , s * mgr.Service ) error {
227+ // Stop the service
228+ _ , err := s .Control (svc .Stop )
229+ if err != nil {
230+ return errors .Wrap (err , "could not stop service" )
231+ }
232+ // Wait for the service to stop
233+ ticker := time .NewTicker (500 * time .Millisecond ) //nolint:gomnd // 500ms
234+ defer ticker .Stop ()
235+ for { // hacky cancellable do-while
236+ status , err := s .Query ()
211237 if err != nil {
212- return err
238+ return errors . Wrap ( err , "could not query service status" )
213239 }
214-
215- // Set the reg key if not already set or has incorrect value
216- if result != SDNRemoteArpMacAddress {
217- if _ , err = execClient .ExecutePowershellCommand (SetSdnRemoteArpMacAddressCommand ); err != nil {
218- log .Printf ("Failed to set SDNRemoteArpMacAddress due to error %s" , err .Error ())
219- return err
220- }
221-
222- log .Printf ("[Azure CNS] SDNRemoteArpMacAddress regKey set successfully. Restarting hns service." )
223- if _ , err := execClient .ExecutePowershellCommand (RestartHnsServiceCommand ); err != nil {
224- log .Printf ("Failed to Restart HNS Service due to error %s" , err .Error ())
225- return err
226- }
240+ if status .State == svc .Stopped {
241+ break
242+ }
243+ select {
244+ case <- ctx .Done ():
245+ return errors .New ("context cancelled" )
246+ case <- ticker .C :
227247 }
228-
229- sdnRemoteArpMacAddressSet = true
230248 }
231-
249+ // Start the service again
250+ if err := s .Start (); err != nil {
251+ return errors .Wrap (err , "could not start service" )
252+ }
232253 return nil
233254}
234255
0 commit comments