arblant
firefox-vpn-client / core / internal / desktop / endsession_windows_test.go
endsession_windows_test.go64 lines · 1.6 KB
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64
//go:build windows

package desktop

import (
	"sync/atomic"
	"testing"
	"time"
)

var (
	procSendMessageW = user32End.NewProc("SendMessageW")
	procIsWindow     = user32End.NewProc("IsWindow")
)

func sendMessage(hwnd uintptr, msg, wParam uintptr) uintptr {
	r, _, _ := procSendMessageW.Call(hwnd, msg, wParam, 0)
	return r
}

// The hidden window agrees to end the session, calls onEnd once when the
// session really ends (not when it is cancelled), and goes away on stop.
func TestEndSessionWindow(t *testing.T) {
	var calls atomic.Int32
	hwnd, stop, err := startEndSessionWindow(func() { calls.Add(1) })
	if err != nil {
		t.Fatal(err)
	}
	if r := sendMessage(hwnd, wmQueryEndSession, 0); r != 1 {
		t.Fatalf("WM_QUERYENDSESSION returned %d, want 1", r)
	}
	sendMessage(hwnd, wmEndSession, 0) // another program cancelled it
	if n := calls.Load(); n != 0 {
		t.Fatalf("onEnd ran %d times for a cancelled end", n)
	}
	sendMessage(hwnd, wmEndSession, 1)
	sendMessage(hwnd, wmEndSession, 1)
	if n := calls.Load(); n != 1 {
		t.Fatalf("onEnd ran %d times, want 1", n)
	}

	stop()
	deadline := time.Now().Add(2 * time.Second)
	for {
		if r, _, _ := procIsWindow.Call(hwnd); r == 0 {
			break
		}
		if time.Now().After(deadline) {
			t.Fatal("the window is still there after stop")
		}
		time.Sleep(10 * time.Millisecond)
	}

	// A second window (the class is registered once per process) works too.
	hwnd2, stop2, err := startEndSessionWindow(func() { calls.Add(1) })
	if err != nil {
		t.Fatal(err)
	}
	defer stop2()
	sendMessage(hwnd2, wmEndSession, 1)
	if n := calls.Load(); n != 2 {
		t.Fatalf("the second window's onEnd did not run (calls=%d)", n)
	}
}