winauth-go/internal/win32/tray_windows.go

//go:build windows

package win32

import (
	"fmt"
	"runtime"
	"sync"
	"sync/atomic"
	"syscall"
	"unsafe"

	"golang.org/x/sys/windows"
)

// TrayItem describes a single popup-menu entry. ID is what comes back on
// Events() when the user clicks it; Label is the visible text.
type TrayItem struct {
	ID    uint32
	Label string
}

// TrayEvent fires when the user activates a menu item or left-clicks the
// icon. ItemID == 0 means a bare icon click (caller convention: treat as
// "show main window").
type TrayEvent struct {
	ItemID uint32
}

// TrayManager owns a dedicated OS thread that runs a message-only window
// to host Shell_NotifyIcon and dispatch popup-menu commands. All public
// methods are safe to call from any goroutine; they post work onto the
// tray thread via PostMessage.
type TrayManager struct {
	tooltip string
	items   []TrayItem
	events  chan TrayEvent

	mu      sync.Mutex
	hwnd    uintptr // message-only window owned by tray thread
	menu    uintptr // HMENU
	stopped atomic.Bool
	started chan struct{}
	done    chan struct{}
}

// NewTrayManager spawns the tray thread, creates the icon, and returns
// once the icon is visible. tooltip is the hover string; items become
// popup-menu entries in order. The reserved ItemID values 0 (icon click)
// and 0xFFFF are off-limits to avoid colliding with our internal codes.
func NewTrayManager(tooltip string, items []TrayItem) (*TrayManager, error) {
	for _, it := range items {
		if it.ID == 0 || it.ID == cmdQuitSentinel {
			return nil, fmt.Errorf("win32: tray item ID %d is reserved", it.ID)
		}
	}
	m := &TrayManager{
		tooltip: tooltip,
		items:   items,
		events:  make(chan TrayEvent, 16),
		started: make(chan struct{}),
		done:    make(chan struct{}),
	}
	initErr := make(chan error, 1)
	go m.run(initErr)
	if err := <-initErr; err != nil {
		return nil, err
	}
	return m, nil
}

// Events yields a stream of clicks until Stop closes the channel.
func (m *TrayManager) Events() <-chan TrayEvent { return m.events }

// Stop tears the icon down and exits the tray thread. Safe to call more
// than once.
func (m *TrayManager) Stop() {
	if !m.stopped.CompareAndSwap(false, true) {
		return
	}
	m.mu.Lock()
	hwnd := m.hwnd
	m.mu.Unlock()
	if hwnd != 0 {
		procPostMessageW.Call(hwnd, uintptr(wmTrayQuit), 0, 0)
	}
	<-m.done
}

const (
	wmUser          uint32 = 0x0400
	wmTrayCallback         = wmUser + 1
	wmTrayQuit             = wmUser + 2
	cmdQuitSentinel uint32 = 0xFFFF
	nimAdd          uint32 = 0x00000000
	nimDelete       uint32 = 0x00000002
	nifMessage      uint32 = 0x00000001
	nifIcon         uint32 = 0x00000002
	nifTip          uint32 = 0x00000004
	wmCommand       uint32 = 0x0111
	wmRButtonUp     uint32 = 0x0205
	wmLButtonUp     uint32 = 0x0202
	mfString        uint32 = 0x00000000
	tpmLeftAlign    uint32 = 0x0000
	tpmRightButton  uint32 = 0x0002
	// idiApplication is the stock placeholder; only used if loading
	// our embedded icon resource fails. Our resource (assigned by
	// `akavel/rsrc`) lives at id 1.
	idiApplication uintptr = 32512
	appIconID      uintptr = 1
)

// trayClassName is unique enough to avoid clashing if another component
// in the same process happens to register classes.
var trayClassName, _ = syscall.UTF16PtrFromString("WinAuthGoTrayHost")

