winipcfg: port more granular locking from route change to others

Signed-off-by: Jason A. Donenfeld <Jason@zx2c4.com>

2 files changed, 36 insertions, 22 deletions

diff --git a/tunnel/winipcfg/interface_change_handler.go b/tunnel/winipcfg/interface_change_handler.go
index 57041719..6af5ea3c 100644
--- a/tunnel/winipcfg/interface_change_handler.go
+++ b/tunnel/winipcfg/interface_change_handler.go
@@ -17,42 +17,49 @@ type InterfaceChangeCallback struct {
 }
 
 var (
-	interfaceChangeMutex     = sync.Mutex{}
-	interfaceChangeCallbacks = make(map[*InterfaceChangeCallback]bool)
-	interfaceChangeHandle    = windows.Handle(0)
+	interfaceChangeAddRemoveMutex = sync.Mutex{}
+	interfaceChangeMutex          = sync.Mutex{}
+	interfaceChangeCallbacks      = make(map[*InterfaceChangeCallback]bool)
+	interfaceChangeHandle         = windows.Handle(0)
 )
 
 // RegisterInterfaceChangeCallback registers a new InterfaceChangeCallback. If this particular callback is already
 // registered, the function will silently return. Returned InterfaceChangeCallback.Unregister method should be used
 // to unregister.
 func RegisterInterfaceChangeCallback(callback func(notificationType MibNotificationType, iface *MibIPInterfaceRow)) (*InterfaceChangeCallback, error) {
-	cb := &InterfaceChangeCallback{callback}
+	s := &InterfaceChangeCallback{callback}
+
+	interfaceChangeAddRemoveMutex.Lock()
+	defer interfaceChangeAddRemoveMutex.Unlock()
 
 	interfaceChangeMutex.Lock()
 	defer interfaceChangeMutex.Unlock()
 
-	interfaceChangeCallbacks[cb] = true
+	interfaceChangeCallbacks[s] = true
 
 	if interfaceChangeHandle == 0 {
 		err := notifyIPInterfaceChange(windows.AF_UNSPEC, windows.NewCallback(interfaceChanged), 0, false, &interfaceChangeHandle)
 		if err != nil {
-			delete(interfaceChangeCallbacks, cb)
+			delete(interfaceChangeCallbacks, s)
 			interfaceChangeHandle = 0
 			return nil, err
 		}
 	}
 
-	return cb, nil
+	return s, nil
 }
 
 // Unregister unregisters the callback.
 func (callback *InterfaceChangeCallback) Unregister() error {
-	interfaceChangeMutex.Lock()
-	defer interfaceChangeMutex.Unlock()
+	interfaceChangeAddRemoveMutex.Lock()
+	defer interfaceChangeAddRemoveMutex.Unlock()
 
+	interfaceChangeMutex.Lock()
 	delete(interfaceChangeCallbacks, callback)
+	removeIt := len(interfaceChangeCallbacks) == 0 && interfaceChangeHandle != 0
+	interfaceChangeMutex.Unlock()
 
-	if len(interfaceChangeCallbacks) < 1 && interfaceChangeHandle != 0 {
+	if removeIt {
 		err := cancelMibChangeNotify2(interfaceChangeHandle)
 		if err != nil {
 			return err
diff --git a/tunnel/winipcfg/unicast_address_change_handler.go b/tunnel/winipcfg/unicast_address_change_handler.go
index fdd2e058..1ec34cd4 100644
--- a/tunnel/winipcfg/unicast_address_change_handler.go
+++ b/tunnel/winipcfg/unicast_address_change_handler.go
@@ -13,46 +13,53 @@ import (
 
 // UnicastAddressChangeCallback structure allows unicast address change callback handling.
 type UnicastAddressChangeCallback struct {
-	cb func(notificationType MibNotificationType, addr *MibUnicastIPAddressRow)
+	cb func(notificationType MibNotificationType, unicastAddress *MibUnicastIPAddressRow)
 }
 
 var (
-	unicastAddressChangeMutex     = sync.Mutex{}
-	unicastAddressChangeCallbacks = make(map[*UnicastAddressChangeCallback]bool)
-	unicastAddressChangeHandle    = windows.Handle(0)
+	unicastAddressChangeAddRemoveMutex = sync.Mutex{}
+	unicastAddressChangeMutex          = sync.Mutex{}
+	unicastAddressChangeCallbacks      = make(map[*UnicastAddressChangeCallback]bool)
+	unicastAddressChangeHandle         = windows.Handle(0)
 )
 
 // RegisterUnicastAddressChangeCallback registers a new UnicastAddressChangeCallback. If this particular callback is already
 // registered, the function will silently return. Returned UnicastAddressChangeCallback.Unregister method should be used
 // to unregister.
-func RegisterUnicastAddressChangeCallback(callback func(notificationType MibNotificationType, addr *MibUnicastIPAddressRow)) (*UnicastAddressChangeCallback, error) {
-	cb := &UnicastAddressChangeCallback{callback}
+func RegisterUnicastAddressChangeCallback(callback func(notificationType MibNotificationType, unicastAddress *MibUnicastIPAddressRow)) (*UnicastAddressChangeCallback, error) {
+	s := &UnicastAddressChangeCallback{callback}
+
+	unicastAddressChangeAddRemoveMutex.Lock()
+	defer unicastAddressChangeAddRemoveMutex.Unlock()
 
 	unicastAddressChangeMutex.Lock()
 	defer unicastAddressChangeMutex.Unlock()
 
-	unicastAddressChangeCallbacks[cb] = true
+	unicastAddressChangeCallbacks[s] = true
 
 	if unicastAddressChangeHandle == 0 {
 		err := notifyUnicastIPAddressChange(windows.AF_UNSPEC, windows.NewCallback(unicastAddressChanged), 0, false, &unicastAddressChangeHandle)
 		if err != nil {
-			delete(unicastAddressChangeCallbacks, cb)
+			delete(unicastAddressChangeCallbacks, s)
 			unicastAddressChangeHandle = 0
 			return nil, err
 		}
 	}
 
-	return cb, nil
+	return s, nil
 }
 
 // Unregister unregisters the callback.
 func (callback *UnicastAddressChangeCallback) Unregister() error {
-	unicastAddressChangeMutex.Lock()
-	defer unicastAddressChangeMutex.Unlock()
+	unicastAddressChangeAddRemoveMutex.Lock()
+	defer unicastAddressChangeAddRemoveMutex.Unlock()
 
+	unicastAddressChangeMutex.Lock()
 	delete(unicastAddressChangeCallbacks, callback)
+	removeIt := len(unicastAddressChangeCallbacks) == 0 && unicastAddressChangeHandle != 0
+	unicastAddressChangeMutex.Unlock()
 
-	if len(unicastAddressChangeCallbacks) < 1 && unicastAddressChangeHandle != 0 {
+	if removeIt {
 		err := cancelMibChangeNotify2(unicastAddressChangeHandle)
 		if err != nil {
 			return err