Reuse winpowrprof callback

Author: nekohasekai

common/winpowrprof/event.go

@@ -6,6 +6,8 @@ const (
 	EVENT_RESUME_AUTOMATIC // Because the user is not present, most applications should do nothing.
 )
 
+type EventCallback = func(event int)
+
 type EventListener interface {
 	Start() error
 	Close() error

common/winpowrprof/event_stub.go

@@ -6,6 +6,6 @@ import (
 	"os"
 )
 
-func NewEventListener(callback func(event int)) (EventListener, error) {
+func NewEventListener(callback EventCallback) (EventListener, error) {
 	return nil, os.ErrInvalid
 }

common/winpowrprof/event_windows.go

@@ -3,81 +3,134 @@ package winpowrprof
 // modify from https://github.com/golang/go/blob/b634f6fdcbebee23b7da709a243f3db217b64776/src/runtime/os_windows.go#L257
 
 import (
+	"sync"
 	"syscall"
 	"unsafe"
 
+	"github.com/sagernet/sing/common/x/list"
+
 	"golang.org/x/sys/windows"
 )
 
+type powerEventListener struct {
+	element *list.Element[EventCallback]
+}
+
+func NewEventListener(callback EventCallback) (EventListener, error) {
+	err := initCallback()
+	if err != nil {
+		return nil, err
+	}
+	access.Lock()
+	defer access.Unlock()
+	return &powerEventListener{
+		element: callbackList.PushBack(callback),
+	}, nil
+}
+
+func (l *powerEventListener) Start() error {
+	access.Lock()
+	defer access.Unlock()
+	if handle != 0 {
+		return nil
+	}
+	return startListener()
+}
+
+func (l *powerEventListener) Close() error {
+	access.Lock()
+	defer access.Unlock()
+	if l.element != nil {
+		callbackList.Remove(l.element)
+	}
+	if callbackList.Len() > 0 {
+		return nil
+	}
+	return closeListener()
+}
+
 var (
 	modpowerprof                                 = windows.NewLazySystemDLL("powrprof.dll")
 	procPowerRegisterSuspendResumeNotification   = modpowerprof.NewProc("PowerRegisterSuspendResumeNotification")
 	procPowerUnregisterSuspendResumeNotification = modpowerprof.NewProc("PowerUnregisterSuspendResumeNotification")
 )
 
-const (
-	PBT_APMSUSPEND         uint32 = 4
-	PBT_APMRESUMESUSPEND   uint32 = 7
-	PBT_APMRESUMEAUTOMATIC uint32 = 18
-)
-
-const (
-	_DEVICE_NOTIFY_CALLBACK = 2
+var (
+	access           sync.Mutex
+	callbackList     list.List[EventCallback]
+	initCallbackOnce sync.Once
+	rawCallback      uintptr
+	handle           uintptr
 )
 
-type _DEVICE_NOTIFY_SUBSCRIBE_PARAMETERS struct {
-	callback uintptr
-	context  uintptr
-}
-
-type eventListener struct {
-	params _DEVICE_NOTIFY_SUBSCRIBE_PARAMETERS
-	handle uintptr
-}
-
-func NewEventListener(callback func(event int)) (EventListener, error) {
-	if err := procPowerRegisterSuspendResumeNotification.Find(); err != nil {
-		return nil, err // Running on Windows 7, where we don't need it anyway.
+func initCallback() error {
+	err := procPowerRegisterSuspendResumeNotification.Find()
+	if err != nil {
+		return err // Running on Windows 7, where we don't need it anyway.
 	}
-	if err := procPowerUnregisterSuspendResumeNotification.Find(); err != nil {
-		return nil, err // Running on Windows 7, where we don't need it anyway.
+	err = procPowerUnregisterSuspendResumeNotification.Find()
+	if err != nil {
+		return err // Running on Windows 7, where we don't need it anyway.
 	}
-
-	var fn interface{} = func(context uintptr, changeType uint32, setting uintptr) uintptr {
-		switch changeType {
-		case PBT_APMSUSPEND:
-			callback(EVENT_SUSPEND)
-		case PBT_APMRESUMESUSPEND:
-			callback(EVENT_RESUME)
-		case PBT_APMRESUMEAUTOMATIC:
-			callback(EVENT_RESUME_AUTOMATIC)
-		}
-		return 0
-	}
-	return &eventListener{
-		params: _DEVICE_NOTIFY_SUBSCRIBE_PARAMETERS{
-			callback: windows.NewCallback(fn),
-		},
-	}, nil
+	initCallbackOnce.Do(func() {
+		rawCallback = windows.NewCallback(func(context uintptr, changeType uint32, setting uintptr) uintptr {
+			const (
+				PBT_APMSUSPEND         uint32 = 4
+				PBT_APMRESUMESUSPEND   uint32 = 7
+				PBT_APMRESUMEAUTOMATIC uint32 = 18
+			)
+			var event int
+			switch changeType {
+			case PBT_APMSUSPEND:
+				event = EVENT_SUSPEND
+			case PBT_APMRESUMESUSPEND:
+				event = EVENT_RESUME
+			case PBT_APMRESUMEAUTOMATIC:
+				event = EVENT_RESUME_AUTOMATIC
+			default:
+				return 0
+			}
+			access.Lock()
+			callbacks := callbackList.Array()
+			access.Unlock()
+			for _, callback := range callbacks {
+				callback(event)
+			}
+			return 0
+		})
+	})
+	return nil
 }
 
-func (l *eventListener) Start() error {
+func startListener() error {
+	type DEVICE_NOTIFY_SUBSCRIBE_PARAMETERS struct {
+		callback uintptr
+		context  uintptr
+	}
+	const DEVICE_NOTIFY_CALLBACK = 2
+	params := DEVICE_NOTIFY_SUBSCRIBE_PARAMETERS{
+		callback: rawCallback,
+	}
 	_, _, errno := syscall.SyscallN(
 		procPowerRegisterSuspendResumeNotification.Addr(),
-		_DEVICE_NOTIFY_CALLBACK,
-		uintptr(unsafe.Pointer(&l.params)),
-		uintptr(unsafe.Pointer(&l.handle)),
+		DEVICE_NOTIFY_CALLBACK,
+		uintptr(unsafe.Pointer(&params)),
+		uintptr(unsafe.Pointer(&handle)),
 	)
 	if errno != 0 {
 		return errno
 	}
 	return nil
 }
 
-func (l *eventListener) Close() error {
-	_, _, errno := syscall.SyscallN(procPowerUnregisterSuspendResumeNotification.Addr(), uintptr(unsafe.Pointer(&l.handle)))
+func closeListener() error {
+	if handle == 0 {
+		return nil
+	}
+	_, _, errno := syscall.SyscallN(procPowerUnregisterSuspendResumeNotification.Addr(), uintptr(unsafe.Pointer(&handle)))
 	if errno != 0 {
 		return errno
 	}
+	handle = 0
 	return nil
 }