func (m *TrayManager) run(initErr chan<- error) {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()
	defer close(m.done)
	defer close(m.events)

	hinst, _, _ := procGetModuleHandleW.Call(0)
	if hinst == 0 {
		initErr <- fmt.Errorf("win32: GetModuleHandle failed")
		return
	}

	wndProc := syscall.NewCallback(m.wndProc)
	wc := wndClassExW{
		cbSize:        uint32(unsafe.Sizeof(wndClassExW{})),
		lpfnWndProc:   wndProc,
		hInstance:     hinst,
		hbrBackground: 0,
		lpszClassName: trayClassName,
	}
	wc.cbSize = uint32(unsafe.Sizeof(wc))
	atom, _, e := procRegisterClassExW.Call(uintptr(unsafe.Pointer(&wc)))
	if atom == 0 {
		// ERROR_CLASS_ALREADY_EXISTS is OK — a previous tray manager left
		// the class behind. Other errors are fatal.
		if errno, ok := e.(syscall.Errno); ok && errno != 1410 /*ERROR_CLASS_ALREADY_EXISTS*/ {
			initErr <- fmt.Errorf("win32: RegisterClassExW: %w", errno)
			return
		}
	}

	// HWND_MESSAGE = -3. A message-only window never paints but still
	// receives messages, which is exactly what Shell_NotifyIcon needs.
	const hwndMessage = ^uintptr(2) // ((HWND)-3)
	hwnd, _, e := procCreateWindowExW.Call(
		0,
		uintptr(unsafe.Pointer(trayClassName)),
		0,
		0,
		0, 0, 0, 0,
		hwndMessage,
		0,
		hinst,
		0,
	)
	if hwnd == 0 {
		initErr <- fmt.Errorf("win32: CreateWindowExW (tray host): %w", e)
		return
	}
	m.mu.Lock()
	m.hwnd = hwnd
	m.mu.Unlock()

	// Build the popup menu once. Item IDs are passed straight to
	// TrackPopupMenu and come back via WM_COMMAND's wParam low word.
	menu, _, _ := procCreatePopupMenu.Call()
	if menu == 0 {
		procDestroyWindow.Call(hwnd)
		initErr <- fmt.Errorf("win32: CreatePopupMenu failed")
		return
	}
	for _, it := range m.items {
		label, err := syscall.UTF16PtrFromString(it.Label)
		if err != nil {
			continue
		}
		procAppendMenuW.Call(menu, uintptr(mfString), uintptr(it.ID), uintptr(unsafe.Pointer(label)))
	}
	m.mu.Lock()
	m.menu = menu
	m.mu.Unlock()

	// Load the application icon embedded in rsrc_windows.syso. Falls
	// back to the stock application icon if the resource is missing
	// (e.g. cross-compile that omitted the .syso).
	hIcon, _, _ := procLoadIconW.Call(hinst, appIconID)
	if hIcon == 0 {
		hIcon, _, _ = procLoadIconW.Call(0, idiApplication)
	}

	nid := notifyIconData{
		cbSize:           uint32(unsafe.Sizeof(notifyIconData{})),
		hWnd:             hwnd,
		uID:              1,
		uFlags:           nifMessage | nifIcon | nifTip,
		uCallbackMessage: wmTrayCallback,
		hIcon:            hIcon,
	}
	copyTooltip(&nid.szTip, m.tooltip)
	r, _, _ := procShellNotifyIconW.Call(uintptr(nimAdd), uintptr(unsafe.Pointer(&nid)))
	if r == 0 {
		procDestroyMenu.Call(menu)
		procDestroyWindow.Call(hwnd)
		initErr <- fmt.Errorf("win32: Shell_NotifyIcon NIM_ADD failed")
		return
	}

	initErr <- nil

	// Standard GetMessage pump. Exits when our WM_TRAYQUIT handler calls
	// DestroyWindow which posts WM_QUIT via PostQuitMessage.
	var msg msgStruct
	for {
		r, _, _ := procGetMessageW.Call(uintptr(unsafe.Pointer(&msg)), 0, 0, 0)
		if r == 0 || int32(r) == -1 {
			break
		}
		procTranslateMessage.Call(uintptr(unsafe.Pointer(&msg)))
		procDispatchMessageW.Call(uintptr(unsafe.Pointer(&msg)))
	}

	// Cleanup: remove icon, destroy menu, the window already went down
	// during WM_TRAYQUIT.
	procShellNotifyIconW.Call(uintptr(nimDelete), uintptr(unsafe.Pointer(&nid)))
	procDestroyMenu.Call(menu)
}

