github.com/apernet/sing-tun@v0.2.6-0.20240323130332-b9f6511036ad/internal/winipcfg/unicast_address_change_handler.go (about)

     1  /* SPDX-License-Identifier: MIT
     2   *
     3   * Copyright (C) 2019-2022 WireGuard LLC. All Rights Reserved.
     4   */
     5  
     6  package winipcfg
     7  
     8  import (
     9  	"sync"
    10  
    11  	"golang.org/x/sys/windows"
    12  )
    13  
    14  // UnicastAddressChangeCallback structure allows unicast address change callback handling.
    15  type UnicastAddressChangeCallback struct {
    16  	cb   func(notificationType MibNotificationType, unicastAddress *MibUnicastIPAddressRow)
    17  	wait sync.WaitGroup
    18  }
    19  
    20  var (
    21  	unicastAddressChangeAddRemoveMutex = sync.Mutex{}
    22  	unicastAddressChangeMutex          = sync.Mutex{}
    23  	unicastAddressChangeCallbacks      = make(map[*UnicastAddressChangeCallback]bool)
    24  	unicastAddressChangeHandle         = windows.Handle(0)
    25  )
    26  
    27  // RegisterUnicastAddressChangeCallback registers a new UnicastAddressChangeCallback. If this particular callback is already
    28  // registered, the function will silently return. Returned UnicastAddressChangeCallback.Unregister method should be used
    29  // to unregister.
    30  func RegisterUnicastAddressChangeCallback(callback func(notificationType MibNotificationType, unicastAddress *MibUnicastIPAddressRow)) (*UnicastAddressChangeCallback, error) {
    31  	s := &UnicastAddressChangeCallback{cb: callback}
    32  
    33  	unicastAddressChangeAddRemoveMutex.Lock()
    34  	defer unicastAddressChangeAddRemoveMutex.Unlock()
    35  
    36  	unicastAddressChangeMutex.Lock()
    37  	defer unicastAddressChangeMutex.Unlock()
    38  
    39  	unicastAddressChangeCallbacks[s] = true
    40  
    41  	if unicastAddressChangeHandle == 0 {
    42  		err := notifyUnicastIPAddressChange(windows.AF_UNSPEC, windows.NewCallback(unicastAddressChanged), 0, false, &unicastAddressChangeHandle)
    43  		if err != nil {
    44  			delete(unicastAddressChangeCallbacks, s)
    45  			unicastAddressChangeHandle = 0
    46  			return nil, err
    47  		}
    48  	}
    49  
    50  	return s, nil
    51  }
    52  
    53  // Unregister unregisters the callback.
    54  func (callback *UnicastAddressChangeCallback) Unregister() error {
    55  	unicastAddressChangeAddRemoveMutex.Lock()
    56  	defer unicastAddressChangeAddRemoveMutex.Unlock()
    57  
    58  	unicastAddressChangeMutex.Lock()
    59  	delete(unicastAddressChangeCallbacks, callback)
    60  	removeIt := len(unicastAddressChangeCallbacks) == 0 && unicastAddressChangeHandle != 0
    61  	unicastAddressChangeMutex.Unlock()
    62  
    63  	callback.wait.Wait()
    64  
    65  	if removeIt {
    66  		err := cancelMibChangeNotify2(unicastAddressChangeHandle)
    67  		if err != nil {
    68  			return err
    69  		}
    70  		unicastAddressChangeHandle = 0
    71  	}
    72  
    73  	return nil
    74  }
    75  
    76  func unicastAddressChanged(callerContext uintptr, row *MibUnicastIPAddressRow, notificationType MibNotificationType) uintptr {
    77  	rowCopy := *row
    78  	unicastAddressChangeMutex.Lock()
    79  	for cb := range unicastAddressChangeCallbacks {
    80  		cb.wait.Add(1)
    81  		go func(cb *UnicastAddressChangeCallback) {
    82  			cb.cb(notificationType, &rowCopy)
    83  			cb.wait.Done()
    84  		}(cb)
    85  	}
    86  	unicastAddressChangeMutex.Unlock()
    87  	return 0
    88  }