From 01c2abbb00f448beb57fa1b3aca236f27e92f8b3 Mon Sep 17 00:00:00 2001
From: Nuno Cruces <ncruces@users.noreply.github.com>
Date: Sat, 14 Aug 2021 02:50:04 +0100
Subject: [PATCH] Workaround #17.

---
 cmd/zenity/main.go | 34 ++++++++++++++--------------------
 notify_windows.go  | 27 ++++++++++++++++++++++++---
 util_windows.go    |  1 +
 3 files changed, 39 insertions(+), 23 deletions(-)

diff --git a/cmd/zenity/main.go b/cmd/zenity/main.go
index ad3ef14..0e87502 100644
--- a/cmd/zenity/main.go
+++ b/cmd/zenity/main.go
@@ -11,9 +11,11 @@ import (
 	"image/color"
 	"os"
 	"os/exec"
+	"os/signal"
 	"path/filepath"
 	"runtime"
 	"strings"
+	"syscall"
 	"time"
 
 	"github.com/ncruces/zenity"
@@ -113,11 +115,15 @@ func main() {
 	if unixeol {
 		zenutil.LineBreak = "\n"
 	}
+	ctx, cancel := signal.NotifyContext(context.Background(),
+		syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM)
+	defer cancel()
 	if zenutil.Timeout > 0 {
-		ctx, cancel := context.WithTimeout(context.Background(), time.Duration(zenutil.Timeout)*time.Second)
-		opts = append(opts, zenity.Context(ctx))
+		c, cancel := context.WithTimeout(ctx, time.Duration(zenutil.Timeout)*time.Second)
 		defer cancel()
+		ctx = c
 	}
+	opts = append(opts, zenity.Context(ctx))
 
 	switch {
 	case errorDlg:
@@ -458,7 +464,7 @@ func errResult(err error) {
 	if os.IsTimeout(err) {
 		os.Exit(5)
 	}
-	if err == zenity.ErrCanceled {
+	if err == zenity.ErrCanceled || err == context.Canceled {
 		os.Exit(1)
 	}
 	if err == zenity.ErrExtraButton {
@@ -471,36 +477,24 @@ func errResult(err error) {
 		os.Stderr.WriteString(zenutil.LineBreak)
 		os.Exit(-1)
 	}
-	os.Exit(0)
 }
 
 func strResult(s string, err error) {
-	if err != nil {
-		errResult(err)
-	}
+	errResult(err)
 	os.Stdout.WriteString(s)
 	os.Stdout.WriteString(zenutil.LineBreak)
-	os.Exit(0)
 }
 
 func lstResult(l []string, err error) {
-	if err != nil {
-		errResult(err)
-	}
-	if len(l) > 0 {
-		os.Stdout.WriteString(strings.Join(l, zenutil.Separator))
-		os.Stdout.WriteString(zenutil.LineBreak)
-	}
-	os.Exit(0)
+	errResult(err)
+	os.Stdout.WriteString(strings.Join(l, zenutil.Separator))
+	os.Stdout.WriteString(zenutil.LineBreak)
 }
 
 func colResult(c color.Color, err error) {
-	if err != nil {
-		errResult(err)
-	}
+	errResult(err)
 	os.Stdout.WriteString(zenutil.UnparseColor(c))
 	os.Stdout.WriteString(zenutil.LineBreak)
-	os.Exit(0)
 }
 
 func ingestPath(path string) string {
diff --git a/notify_windows.go b/notify_windows.go
index 0c545bd..32c433a 100644
--- a/notify_windows.go
+++ b/notify_windows.go
@@ -1,16 +1,19 @@
 package zenity
 
 import (
+	"math/rand"
 	"runtime"
 	"syscall"
+	"time"
 	"unsafe"
 
 	"github.com/ncruces/zenity/internal/zenutil"
 )
 
 var (
-	shellNotifyIcon = shell32.NewProc("Shell_NotifyIconW")
-	wtsSendMessage  = wtsapi32.NewProc("WTSSendMessageW")
+	rtlGetNtVersionNumbers = ntdll.NewProc("RtlGetNtVersionNumbers")
+	shellNotifyIcon        = shell32.NewProc("Shell_NotifyIconW")
+	wtsSendMessage         = wtsapi32.NewProc("WTSSendMessageW")
 )
 
 func notify(text string, opts options) error {
@@ -20,7 +23,7 @@ func notify(text string, opts options) error {
 
 	var args _NOTIFYICONDATA
 	args.StructSize = uint32(unsafe.Sizeof(args))
-	args.ID = 0x378eb49c    // random
+	args.ID = rand.Uint32()
 	args.Flags = 0x00000010 // NIF_INFO
 	args.State = 0x00000001 // NIS_HIDDEN
 
@@ -52,6 +55,24 @@ func notify(text string, opts options) error {
 		return err
 	}
 
+	var major, minor, build uint32
+	rtlGetNtVersionNumbers.Call(
+		uintptr(unsafe.Pointer(&major)),
+		uintptr(unsafe.Pointer(&minor)),
+		uintptr(unsafe.Pointer(&build)))
+
+	// On Windows 7 (6.1) and lower, wait up to 10 seconds to clean up.
+	if major < 6 || major == 6 && minor < 2 {
+		if opts.ctx != nil {
+			select {
+			case <-opts.ctx.Done():
+			case <-time.After(10 * time.Second):
+			}
+		} else {
+			time.Sleep(10 * time.Second)
+		}
+	}
+
 	shellNotifyIcon.Call(2 /* NIM_DELETE */, uintptr(unsafe.Pointer(&args)))
 	return nil
 }
diff --git a/util_windows.go b/util_windows.go
index e7789c1..22cfd90 100644
--- a/util_windows.go
+++ b/util_windows.go
@@ -17,6 +17,7 @@ var (
 	comdlg32 = syscall.NewLazyDLL("comdlg32.dll")
 	gdi32    = syscall.NewLazyDLL("gdi32.dll")
 	kernel32 = syscall.NewLazyDLL("kernel32.dll")
+	ntdll    = syscall.NewLazyDLL("ntdll.dll")
 	ole32    = syscall.NewLazyDLL("ole32.dll")
 	shell32  = syscall.NewLazyDLL("shell32.dll")
 	user32   = syscall.NewLazyDLL("user32.dll")