// wndProc handles the small set of messages we care about. Anything else
// falls through to DefWindowProcW.
func (m *TrayManager) wndProc(hwnd uintptr, msg uint32, wParam, lParam uintptr) uintptr {
	switch msg {
	case wmTrayQuit:
		procDestroyWindow.Call(hwnd)
		procPostQuitMessage.Call(0)
		return 0
	case wmTrayCallback:
		// lParam low word holds the mouse message; wParam is the icon ID
		// (we only have one icon so we ignore it).
		mouse := uint32(lParam) & 0xFFFF
		switch mouse {
		case wmLButtonUp:
			select {
			case m.events <- TrayEvent{ItemID: 0}:
			default:
			}
		case wmRButtonUp:
			m.showPopupMenu(hwnd)
		}
		return 0
	case wmCommand:
		id := uint32(wParam) & 0xFFFF
		if id != 0 {
			select {
			case m.events <- TrayEvent{ItemID: id}:
			default:
			}
		}
		return 0
	}
	r, _, _ := procDefWindowProcW.Call(hwnd, uintptr(msg), wParam, lParam)
	return r
}

// showPopupMenu positions the menu at the cursor and posts the selection
// back to us as a WM_COMMAND. SetForegroundWindow before TrackPopupMenu
// is the documented dance to avoid the menu lingering when the user
// clicks elsewhere.
func (m *TrayManager) showPopupMenu(hwnd uintptr) {
	var pt struct{ X, Y int32 }
	procGetCursorPos.Call(uintptr(unsafe.Pointer(&pt)))
	procSetForegroundWindowProc.Call(hwnd)
	procTrackPopupMenu.Call(
		m.menu,
		uintptr(tpmLeftAlign|tpmRightButton),
		uintptr(pt.X), uintptr(pt.Y),
		0, hwnd, 0,
	)
	// MSDN: post a benign null message so the menu's modal loop exits
	// cleanly when the user clicks outside it.
	procPostMessageW.Call(hwnd, 0, 0, 0)
}

// copyTooltip writes s into the fixed-size buffer used by NOTIFYICONDATAW,
// truncating to 127 chars + NUL.
func copyTooltip(dst *[128]uint16, s string) {
	utf := utf16FromString(s)
	if len(utf) >= len(dst) {
		utf = utf[:len(dst)-1]
	}
	copy(dst[:], utf)
	dst[len(utf)] = 0
}

func utf16FromString(s string) []uint16 {
	r, _ := syscall.UTF16FromString(s)
	if n := len(r); n > 0 && r[n-1] == 0 {
		return r[:n-1]
	}
	return r
}

// -----------------------------------------------------------------------------
// raw syscalls + structs

var (
	procGetModuleHandleW = kernel32.NewProc("GetModuleHandleW")
	procRegisterClassExW = user32.NewProc("RegisterClassExW")
	procCreateWindowExW  = user32.NewProc("CreateWindowExW")
	procDestroyWindow    = user32.NewProc("DestroyWindow")
	procDefWindowProcW   = user32.NewProc("DefWindowProcW")
	procGetMessageW      = user32.NewProc("GetMessageW")
	procTranslateMessage = user32.NewProc("TranslateMessage")
	procDispatchMessageW = user32.NewProc("DispatchMessageW")
	procPostMessageW     = user32.NewProc("PostMessageW")
	procPostQuitMessage  = user32.NewProc("PostQuitMessage")
	procLoadIconW        = user32.NewProc("LoadIconW")
	procCreatePopupMenu  = user32.NewProc("CreatePopupMenu")
	procAppendMenuW      = user32.NewProc("AppendMenuW")
	procDestroyMenu      = user32.NewProc("DestroyMenu")
	procTrackPopupMenu   = user32.NewProc("TrackPopupMenu")
	procGetCursorPos     = user32.NewProc("GetCursorPos")

	shell32              = windows.NewLazySystemDLL("shell32.dll")
	procShellNotifyIconW = shell32.NewProc("Shell_NotifyIconW")
)

type wndClassExW struct {
	cbSize        uint32
	style         uint32
	lpfnWndProc   uintptr
	cbClsExtra    int32
	cbWndExtra    int32
	hInstance     uintptr
	hIcon         uintptr
	hCursor       uintptr
	hbrBackground uintptr
	lpszMenuName  *uint16
	lpszClassName *uint16
	hIconSm       uintptr
}

type notifyIconData struct {
	cbSize           uint32
	hWnd             uintptr
	uID              uint32
	uFlags           uint32
	uCallbackMessage uint32
	hIcon            uintptr
	szTip            [128]uint16
	dwState          uint32
	dwStateMask      uint32
	szInfo           [256]uint16
	uVersion         uint32
	szInfoTitle      [64]uint16
	dwInfoFlags      uint32
	guidItem         [16]byte
	hBalloonIcon     uintptr
}