From fd22f9bf8e941260cd8b3c560c3afa82de925229 Mon Sep 17 00:00:00 2001
From: Nuno Cruces <ncruces@gmail.com>
Date: Wed, 8 Jan 2020 00:44:26 +0000
Subject: [PATCH] Message dialogs improvements (windows).

---
 file_windows.go |  8 ------
 init_windows.go | 34 ++++++++++++++++++++++
 msg_windows.go  | 75 +++++++++++++++++++++++++++++++++++++++++++++++--
 3 files changed, 106 insertions(+), 11 deletions(-)
 create mode 100644 init_windows.go

diff --git a/file_windows.go b/file_windows.go
index 1c75dcd..98113b2 100644
--- a/file_windows.go
+++ b/file_windows.go
@@ -11,17 +11,9 @@ import (
 )
 
 var (
-	comdlg32 = syscall.NewLazyDLL("comdlg32.dll")
-	ole32    = syscall.NewLazyDLL("ole32.dll")
-	shell32  = syscall.NewLazyDLL("shell32.dll")
-
 	getOpenFileName             = comdlg32.NewProc("GetOpenFileNameW")
 	getSaveFileName             = comdlg32.NewProc("GetSaveFileNameW")
 	commDlgExtendedError        = comdlg32.NewProc("CommDlgExtendedError")
-	coInitializeEx              = ole32.NewProc("CoInitializeEx")
-	coUninitialize              = ole32.NewProc("CoUninitialize")
-	coCreateInstance            = ole32.NewProc("CoCreateInstance")
-	coTaskMemFree               = ole32.NewProc("CoTaskMemFree")
 	shBrowseForFolder           = shell32.NewProc("SHBrowseForFolderW")
 	shGetPathFromIDListEx       = shell32.NewProc("SHGetPathFromIDListEx")
 	shCreateItemFromParsingName = shell32.NewProc("SHCreateItemFromParsingName")
diff --git a/init_windows.go b/init_windows.go
new file mode 100644
index 0000000..7089f48
--- /dev/null
+++ b/init_windows.go
@@ -0,0 +1,34 @@
+package zenity
+
+import "syscall"
+
+var (
+	comdlg32 = syscall.NewLazyDLL("comdlg32.dll")
+	kernel32 = syscall.NewLazyDLL("kernel32.dll")
+	ole32    = syscall.NewLazyDLL("ole32.dll")
+	shell32  = syscall.NewLazyDLL("shell32.dll")
+	user32   = syscall.NewLazyDLL("user32.dll")
+
+	getCurrentThreadId = kernel32.NewProc("GetCurrentThreadId")
+
+	coInitializeEx   = ole32.NewProc("CoInitializeEx")
+	coUninitialize   = ole32.NewProc("CoUninitialize")
+	coCreateInstance = ole32.NewProc("CoCreateInstance")
+	coTaskMemFree    = ole32.NewProc("CoTaskMemFree")
+
+	getClassName        = user32.NewProc("GetClassNameA")
+	setWindowsHookEx    = user32.NewProc("SetWindowsHookExW")
+	unhookWindowsHookEx = user32.NewProc("UnhookWindowsHookEx")
+	callNextHookEx      = user32.NewProc("CallNextHookEx")
+	enumChildWindows    = user32.NewProc("EnumChildWindows")
+	getDlgCtrlID        = user32.NewProc("GetDlgCtrlID")
+	setWindowText       = user32.NewProc("SetWindowTextW")
+)
+
+type _CWPRETSTRUCT struct {
+	Result  uintptr
+	LParam  uintptr
+	WParam  uintptr
+	Message uint32
+	HWnd    uintptr
+}
diff --git a/msg_windows.go b/msg_windows.go
index eba781b..3a93753 100644
--- a/msg_windows.go
+++ b/msg_windows.go
@@ -1,12 +1,12 @@
 package zenity
 
 import (
+	"runtime"
 	"syscall"
 	"unsafe"
 )
 
 var (
-	user32     = syscall.NewLazyDLL("user32.dll")
 	messageBox = user32.NewProc("MessageBoxW")
 )
 
@@ -49,17 +49,86 @@ func message(typ int, text string, options []Option) (bool, error) {
 		flags |= 0x40 // MB_ICONINFORMATION
 	}
 
+	if typ == 2 && opts.defcancel {
+		if opts.extra == "" {
+			flags |= 0x100 // MB_DEFBUTTON2
+		} else {
+			flags |= 0x200 // MB_DEFBUTTON3
+		}
+	}
+
 	if opts.title != "" {
 		caption = uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(opts.title)))
 	}
 
+	if opts.ok != "" || opts.cancel != "" || opts.extra != "" {
+		runtime.LockOSThread()
+		defer runtime.UnlockOSThread()
+
+		hook, err := hookMessageLabels(typ, opts)
+		if hook == 0 {
+			return false, err
+		}
+		defer unhookWindowsHookEx.Call(hook)
+	}
+
 	n, _, err := messageBox.Call(0,
 		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(text))),
 		caption, flags)
 
 	if n == 0 {
 		return false, err
-	} else {
-		return n == 1 /* IDOK */ || n == 6 /* IDYES */, nil
 	}
+	if n == 7 || n == 2 && typ != 2 { // IDNO
+		return false, ErrExtraButton
+	}
+	if n == 1 || n == 6 { // IDOK, IDYES
+		return true, nil
+	}
+	return false, nil
+}
+
+func hookMessageLabels(typ int, opts options) (hook uintptr, err error) {
+	tid, _, _ := getCurrentThreadId.Call()
+	hook, _, err = setWindowsHookEx.Call(12, // WH_CALLWNDPROCRET
+		syscall.NewCallback(func(code int, wparam, lparam uintptr) uintptr {
+			msg := *(*_CWPRETSTRUCT)(unsafe.Pointer(lparam))
+			if msg.Message == 0x0110 { // WM_INITDIALOG
+				name := [7]byte{}
+				n, _, _ := getClassName.Call(msg.HWnd, uintptr(unsafe.Pointer(&name)), uintptr(len(name)))
+				if string(name[:n]) == "#32770" {
+					enumChildWindows.Call(msg.HWnd,
+						syscall.NewCallback(func(hwnd, lparam uintptr) uintptr {
+							name := [7]byte{}
+							n, _, _ := getClassName.Call(hwnd, uintptr(unsafe.Pointer(&name)), uintptr(len(name)))
+							if string(name[:n]) == "Button" {
+								ctl, _, _ := getDlgCtrlID.Call(hwnd)
+								var text string
+								switch ctl {
+								case 1, 6: // IDOK, IDYES
+									text = opts.ok
+								case 2: // IDCANCEL
+									if typ == 2 {
+										text = opts.cancel
+									} else if opts.extra != "" {
+										text = opts.extra
+									} else {
+										text = opts.ok
+									}
+								case 7: // IDNO
+									text = opts.extra
+								}
+								if text != "" {
+									ptr := syscall.StringToUTF16Ptr(text)
+									setWindowText.Call(hwnd, uintptr(unsafe.Pointer(ptr)))
+								}
+							}
+							return 1
+						}), 0)
+				}
+			}
+			next, _, _ := callNextHookEx.Call(hook, uintptr(code), wparam, lparam)
+			return next
+		}), 0, tid)
+	return
 }