Add Windows sleep detection for pre-sleep down/wake up

client/internal/sleep/detector_notsupported.go

@@ -1,4 +1,4 @@
-//go:build !darwin || ios
+//go:build (!darwin && !windows) || ios
 
 package sleep
 

client/internal/sleep/detector_windows.go

@@ -0,0 +1,249 @@
+//go:build windows
+
+package sleep
+
+import (
+	"fmt"
+	"sync"
+	"time"
+	"unsafe"
+
+	log "github.com/sirupsen/logrus"
+	"golang.org/x/sys/windows"
+)
+
+// Power broadcast event types delivered to the notification callback.
+// https://learn.microsoft.com/en-us/windows/win32/power/power-management-events
+const (
+	pbtAPMSuspend         uintptr = 0x0004
+	pbtAPMResumeSuspend   uintptr = 0x0007
+	pbtAPMResumeAutomatic uintptr = 0x0012
+	deviceNotifyCallback  uintptr = 0x00000002
+)
+
+var (
+	// The callback flavor of these functions is exported by user32.dll, not
+	// powrprof.dll (which exports the differently-shaped Power* variants).
+	user32 = windows.NewLazySystemDLL("user32.dll")
+
+	// https://learn.microsoft.com/en-us/windows/win32/api/winuser/nf-winuser-registersuspendresumenotification
+	registerSuspendResumeNotification = user32.NewProc("RegisterSuspendResumeNotification")
+	// https://learn.microsoft.com/en-us/windows/win32/api/winuser/nf-winuser-unregistersuspendresumenotification
+	unregisterSuspendResumeNotification = user32.NewProc("UnregisterSuspendResumeNotification")
+
+	libInitOnce sync.Once
+	libInitErr  error
+
+	// callbackThunk is the single C-callable trampoline registered with the OS.
+	// windows.NewCallback slots are a finite, non-reclaimable resource, so a
+	// single thunk dispatches to the Detector identified by the callback Context.
+	callbackThunk uintptr
+
+	// registry maps the Context value handed to the OS callback back to the
+	// Detector that registered it, mirroring darwin's serviceRegistry.
+	registry   = make(map[int]*Detector)
+	registryMu sync.Mutex
+	nextHandle int
+
+	// lifecycleMu serializes Register/Deregister so concurrent lifecycle calls
+	// can't race on the shared registry or the OS registration handle.
+	lifecycleMu sync.Mutex
+)
+
+// deviceNotifySubscribeParameters is DEVICE_NOTIFY_SUBSCRIBE_PARAMETERS.
+// https://learn.microsoft.com/en-us/windows/win32/api/winuser/ns-winuser-device_notify_subscribe_parameters
+type deviceNotifySubscribeParameters struct {
+	Callback uintptr
+	Context  uintptr
+}
+
+// Detector delivers sleep and wake events to a registered callback.
+type Detector struct {
+	callback func(event EventType)
+	done     chan struct{}
+
+	// handle keys this detector in the package registry and is passed to the OS
+	// as the callback Context. Zero means the detector is not registered.
+	handle int
+	// hPowerNotify is the HPOWERNOTIFY returned by RegisterSuspendResumeNotification.
+	hPowerNotify uintptr
+	// params is kept alive for the lifetime of the registration so the OS never
+	// dereferences freed memory.
+	params deviceNotifySubscribeParameters
+}
+
+// NewDetector resolves powrprof.dll symbols and returns a Detector.
+func NewDetector() (*Detector, error) {
+	if err := initLibs(); err != nil {
+		return nil, err
+	}
+	return &Detector{}, nil
+}
+
+func initLibs() error {
+	libInitOnce.Do(func() {
+		if err := registerSuspendResumeNotification.Find(); err != nil {
+			libInitErr = fmt.Errorf("resolve RegisterSuspendResumeNotification: %w", err)
+			return
+		}
+		if err := unregisterSuspendResumeNotification.Find(); err != nil {
+			libInitErr = fmt.Errorf("resolve UnregisterSuspendResumeNotification: %w", err)
+			return
+		}
+		callbackThunk = windows.NewCallback(powerCallback)
+	})
+	return libInitErr
+}
+
+// Register installs callback for power events and subscribes to suspend/resume
+// notifications via powrprof.dll.
+func (d *Detector) Register(callback func(event EventType)) error {
+	lifecycleMu.Lock()
+	defer lifecycleMu.Unlock()
+
+	registryMu.Lock()
+	if d.handle != 0 {
+		registryMu.Unlock()
+		return fmt.Errorf("detector service already registered")
+	}
+	d.callback = callback
+	d.done = make(chan struct{})
+	nextHandle++
+	handle := nextHandle
+	d.handle = handle
+	registry[handle] = d
+	registryMu.Unlock()
+
+	d.params = deviceNotifySubscribeParameters{
+		Callback: callbackThunk,
+		Context:  uintptr(handle),
+	}
+
+	ret, _, callErr := registerSuspendResumeNotification.Call(
+		uintptr(unsafe.Pointer(&d.params)),
+		deviceNotifyCallback,
+	)
+	if ret == 0 {
+		registryMu.Lock()
+		delete(registry, handle)
+		close(d.done)
+		d.done = nil
+		d.handle = 0
+		registryMu.Unlock()
+		return fmt.Errorf("RegisterSuspendResumeNotification failed: %w", callErr)
+	}
+	d.hPowerNotify = ret
+
+	log.Info("sleep detection service started on Windows")
+	return nil
+}
+
+// Deregister unsubscribes from power notifications and removes the detector.
+func (d *Detector) Deregister() error {
+	lifecycleMu.Lock()
+	defer lifecycleMu.Unlock()
+
+	registryMu.Lock()
+	if d.handle == 0 {
+		registryMu.Unlock()
+		return nil
+	}
+	close(d.done)
+	delete(registry, d.handle)
+	hPowerNotify := d.hPowerNotify
+	d.handle = 0
+	d.hPowerNotify = 0
+	registryMu.Unlock()
+
+	log.Info("sleep detection service stopping (deregister)")
+
+	if hPowerNotify == 0 {
+		return nil
+	}
+
+	ret, _, callErr := unregisterSuspendResumeNotification.Call(hPowerNotify)
+	if ret == 0 {
+		return fmt.Errorf("UnregisterSuspendResumeNotification failed: %w", callErr)
+	}
+	return nil
+}
+
+func (d *Detector) triggerCallback(event EventType, cb func(event EventType), done <-chan struct{}) {
+	if cb == nil || done == nil {
+		return
+	}
+
+	select {
+	case <-done:
+		return
+	default:
+	}
+
+	doneChan := make(chan struct{})
+	// The OS invokes this callback synchronously on the suspend path, so run the
+	// teardown inline with a bounded budget (mirroring the macOS detector) so
+	// Down completes before the machine suspends without blocking indefinitely.
+	timeout := time.NewTimer(20 * time.Second)
+	defer timeout.Stop()
+
+	go func() {
+		defer close(doneChan)
+		defer func() {
+			if r := recover(); r != nil {
+				log.Errorf("panic in sleep callback: %v", r)
+			}
+		}()
+		log.Info("sleep detection event fired")
+		cb(event)
+	}()
+
+	select {
+	case <-doneChan:
+	case <-done:
+	case <-timeout.C:
+		log.Warn("sleep callback timed out")
+	}
+}
+
+// powerCallback is the DEVICE_NOTIFY_CALLBACK_ROUTINE trampoline, invoked by the
+// OS on a system thread. A Go panic crossing the syscall boundary has undefined
+// behavior, so contain it here. It must return ERROR_SUCCESS (0).
+func powerCallback(context uintptr, msgType uintptr, setting uintptr) uintptr {
+	defer func() {
+		if r := recover(); r != nil {
+			log.Errorf("panic in sleep powerCallback: %v", r)
+		}
+	}()
+
+	var event EventType
+	switch msgType {
+	case pbtAPMSuspend:
+		event = EventTypeSleep
+	case pbtAPMResumeAutomatic, pbtAPMResumeSuspend:
+		event = EventTypeWakeUp
+	default:
+		return 0
+	}
+
+	dispatchEvent(int(context), event)
+	return 0
+}
+
+func dispatchEvent(handle int, event EventType) {
+	registryMu.Lock()
+	d := registry[handle]
+	var (
+		cb   func(event EventType)
+		done <-chan struct{}
+	)
+	if d != nil {
+		cb = d.callback
+		done = d.done
+	}
+	registryMu.Unlock()
+
+	if d == nil {
+		return
+	}
+	d.triggerCallback(event, cb, done)
+}

