From 7220d62e96d214c790175aee5bce3ab935e43e82 Mon Sep 17 00:00:00 2001
From: "Jason A. Donenfeld" <Jason@zx2c4.com>
Date: Fri, 10 May 2019 14:00:33 +0200
Subject: ui: dont use main window and run our own message loop

This allows us to do custom things like send raise click messages.
---
 ui/managewindow.go | 57 +++++++++++++++++++++++++++++++++++++++++-------------
 ui/raise.go        | 10 +++++-----
 ui/tray.go         |  2 +-
 3 files changed, 50 insertions(+), 19 deletions(-)

diff --git a/ui/managewindow.go b/ui/managewindow.go
index b13e0afd..7f0bd608 100644
--- a/ui/managewindow.go
+++ b/ui/managewindow.go
@@ -12,7 +12,7 @@ import (
 )
 
 type ManageTunnelsWindow struct {
-	*walk.MainWindow
+	walk.FormBase
 
 	tabs        *walk.TabWidget
 	tunnelsPage *TunnelsPage
@@ -22,28 +22,36 @@ type ManageTunnelsWindow struct {
 	tunnelChangedCB *service.TunnelChangeCallback
 }
 
+const (
+	manageWindowWindowClass = "WireGuard UI - Manage Tunnels"
+	selectCorrectTabMsg     = win.WM_USER + 8993
+)
+
+func init() {
+	walk.MustRegisterWindowClass(manageWindowWindowClass)
+}
+
 func NewManageTunnelsWindow() (*ManageTunnelsWindow, error) {
 	var err error
 
-	var disposables walk.Disposables
-	defer disposables.Treat()
+	font, err := walk.NewFont("Segoe UI", 9, 0)
+	if err != nil {
+		return nil, err
+	}
 
-	mtw := &ManageTunnelsWindow{}
+	mtw := new(ManageTunnelsWindow)
+	mtw.SetName("WireGuard")
 
-	mtw.MainWindow, err = walk.NewMainWindowWithName("WireGuard")
+	err = walk.InitWindow(mtw, nil, manageWindowWindowClass, win.WS_OVERLAPPEDWINDOW, win.WS_EX_CONTROLPARENT)
 	if err != nil {
 		return nil, err
 	}
-	disposables.Add(mtw)
+	mtw.SetPersistent(true)
 
 	if icon, err := loadLogoIcon(mtw.DPI() / 3); err == nil { //TODO: calculate DPI dynamically
 		mtw.SetIcon(icon)
 	}
 	mtw.SetTitle("WireGuard")
-	font, err := walk.NewFont("Segoe UI", 9, 0)
-	if err != nil {
-		return nil, err
-	}
 	mtw.AddDisposable(font)
 	mtw.SetFont(font)
 	mtw.SetSize(walk.Size{640, 480})
@@ -77,8 +85,6 @@ func NewManageTunnelsWindow() (*ManageTunnelsWindow, error) {
 	}
 	mtw.tabs.Pages().Add(mtw.logPage.TabPage)
 
-	disposables.Spare()
-
 	mtw.tunnelChangedCB = service.IPCClientRegisterTunnelChange(mtw.onTunnelChange)
 	globalState, _ := service.IPCClientGlobalState()
 	mtw.onTunnelChange(nil, service.TunnelUnknown, globalState, nil)
@@ -91,7 +97,7 @@ func (mtw *ManageTunnelsWindow) Dispose() {
 		mtw.tunnelChangedCB.Unregister()
 		mtw.tunnelChangedCB = nil
 	}
-	mtw.MainWindow.Dispose()
+	mtw.FormBase.Dispose()
 }
 
 func (mtw *ManageTunnelsWindow) onTunnelChange(tunnel *service.Tunnel, state service.TunnelState, globalState service.TunnelState, err error) {
@@ -122,3 +128,28 @@ func (mtw *ManageTunnelsWindow) UpdateFound() {
 		mtw.tabs.Pages().Add(updatePage.TabPage)
 	}
 }
+
+func (mtw *ManageTunnelsWindow) WndProc(hwnd win.HWND, msg uint32, wParam, lParam uintptr) uintptr {
+	switch msg {
+	case win.WM_QUERYENDSESSION:
+		if lParam == win.ENDSESSION_CLOSEAPP {
+			return win.TRUE
+		}
+	case win.WM_ENDSESSION:
+		if lParam == win.ENDSESSION_CLOSEAPP && wParam == 1 {
+			walk.App().Exit(198)
+		}
+	case selectCorrectTabMsg:
+		if !mtw.Visible() {
+			mtw.tunnelsPage.listView.SelectFirstActiveTunnel()
+			if mtw.tabs.Pages().Len() != 3 {
+				mtw.tabs.SetCurrentIndex(0)
+			}
+		}
+		if mtw.tabs.Pages().Len() == 3 {
+			mtw.tabs.SetCurrentIndex(2)
+		}
+	}
+
+	return mtw.FormBase.WndProc(hwnd, msg, wParam, lParam)
+}
diff --git a/ui/raise.go b/ui/raise.go
index f902546e..c4a45833 100644
--- a/ui/raise.go
+++ b/ui/raise.go
@@ -15,8 +15,6 @@ import (
 	"golang.org/x/sys/windows"
 )
 
-const wireguardUIClass = "WireGuard UI - MainWindow"
-
 func raise(hwnd win.HWND) {
 	if win.IsIconic(hwnd) {
 		win.ShowWindow(hwnd, win.SW_RESTORE)
@@ -30,10 +28,11 @@ func raise(hwnd win.HWND) {
 }
 
 func RaiseUI() bool {
-	hwnd := win.FindWindow(windows.StringToUTF16Ptr(wireguardUIClass), nil)
+	hwnd := win.FindWindow(windows.StringToUTF16Ptr(manageWindowWindowClass), nil)
 	if hwnd == 0 {
 		return false
 	}
+	win.SendMessage(hwnd, selectCorrectTabMsg, 0, 0)
 	raise(hwnd)
 	return true
 }
@@ -42,12 +41,13 @@ func WaitForRaiseUIThenQuit() {
 	var handle win.HWINEVENTHOOK
 	runtime.LockOSThread()
 	handle, err := win.SetWinEventHook(win.EVENT_OBJECT_CREATE, win.EVENT_OBJECT_CREATE, 0, func(hWinEventHook win.HWINEVENTHOOK, event uint32, hwnd win.HWND, idObject int32, idChild int32, idEventThread uint32, dwmsEventTime uint32) uintptr {
-		class := make([]uint16, len(wireguardUIClass)+2) /* Plus 2, one for the null terminator, and one to see if this is only a prefix */
+		class := make([]uint16, len(manageWindowWindowClass)+2) /* Plus 2, one for the null terminator, and one to see if this is only a prefix */
 		n, err := win.GetClassName(hwnd, &class[0], len(class))
-		if err != nil || n != len(wireguardUIClass) || windows.UTF16ToString(class) != wireguardUIClass {
+		if err != nil || n != len(manageWindowWindowClass) || windows.UTF16ToString(class) != manageWindowWindowClass {
 			return 0
 		}
 		win.UnhookWinEvent(handle)
+		win.SendMessage(hwnd, selectCorrectTabMsg, 0, 0)
 		raise(hwnd)
 		os.Exit(0)
 		return 0
diff --git a/ui/tray.go b/ui/tray.go
index 668163da..f2c387e4 100644
--- a/ui/tray.go
+++ b/ui/tray.go
@@ -41,7 +41,7 @@ func NewTray(mtw *ManageTunnelsWindow) (*Tray, error) {
 		tunnels: make(map[string]*walk.Action),
 	}
 
-	tray.NotifyIcon, err = walk.NewNotifyIcon(mtw.MainWindow)
+	tray.NotifyIcon, err = walk.NewNotifyIcon(mtw)
 	if err != nil {
 		return nil, err
 	}
-- 
cgit v1.2.3-59-g8ed1b