Duplicate Process Handling (#13)

* Added duplicate Pageant window checking

* Added custom `LAUNCH_APP` msi switch for winget

* Added fix so users can still get error messages if running from console

go.mod

@@ -4,5 +4,5 @@ go 1.17
 
 require (
 	github.com/Microsoft/go-winio v0.5.2
-	golang.org/x/sys v0.0.0-20220317061510-51cd9980dadf
+	golang.org/x/sys v0.0.0-20220319134239-a9b59b0215f8
 )

go.sum

@@ -6,5 +6,5 @@ github.com/sirupsen/logrus v1.7.0/go.mod h1:yWOB1SBYBC5VeMP7gHvWumXLIWorT60ONWic
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
 golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210124154548-22da62e12c0c/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20220317061510-51cd9980dadf h1:Fm4IcnUL803i92qDlmB0obyHmosDrxZWxJL3gIeNqOw=
-golang.org/x/sys v0.0.0-20220317061510-51cd9980dadf/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20220319134239-a9b59b0215f8 h1:OH54vjqzRWmbJ62fjuhxy7AxFFgoHN0/DPc/UrL8cAs=
+golang.org/x/sys v0.0.0-20220319134239-a9b59b0215f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=

internal/win/consoleoutput.go

@@ -0,0 +1,68 @@
+package win
+
+import (
+	"fmt"
+	"log"
+	"os"
+	"syscall"
+
+	"golang.org/x/sys/windows"
+)
+
+func AttachConsole() error {
+	r1, _, err := attachConsole.Call(ATTACH_PARENT_PROCESS)
+	if r1 == 0 {
+		errno, ok := err.(syscall.Errno)
+		if ok && errno == ERROR_INVALID_HANDLE {
+			// console handle doesn't exist; not a real error, but the console handle will be invalid.
+			return nil
+		}
+		return err
+	}
+	return nil
+}
+
+var oldStdout *os.File
+
+func FixConsoleIfNeeded() error {
+	// Keep old os.Stdout reference so it dont get GC'd and cleaned up
+	// You never want to close file descriptors 0, 1, and 2.
+	oldStdout = os.Stdout
+	stdout, _ := syscall.GetStdHandle(syscall.STD_OUTPUT_HANDLE)
+
+	var invalid syscall.Handle
+	con := invalid
+
+	if stdout == invalid {
+		err := AttachConsole()
+		if err != nil {
+			return fmt.Errorf("attachconsole: %v", err)
+		}
+		if stdout == invalid {
+			stdout, _ = syscall.GetStdHandle(syscall.STD_OUTPUT_HANDLE)
+			con = stdout
+		}
+	}
+
+	if con != invalid {
+		// Make sure the console is configured to convert
+		// \n to \r\n, like Go programs expect.
+		h := windows.Handle(con)
+		var st uint32
+		err := windows.GetConsoleMode(h, &st)
+		if err != nil {
+			return fmt.Errorf("GetConsoleMode: %v", err)
+		}
+		err = windows.SetConsoleMode(h, st&^windows.DISABLE_NEWLINE_AUTO_RETURN)
+		if err != nil {
+			return fmt.Errorf("SetConsoleMode: %v", err)
+		}
+	}
+
+	if stdout != invalid {
+		os.Stdout = os.NewFile(uintptr(stdout), "stdout")
+	}
+
+	log.SetOutput(os.Stdout)
+	return nil
+}

internal/win/defs.go

@@ -2,16 +2,18 @@ package win
 
 const (
 	//revive:disable:var-naming,exported
-	IDI_APPLICATION    = 32512
-	IDC_IBEAM          = 32513
-	BLACK_BRUSH        = 4
-	WM_COPYDATA        = 74
-	WM_DESTROY         = 2
-	WM_CLOSE           = 16
-	WM_QUERYENDSESSION = 17
-	WM_QUIT            = 18
-	WS_EX_TOOLWINDOW   = 0x00000080
-	WS_EX_APPWINDOW    = 0x00040000
+	IDI_APPLICATION       = 32512
+	IDC_IBEAM             = 32513
+	BLACK_BRUSH           = 4
+	WM_COPYDATA           = 74
+	WM_DESTROY            = 2
+	WM_CLOSE              = 16
+	WM_QUERYENDSESSION    = 17
+	WM_QUIT               = 18
+	WS_EX_TOOLWINDOW      = 0x00000080
+	WS_EX_APPWINDOW       = 0x00040000
+	ATTACH_PARENT_PROCESS = uintptr(^uint32(0)) // (DWORD)-1
+	ERROR_INVALID_HANDLE  = 6
 )
 
 type (

internal/win/kernel32.go

@@ -15,6 +15,7 @@ var (
 	getModuleHandle = libkernel32.NewProc("GetModuleHandleW")
 	globalAlloc     = libkernel32.NewProc("GlobalAlloc")
 	globalFree      = libkernel32.NewProc("GlobalFree")
+	attachConsole   = libkernel32.NewProc("AttachConsole")
 )
 
 func GetLastError() uint32 {

internal/win/user32.go

@@ -21,6 +21,7 @@ var (
 	translateMessage = libuser32.NewProc("TranslateMessage")
 	dispatchMessage  = libuser32.NewProc("DispatchMessageW")
 	postQuitMessage  = libuser32.NewProc("PostQuitMessage")
+	findWindow       = libuser32.NewProc("FindWindowW")
 )
 
 func MAKEINTRESOURCE(id uintptr) *uint16 {
@@ -130,3 +131,12 @@ func PostQuitMessage(exitCode int32) {
 		0,
 		0)
 }
+
+func FindWindow(lpClassName, lpWindowName *uint16) HWND {
+	ret, _, _ := syscall.Syscall(findWindow.Addr(), 2,
+		uintptr(unsafe.Pointer(lpClassName)),
+		uintptr(unsafe.Pointer(lpWindowName)),
+		0)
+
+	return HWND(ret)
+}

main.go

@@ -4,6 +4,7 @@ import (
 	"flag"
 	"fmt"
 	"log"
+	"os"
 	"runtime"
 	"unsafe"
 
@@ -15,15 +16,29 @@ var (
 	noPageantPipe = flag.Bool("no-pageant-pipe", false, "Toggle pageant named pipe proxying")
 )
 
+var oldStdin, oldStdout, oldStderr *os.File
+
 func main() {
 	flag.Parse()
 
+	err := win.FixConsoleIfNeeded()
+	if err != nil {
+		log.Fatalf("FixConsoleOutput: %v\n", err)
+	}
+
+	// Check if any application claiming to be a Pageant Window is already running
+	if doesPagentWindowExist() {
+		log.Println("This application is already running, exiting.")
+		return
+	}
+
 	// Start a proxy/redirector for the pageant named pipes
 	if !*noPageantPipe {
 		go pipeProxy()
 	}
 
 	runtime.LockOSThread()
+	defer runtime.UnlockOSThread()
 
 	pageantWindow := createPageantWindow()
 	if pageantWindow == 0 {
@@ -42,5 +57,4 @@ func main() {
 
 	// Explicitly release the global memory handle
 	win.GlobalFree(hglobal)
-	runtime.UnlockOSThread()
 }

pageant.go

@@ -13,12 +13,11 @@ import (
 	"unsafe"
 
 	"encoding/binary"
+	"encoding/hex"
 
 	"github.com/Microsoft/go-winio"
 	"golang.org/x/sys/windows"
 
-	"encoding/hex"
-
 	"github.com/ndbeals/winssh-pageant/internal/security"
 	"github.com/ndbeals/winssh-pageant/internal/sshagent"
 	"github.com/ndbeals/winssh-pageant/internal/win"
@@ -64,6 +63,10 @@ func openFileMap(dwDesiredAccess, bInheritHandle uint32, mapNamePtr uintptr) (wi
 	return windows.Handle(mapPtr), err
 }
 
+func doesPagentWindowExist() bool {
+	return win.FindWindow(wndClassNamePtr, nil) != 0
+}
+
 func registerPageantWindow(hInstance win.HINSTANCE) (atom win.ATOM) {
 	var wc win.WNDCLASSEX
 	wc.Style = 0
@@ -215,16 +218,17 @@ func pipeProxy() {
 	listener, err := winio.ListenPipe(pipeName, nil)
 	if err != nil {
 		log.Println(err)
-	}
-	defer listener.Close()
+	} else {
+		defer listener.Close()
 
-	for {
-		conn, err := listener.Accept()
-		if err != nil {
-			log.Println(err)
-			return
+		for {
+			conn, err := listener.Accept()
+			if err != nil {
+				log.Println(err)
+				return
+			}
+			go pipeListen(conn)
 		}
-		go pipeListen(conn)
 	}
 }
 

resources/templates/product.wxs

@@ -15,6 +15,7 @@
         <Package InstallerVersion="500" Compressed="yes" Comments="WinSSH-Pageant Installer Package" InstallScope="perUser" />
         <Property Id="ALLUSERS" Secure="yes" Value="2"/>
         <Property Id="MSIINSTALLPERUSER" Secure="yes" Value="1" />
+        <Property Id="LAUNCH_APP" Hidden="no" Secure="yes">1</Property>
 
         <Media Id="1" Cabinet="product.cab" EmbedCab="yes" />
 
@@ -77,7 +78,7 @@
         <Property Id="WixShellExecTarget" Value="[#ApplicationFile0]" />
         <CustomAction Id="LaunchApplication" BinaryKey="WixCA" DllEntry="WixShellExec" Impersonate="yes" />
 
-        <Property Id="WixQuietExecCmdLine" Value='"taskkill.exe" /F /IM winssh-pageant.exe'/>
+        <Property Id="WixQuietExecCmdLine" Value='"taskkill.exe" /F /IM {{ slice (index .Files.Items 0) 3}}'/>
         <CustomAction Id="AppKill" BinaryKey="WixCA" DllEntry="WixQuietExec" Execute="immediate" Return="ignore" Impersonate="yes" />
 
         {{range $i, $e := .InstallHooks}}
@@ -92,12 +93,12 @@
         <InstallExecuteSequence>
             <RemoveExistingProducts After="InstallValidate" />
             <Custom Action="AppKill" Before="InstallValidate">REMOVE OR REINSTALL</Custom>
-            <Custom Action="LaunchApplication" After="InstallFinalize">(NOT Installed OR REINSTALL) AND NOT REMOVE</Custom>
+            <Custom Action="LaunchApplication" After="InstallFinalize">((NOT Installed OR REINSTALL) AND NOT REMOVE) AND (<![CDATA[LAUNCH_APP="true" OR LAUNCH_APP=1]]>)</Custom>
 
             {{range $i, $e := .InstallHooks}}
             <Custom Action="CustomInstallExec{{$i}}" After="{{if eq $i 0}}InstallFiles{{else}}CustomInstallExec{{dec $i}}{{end}}">NOT Installed AND NOT REMOVE</Custom>
             {{end}}
-         {{range $i, $e := .UninstallHooks}}
+            {{range $i, $e := .UninstallHooks}}
             <Custom Action="CustomUninstallExec{{$i}}" After="{{if eq $i 0}}InstallInitialize{{else}}CustomUninstallExec{{dec $i}}{{end}}">REMOVE ~= "ALL"</Custom>
             {{end}}
         </InstallExecuteSequence>