client/internal/sleep/detector_windows_test.go

@@ -0,0 +1,83 @@
+//go:build windows
+
+package sleep
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// registerFake inserts a detector into the registry without touching the OS,
+// so dispatch routing can be exercised in isolation. It returns the handle and
+// a cleanup that removes the entry.
+func registerFake(t *testing.T, cb func(EventType)) (int, func()) {
+	t.Helper()
+
+	registryMu.Lock()
+	nextHandle++
+	handle := nextHandle
+	d := &Detector{callback: cb, done: make(chan struct{}), handle: handle}
+	registry[handle] = d
+	registryMu.Unlock()
+
+	return handle, func() {
+		registryMu.Lock()
+		delete(registry, handle)
+		registryMu.Unlock()
+	}
+}
+
+func TestPowerCallback_MapsMessageTypes(t *testing.T) {
+	tests := []struct {
+		name    string
+		msgType uintptr
+		want    EventType
+		fires   bool
+	}{
+		{"suspend", pbtAPMSuspend, EventTypeSleep, true},
+		{"resume automatic", pbtAPMResumeAutomatic, EventTypeWakeUp, true},
+		{"resume suspend", pbtAPMResumeSuspend, EventTypeWakeUp, true},
+		{"unknown", 0x9999, EventTypeUnknown, false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := make(chan EventType, 1)
+			handle, cleanup := registerFake(t, func(e EventType) { got <- e })
+			defer cleanup()
+
+			ret := powerCallback(uintptr(handle), tt.msgType, 0)
+			require.Equal(t, uintptr(0), ret, "callback must return ERROR_SUCCESS")
+
+			if !tt.fires {
+				assert.Empty(t, got, "no event should fire for unhandled message type")
+				return
+			}
+			select {
+			case e := <-got:
+				assert.Equal(t, tt.want, e, "mapped event type should match")
+			default:
+				t.Fatal("expected callback to fire")
+			}
+		})
+	}
+}
+
+func TestDispatchEvent_UnknownHandleNoPanic(t *testing.T) {
+	require.NotPanics(t, func() {
+		dispatchEvent(-1, EventTypeSleep)
+	}, "dispatch for an unregistered handle must be a no-op")
+}
+
+func TestTriggerCallback_SkipsAfterDone(t *testing.T) {
+	done := make(chan struct{})
+	close(done)
+
+	fired := false
+	d := &Detector{}
+	d.triggerCallback(EventTypeSleep, func(EventType) { fired = true }, done)
+
+	assert.False(t, fired, "callback must not run once the detector is done")
+}