add systray_on_exit callback

3 files changed, 48 insertions, 22 deletions

diff --git a/systray.go b/systray.go
index 7abc072..38cf242 100644
--- a/systray.go
+++ b/systray.go
@@ -1,7 +1,9 @@
 /*
-Package systray is a cross platfrom Go library to place an icon and menu in the notification area.
+Package systray is a cross platfrom 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 at the very beginning of main() to lock at main thread.
+Methods can be called from any goroutine except Run(), which should be called
+at the very beginning of main() to lock at main thread.
 */
 package systray
 
@@ -17,7 +19,7 @@ import (
 // Don't create it directly, use the one systray.AddMenuItem() returned
 type MenuItem struct {
 	// ClickedCh is the channel which will be notified when the menu item is clicked
-	ClickedCh chan interface{}
+	ClickedCh chan struct{}
 
 	// id uniquely identify a menu item, not supposed to be modified
 	id int32
@@ -34,8 +36,8 @@ type MenuItem struct {
 var (
 	log = golog.LoggerFor("systray")
 
-	readyCh       = make(chan interface{})
-	clickedCh     = make(chan interface{})
+	readyCh       = make(chan struct{})
+	exitCh        = make(chan struct{})
 	menuItems     = make(map[int32]*MenuItem)
 	menuItemsLock sync.RWMutex
 
@@ -46,11 +48,18 @@ var (
 // 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()) {
+func Run(onReady func(), onExit func()) {
 	runtime.LockOSThread()
 	go func() {
-		<-readyCh
-		onReady()
+		for {
+			select {
+			case <-readyCh:
+				onReady()
+			case <-exitCh:
+				onExit()
+				return
+			}
+		}
 	}()
 
 	nativeLoop()
@@ -68,7 +77,7 @@ func Quit() {
 func AddMenuItem(title string, tooltip string) *MenuItem {
 	id := atomic.AddInt32(&currentID, 1)
 	item := &MenuItem{nil, id, title, tooltip, false, false}
-	item.ClickedCh = make(chan interface{})
+	item.ClickedCh = make(chan struct{})
 	item.update()
 	return item
 }
@@ -128,7 +137,11 @@ func (item *MenuItem) update() {
 }
 
 func systrayReady() {
-	readyCh <- nil
+	readyCh <- struct{}{}
+}
+
+func systrayExit() {
+	exitCh <- struct{}{}
 }
 
 func systrayMenuItemSelected(id int32) {
@@ -136,7 +149,7 @@ func systrayMenuItemSelected(id int32) {
 	item := menuItems[id]
 	menuItemsLock.RUnlock()
 	select {
-	case item.ClickedCh <- nil:
+	case item.ClickedCh <- struct{}{}:
 	// in case no one waiting for the channel
 	default:
 	}
diff --git a/systray/systray/systray.cpp b/systray/systray/systray.cpp
index 8de2daa..ae6961f 100755
--- a/systray/systray/systray.cpp
+++ b/systray/systray/systray.cpp
@@ -12,6 +12,7 @@ static NOTIFYICONDATA nid;
 static HWND hWnd;

 static HMENU hTrayMenu;

 

+void (*systray_on_exit)(int ignored);

 void (*systray_menu_item_selected)(int menu_id);

 

 void reportWindowsError(const char* action) {

@@ -26,7 +27,7 @@ void reportWindowsError(const char* action) {
 			pErrMsg,

 			0,

 			NULL);

-	printf("Systray error %s: %d %s\n", action, errCode, pErrMsg);

+	printf("Systray error %s: %d %ls\n", action, errCode, pErrMsg);

 }

 

 void ShowMenu(HWND hWnd) {

@@ -62,8 +63,12 @@ LRESULT CALLBACK WndProc(HWND hWnd, UINT message, WPARAM wParam, LPARAM lParam)
 			}

 			break;

 		case WM_DESTROY:

+			systray_on_exit(0/*ignored*/);

 			PostQuitMessage(0);

 			break;

+		case WM_ENDSESSION:

+			systray_on_exit(0/*ignored*/);

+			break;

 		case WM_SYSTRAY_MESSAGE:

 			switch(lParam) {

 				case WM_RBUTTONUP:

@@ -133,7 +138,10 @@ BOOL addNotifyIcon() {
 	return Shell_NotifyIcon(NIM_ADD, &nid);

 }

 

-int nativeLoop(void (*systray_ready)(int ignored), void (*_systray_menu_item_selected)(int menu_id)) {

+int nativeLoop(void (*systray_ready)(int ignored),

+	void (*_systray_on_exit)(int ignored),

+    void (*_systray_menu_item_selected)(int menu_id)) {

+	systray_on_exit = _systray_on_exit;

 	systray_menu_item_selected = _systray_menu_item_selected;

 

 	HINSTANCE hInstance = GetModuleHandle(NULL);

@@ -152,7 +160,7 @@ int nativeLoop(void (*systray_ready)(int ignored), void (*_systray_menu_item_sel
 	while (GetMessage(&msg, NULL, 0, 0)) {

 		TranslateMessage(&msg);

 		DispatchMessage(&msg);

-	}   

+	}

 	return EXIT_SUCCESS;

 }

 

diff --git a/systray_windows.go b/systray_windows.go
index 7630c82..664af0c 100644
--- a/systray_windows.go
+++ b/systray_windows.go
@@ -5,19 +5,18 @@ import (
 	"io/ioutil"
 	"os"
 	"path/filepath"
-	"syscall"
 	"runtime"
+	"syscall"
 	"unsafe"
 
 	"github.com/getlantern/filepersist"
 )
 
 var (
-	iconFiles = make([]*os.File, 0)
-	dllDir    = filepath.Join(os.Getenv("APPDATA"), "systray")
+	iconFiles   = make([]*os.File, 0)
+	dllDir      = filepath.Join(os.Getenv("APPDATA"), "systray")
 	dllFileName = "systray" + runtime.GOARCH + ".dll"
-	dllFile   = filepath.Join(dllDir, dllFileName)
-
+	dllFile     = filepath.Join(dllDir, dllFileName)
 
 	mod                      = syscall.NewLazyDLL(dllFile)
 	_nativeLoop              = mod.NewProc("nativeLoop")
@@ -32,23 +31,24 @@ func init() {
 	// Write DLL to file
 	b, err := Asset(dllFileName)
 	if err != nil {
-		panic(fmt.Errorf("Unable to read " + dllFileName + ": %v", err))
+		panic(fmt.Errorf("Unable to read "+dllFileName+": %v", err))
 	}
 
 	err = os.MkdirAll(dllDir, 0755)
 	if err != nil {
-		panic(fmt.Errorf("Unable to create directory %v to hold " + dllFileName + ": %v", dllDir, err))
+		panic(fmt.Errorf("Unable to create directory %v to hold "+dllFileName+": %v", dllDir, err))
 	}
 
 	err = filepersist.Save(dllFile, b, 0644)
 	if err != nil {
-		panic(fmt.Errorf("Unable to save " + dllFileName + " to %v: %v", dllFile, err))
+		panic(fmt.Errorf("Unable to save "+dllFileName+" to %v: %v", dllFile, err))
 	}
 }
 
 func nativeLoop() {
 	_nativeLoop.Call(
 		syscall.NewCallbackCDecl(systray_ready),
+		syscall.NewCallbackCDecl(systray_on_exit),
 		syscall.NewCallbackCDecl(systray_menu_item_selected))
 }
 
@@ -151,6 +151,11 @@ func systray_ready(ignore uintptr) uintptr {
 	return 0
 }
 
+func systray_on_exit(ignore uintptr) uintptr {
+	systrayExit()
+	return 0
+}
+
 func systray_menu_item_selected(id uintptr) uintptr {
 	systrayMenuItemSelected(int32(id))
 	return 0