Diff from github.com/getlantern/systray@6f0e5a3 to github.com/rupor-github/wsl-ssh-agent@a305054

README.txt

After I read https://github.com/rupor-github/wsl-ssh-agent/tree/a305054739d6ce1fa6261a8b4cb673df083b160e/systray
I was curious about modification made for systray bundled in wsl-ssh-agent and here are diffs.

systray.go.diff

--- github.com/getlantern/systray.6f0e5a3c556c78a9bb175697ca41a193bf74ffc9/systray.go	2020-03-06 09:14:40.922232883 +0900
+++ github.com/rupor-github/wsl-ssh-agent/systray/systray.go	2020-03-06 09:11:40.494702083 +0900
@@ -1,5 +1,5 @@
 /*
-Package systray is a cross platfrom Go library to place an icon and menu in the
+Package systray is a cross platform Go library to place an icon and menu in the
 notification area.
 Supports Windows, Mac OSX and Linux currently.
 Methods can be called from any goroutine except Run(), which should be called
@@ -11,8 +11,6 @@
 	"runtime"
 	"sync"
 	"sync/atomic"
-
-	"github.com/getlantern/golog"
 )
 
 var (
@@ -38,13 +36,53 @@
 	checked bool
 }
 
-var (
-	log = golog.LoggerFor("systray")
+// SessionEvent describes WM_WTSSESSION_CHANGE wParam.
+type SessionEvent int32
+
+// SessionEvent values.
+const (
+	SesConsoleConnect    SessionEvent = iota + 0x1 // WTS_CONSOLE_CONNECT (0x1)
+	SesConsoleDisconnect                           // WTS_CONSOLE_DISCONNECT (0x2)
+	SesRemoteConnect                               // WTS_REMOTE_CONNECT (0x3)
+	SesRemoteDisconnect                            // WTS_REMOTE_DISCONNECT (0x4)
+	SesLogon                                       // WTS_SESSION_LOGON (0x5)
+	SesLogoff                                      // WTS_SESSION_LOGOFF (0x6)
+	SesLock                                        // WTS_SESSION_LOCK (0x7)
+	SesUnlock                                      // WTS_SESSION_UNLOCK (0x8)
+	SesRemoteControl                               // WTS_SESSION_REMOTE_CONTROL (0x9)
+)
+
+func (e SessionEvent) String() string {
+	res := "UNKNOWN"
+	switch e {
+	case SesConsoleConnect:
+		res = "WTS_CONSOLE_CONNECT (0x1)"
+	case SesConsoleDisconnect:
+		res = "WTS_CONSOLE_DISCONNECT (0x2)"
+	case SesRemoteConnect:
+		res = "WTS_REMOTE_CONNECT (0x3)"
+	case SesRemoteDisconnect:
+		res = "WTS_REMOTE_DISCONNECT (0x4)"
+	case SesLogon:
+		res = "WTS_SESSION_LOGON (0x5)"
+	case SesLogoff:
+		res = "WTS_SESSION_LOGOFF (0x6)"
+	case SesLock:
+		res = "WTS_SESSION_LOCK (0x7)"
+	case SesUnlock:
+		res = "WTS_SESSION_UNLOCK (0x8)"
+	case SesRemoteControl:
+		res = "WTS_SESSION_REMOTE_CONTROL (0x9)"
+	}
+	return res
+}
 
-	systrayReady  func()
-	systrayExit   func()
-	menuItems     = make(map[int32]*MenuItem)
-	menuItemsLock sync.RWMutex
+var (
+	systrayReady   func()
+	systrayExit    func()
+	systraySession func(SessionEvent)
+	menuItems      = make(map[int32]*MenuItem)
+	menuItemsLock  sync.RWMutex
 
 	currentID = int32(-1)
 )
@@ -53,7 +91,7 @@
 // callback.
 // It blocks until systray.Quit() is called.
 // Should be called at the very beginning of main() to lock at main thread.
-func Run(onReady func(), onExit func()) {
+func Run(onReady func(), onExit func(), onSessionEvent func(SessionEvent)) {
 	runtime.LockOSThread()
 	atomic.StoreInt64(&hasStarted, 1)
 
@@ -78,6 +116,11 @@
 	}
 	systrayExit = onExit
 
+	if onSessionEvent == nil {
+		onSessionEvent = func(SessionEvent) {}
+	}
+	systraySession = onSessionEvent
+
 	nativeLoop()
 }
 
@@ -161,7 +204,7 @@
 	item.update()
 }
 
-// update propogates changes on a menu item to systray
+// update propagates changes on a menu item to systray
 func (item *MenuItem) update() {
 	menuItemsLock.Lock()
 	defer menuItemsLock.Unlock()

systray_windows.go.diff

--- github.com/getlantern/systray.6f0e5a3c556c78a9bb175697ca41a193bf74ffc9/systray_windows.go	2020-03-06 09:14:40.922232883 +0900
+++ github.com/rupor-github/wsl-ssh-agent/systray/systray_windows.go	2020-03-06 09:11:40.494702083 +0900
@@ -1,11 +1,13 @@
 // +build windows
 
+//nolint
 package systray
 
 import (
 	"crypto/md5"
 	"encoding/hex"
 	"io/ioutil"
+	"log"
 	"os"
 	"path/filepath"
 	"sort"
@@ -18,36 +20,39 @@
 // Helpful sources: https://github.com/golang/exp/blob/master/shiny/driver/internal/win32
 
 var (
-	k32                    = windows.NewLazySystemDLL("Kernel32.dll")
-	s32                    = windows.NewLazySystemDLL("Shell32.dll")
-	u32                    = windows.NewLazySystemDLL("User32.dll")
-	pGetModuleHandle       = k32.NewProc("GetModuleHandleW")
-	pShellNotifyIcon       = s32.NewProc("Shell_NotifyIconW")
-	pCreatePopupMenu       = u32.NewProc("CreatePopupMenu")
-	pCreateWindowEx        = u32.NewProc("CreateWindowExW")
-	pDefWindowProc         = u32.NewProc("DefWindowProcW")
-	pDeleteMenu            = u32.NewProc("DeleteMenu")
-	pDestroyWindow         = u32.NewProc("DestroyWindow")
-	pDispatchMessage       = u32.NewProc("DispatchMessageW")
-	pGetCursorPos          = u32.NewProc("GetCursorPos")
-	pGetMenuItemID         = u32.NewProc("GetMenuItemID")
-	pGetMessage            = u32.NewProc("GetMessageW")
-	pInsertMenuItem        = u32.NewProc("InsertMenuItemW")
-	pLoadIcon              = u32.NewProc("LoadIconW")
-	pLoadImage             = u32.NewProc("LoadImageW")
-	pLoadCursor            = u32.NewProc("LoadCursorW")
-	pPostMessage           = u32.NewProc("PostMessageW")
-	pPostQuitMessage       = u32.NewProc("PostQuitMessage")
-	pRegisterClass         = u32.NewProc("RegisterClassExW")
-	pRegisterWindowMessage = u32.NewProc("RegisterWindowMessageW")
-	pSetForegroundWindow   = u32.NewProc("SetForegroundWindow")
-	pSetMenuInfo           = u32.NewProc("SetMenuInfo")
-	pSetMenuItemInfo       = u32.NewProc("SetMenuItemInfoW")
-	pShowWindow            = u32.NewProc("ShowWindow")
-	pTrackPopupMenu        = u32.NewProc("TrackPopupMenu")
-	pTranslateMessage      = u32.NewProc("TranslateMessage")
-	pUnregisterClass       = u32.NewProc("UnregisterClassW")
-	pUpdateWindow          = u32.NewProc("UpdateWindow")
+	k32                            = windows.NewLazySystemDLL("Kernel32.dll")
+	s32                            = windows.NewLazySystemDLL("Shell32.dll")
+	u32                            = windows.NewLazySystemDLL("User32.dll")
+	wts32                          = windows.NewLazySystemDLL("Wtsapi32.dll")
+	pGetModuleHandle               = k32.NewProc("GetModuleHandleW")
+	pShellNotifyIcon               = s32.NewProc("Shell_NotifyIconW")
+	pCreatePopupMenu               = u32.NewProc("CreatePopupMenu")
+	pCreateWindowEx                = u32.NewProc("CreateWindowExW")
+	pDefWindowProc                 = u32.NewProc("DefWindowProcW")
+	pDeleteMenu                    = u32.NewProc("DeleteMenu")
+	pDestroyWindow                 = u32.NewProc("DestroyWindow")
+	pDispatchMessage               = u32.NewProc("DispatchMessageW")
+	pGetCursorPos                  = u32.NewProc("GetCursorPos")
+	pGetMenuItemID                 = u32.NewProc("GetMenuItemID")
+	pGetMessage                    = u32.NewProc("GetMessageW")
+	pInsertMenuItem                = u32.NewProc("InsertMenuItemW")
+	pLoadIcon                      = u32.NewProc("LoadIconW")
+	pLoadImage                     = u32.NewProc("LoadImageW")
+	pLoadCursor                    = u32.NewProc("LoadCursorW")
+	pPostMessage                   = u32.NewProc("PostMessageW")
+	pPostQuitMessage               = u32.NewProc("PostQuitMessage")
+	pRegisterClass                 = u32.NewProc("RegisterClassExW")
+	pRegisterWindowMessage         = u32.NewProc("RegisterWindowMessageW")
+	pSetForegroundWindow           = u32.NewProc("SetForegroundWindow")
+	pSetMenuInfo                   = u32.NewProc("SetMenuInfo")
+	pSetMenuItemInfo               = u32.NewProc("SetMenuItemInfoW")
+	pShowWindow                    = u32.NewProc("ShowWindow")
+	pTrackPopupMenu                = u32.NewProc("TrackPopupMenu")
+	pTranslateMessage              = u32.NewProc("TranslateMessage")
+	pUnregisterClass               = u32.NewProc("UnregisterClassW")
+	pUpdateWindow                  = u32.NewProc("UpdateWindow")
+	pRegisterSessionNotification   = wts32.NewProc("WTSRegisterSessionNotification")
+	pUnRegisterSessionNotification = wts32.NewProc("WTSUnRegisterSessionNotification")
 )
 
 // Contains window class information.
@@ -175,6 +180,8 @@
 	wmTaskbarCreated uint32
 
 	visibleItems []uint32
+
+	sessionRegistered bool
 }
 
 // Loads an image from file and shows it in tray.
@@ -236,11 +243,12 @@
 // https://msdn.microsoft.com/en-us/library/windows/desktop/ms633573(v=vs.85).aspx
 func (t *winTray) wndProc(hWnd windows.Handle, message uint32, wParam, lParam uintptr) (lResult uintptr) {
 	const (
-		WM_COMMAND    = 0x0111
-		WM_DESTROY    = 0x0002
-		WM_ENDSESSION = 0x16
-		WM_RBUTTONUP  = 0x0205
-		WM_LBUTTONUP  = 0x0202
+		WM_COMMAND           = 0x0111
+		WM_DESTROY           = 0x0002
+		WM_ENDSESSION        = 0x16
+		WM_RBUTTONUP         = 0x0205
+		WM_LBUTTONUP         = 0x0202
+		WM_WTSSESSION_CHANGE = 0x02B1
 	)
 	switch message {
 	case WM_COMMAND:
@@ -249,6 +257,12 @@
 			systrayMenuItemSelected(menuId)
 		}
 	case WM_DESTROY:
+		if t.sessionRegistered {
+			unregistered, _, _ := pUnRegisterSessionNotification.Call(uintptr(hWnd))
+			if unregistered != 0 {
+				t.sessionRegistered = false
+			}
+		}
 		// same as WM_ENDSESSION, but throws 0 exit code after all
 		defer pPostQuitMessage.Call(uintptr(int32(0)))
 		fallthrough
@@ -264,6 +278,9 @@
 		}
 	case t.wmTaskbarCreated: // on explorer.exe restarts
 		t.nid.add()
+	case WM_WTSSESSION_CHANGE:
+		systraySession(SessionEvent(wParam))
+		fallthrough
 	default:
 		// Calls the default window procedure to provide default processing for any window messages that an application does not process.
 		// https://msdn.microsoft.com/en-us/library/windows/desktop/ms633572(v=vs.85).aspx
@@ -300,6 +317,7 @@
 		CS_VREDRAW = 0x0001
 	)
 	const NIF_MESSAGE = 0x00000001
+	const NOTIFY_FOR_THIS_SESSION = 0
 
 	// https://msdn.microsoft.com/en-us/library/windows/desktop/ms644931(v=vs.85).aspx
 	const WM_USER = 0x0400
@@ -392,6 +410,15 @@
 		uintptr(t.window),
 	)
 
+	registered, _, err := pRegisterSessionNotification.Call(
+		uintptr(t.window),
+		uintptr(NOTIFY_FOR_THIS_SESSION),
+	)
+	if registered == 0 {
+		return err
+	}
+	t.sessionRegistered = true
+
 	t.nid = &notifyIconData{
 		Wnd:             windows.Handle(t.window),
 		ID:              100,
@@ -596,12 +623,12 @@
 
 func nativeLoop() {
 	if err := wt.initInstance(); err != nil {
-		log.Errorf("Unable to init instance: %v", err)
+		log.Printf("Unable to init instance: %v", err)
 		return
 	}
 
 	if err := wt.createMenu(); err != nil {
-		log.Errorf("Unable to create menu: %v", err)
+		log.Printf("Unable to create menu: %v", err)
 		return
 	}
 
@@ -630,7 +657,7 @@
 		// https://msdn.microsoft.com/en-us/library/windows/desktop/ms644936(v=vs.85).aspx
 		switch int32(ret) {
 		case -1:
-			log.Errorf("Error at message loop: %v", err)
+			log.Printf("Error at message loop: %v", err)
 			return
 		case 0:
 			return
@@ -662,13 +689,13 @@
 
 	if _, err := os.Stat(iconFilePath); os.IsNotExist(err) {
 		if err := ioutil.WriteFile(iconFilePath, iconBytes, 0644); err != nil {
-			log.Errorf("Unable to write icon data to temp file: %v", err)
+			log.Printf("Unable to write icon data to temp file: %v", err)
 			return
 		}
 	}
 
 	if err := wt.setIcon(iconFilePath); err != nil {
-		log.Errorf("Unable to set icon: %v", err)
+		log.Printf("Unable to set icon: %v", err)
 		return
 	}
 }
@@ -687,7 +714,7 @@
 // only available on Mac and Windows.
 func SetTooltip(tooltip string) {
 	if err := wt.setTooltip(tooltip); err != nil {
-		log.Errorf("Unable to set tooltip: %v", err)
+		log.Printf("Unable to set tooltip: %v", err)
 		return
 	}
 }
@@ -695,7 +722,7 @@
 func addOrUpdateMenuItem(item *MenuItem) {
 	err := wt.addOrUpdateMenuItem(item.id, item.title, item.disabled, item.checked)
 	if err != nil {
-		log.Errorf("Unable to addOrUpdateMenuItem: %v", err)
+		log.Printf("Unable to addOrUpdateMenuItem: %v", err)
 		return
 	}
 }
@@ -703,7 +730,7 @@
 func addSeparator(id int32) {
 	err := wt.addSeparatorMenuItem(id)
 	if err != nil {
-		log.Errorf("Unable to addSeparator: %v", err)
+		log.Printf("Unable to addSeparator: %v", err)
 		return
 	}
 }
@@ -711,7 +738,7 @@
 func hideMenuItem(item *MenuItem) {
 	err := wt.hideMenuItem(item.id)
 	if err != nil {
-		log.Errorf("Unable to hideMenuItem: %v", err)
+		log.Printf("Unable to hideMenuItem: %v", err)
 		return
 	}
 }