From 37dd2808dd1173c261d1caee3bec70fe5ac31703 Mon Sep 17 00:00:00 2001
From: buptczq <czq@qq.com>
Date: Fri, 20 Nov 2020 16:22:07 +0800
Subject: [PATCH] fix: hide WSL2 menu when Hyper-V disabled

---
 app/vsock.go     | 40 ++++++++++++++++++++++++++--------------
 main.go          |  2 ++
 utils/wsl2.go    | 13 +++++++++++++
 versioninfo.json |  6 +++---
 4 files changed, 44 insertions(+), 17 deletions(-)

diff --git a/app/vsock.go b/app/vsock.go
index 9980ff9..652f44c 100644
--- a/app/vsock.go
+++ b/app/vsock.go
@@ -124,6 +124,15 @@ func (s *VSock) Run(ctx context.Context, handler func(conn io.ReadWriteCloser))
 	if isHV {
 		return nil
 	}
+
+	if !utils.CheckHvSocket() {
+		return nil
+	}
+
+	if !utils.CheckHVService() {
+		return nil
+	}
+
 	agentSrvGUID := winio.VsockServiceID(utils.ServicePort)
 	pipe, err := winio.ListenHvsock(&winio.HvsockAddr{
 		VMID:      vmWildCard,
@@ -163,6 +172,9 @@ func (*VSock) AppId() AppId {
 }
 
 func (s *VSock) Menu(register func(id AppId, name string, handler func())) {
+	if !utils.CheckHvSocket() {
+		return
+	}
 	wsl2 := AppId(APP_WSL2)
 	register(wsl2, "Show WSL2 / Linux On Hyper-V Settings", s.onClick)
 	register(s.AppId(), "Check Hyper-V Agent Status", s.onCheckClick)
@@ -170,36 +182,36 @@ func (s *VSock) Menu(register func(id AppId, name string, handler func())) {
 
 func (s *VSock) onClick() {
 	if s.running {
+		help := "socat UNIX-LISTEN:/tmp/wincrypt-hv.sock,fork,mode=777 SOCKET-CONNECT:40:0:x0000x33332222x02000000x00000000,forever,interval=5 &\n"
+		help += "export SSH_AUTH_SOCK=/tmp/wincrypt-hv.sock\n"
+		if utils.MessageBox(s.AppId().FullName()+" (OK to copy):", help, utils.MB_OKCANCEL) == utils.IDOK {
+			utils.SetClipBoard(help)
+		}
+	} else {
 		if !utils.CheckHVService() {
 			if utils.MessageBox(s.AppId().FullName()+":", s.AppId().String()+" agent is not working! Don you want to enable it?", utils.MB_OKCANCEL) == utils.IDOK {
 				if err := utils.RunMeElevatedWithArgs("-i"); err != nil {
 					utils.MessageBox("Install Service Error:", err.Error(), utils.MB_ICONERROR)
-					return
 				}
-			} else {
-				return
 			}
+		} else {
+			utils.MessageBox("Error:", s.AppId().String()+" agent doesn't work!", utils.MB_ICONWARNING)
 		}
-		help := "socat UNIX-LISTEN:/tmp/wincrypt-hv.sock,fork,mode=777 SOCKET-CONNECT:40:0:x0000x33332222x02000000x00000000,forever,interval=5 &\n"
-		help += "export SSH_AUTH_SOCK=/tmp/wincrypt-hv.sock\n"
-		if utils.MessageBox(s.AppId().FullName()+" (OK to copy):", help, utils.MB_OKCANCEL) == utils.IDOK {
-			utils.SetClipBoard(help)
-		}
-	} else {
-		utils.MessageBox("Error:", s.AppId().String()+" agent doesn't work!", utils.MB_ICONWARNING)
 	}
 }
 
 func (s *VSock) onCheckClick() {
 	if s.running {
+		utils.MessageBox(s.AppId().FullName()+":", s.AppId().String()+" agent is working!", 0)
+	} else {
 		if !utils.CheckHVService() {
 			if utils.MessageBox(s.AppId().FullName()+":", s.AppId().String()+" agent is not working! Don you want to enable it?", utils.MB_OKCANCEL) == utils.IDOK {
-				utils.RunMeElevatedWithArgs("-i")
+				if err := utils.RunMeElevatedWithArgs("-i"); err != nil {
+					utils.MessageBox("Install Service Error:", err.Error(), utils.MB_ICONERROR)
+				}
 			}
 		} else {
-			utils.MessageBox(s.AppId().FullName()+":", s.AppId().String()+" agent is working!", 0)
+			utils.MessageBox("Error:", s.AppId().String()+" agent doesn't work!", utils.MB_ICONWARNING)
 		}
-	} else {
-		utils.MessageBox("Error:", s.AppId().String()+" agent doesn't work!", utils.MB_ICONWARNING)
 	}
 }
diff --git a/main.go b/main.go
index 25f5c41..063d430 100644
--- a/main.go
+++ b/main.go
@@ -58,6 +58,8 @@ func installService() {
 	})
 	if err != nil {
 		utils.MessageBox("Install Service Error:", err.Error(), utils.MB_ICONERROR)
+	} else {
+		utils.MessageBox("Install Service Success:", "Please reboot your computer to take effect!", utils.MB_ICONINFORMATION)
 	}
 	return
 
diff --git a/utils/wsl2.go b/utils/wsl2.go
index a990b6e..dee98dd 100644
--- a/utils/wsl2.go
+++ b/utils/wsl2.go
@@ -5,8 +5,11 @@ import (
 	"github.com/StackExchange/wmi"
 	"golang.org/x/sys/windows/registry"
 	"strings"
+	"syscall"
 )
 
+const afHvSock = 34 // AF_HYPERV
+
 func CheckHVService() bool {
 	gcs, err := registry.OpenKey(registry.LOCAL_MACHINE, `SOFTWARE\Microsoft\Windows NT\CurrentVersion\Virtualization\GuestCommunicationServices`, registry.READ)
 	if err != nil {
@@ -50,3 +53,13 @@ func GetVMID() []string {
 	}
 	return results
 }
+
+func CheckHvSocket() bool {
+	fd, err := syscall.Socket(afHvSock, syscall.SOCK_STREAM, 1)
+	if err != nil {
+		println(err.Error())
+		return false
+	}
+	syscall.Close(fd)
+	return true
+}
diff --git a/versioninfo.json b/versioninfo.json
index 860f184..6f46244 100644
--- a/versioninfo.json
+++ b/versioninfo.json
@@ -3,13 +3,13 @@
         "FileVersion": {
             "Major": 1,
             "Minor": 1,
-            "Patch": 1,
+            "Patch": 2,
             "Build": 0
         },
         "ProductVersion": {
             "Major": 1,
             "Minor": 1,
-            "Patch": 1,
+            "Patch": 2,
             "Build": 0
         },
         "FileFlagsMask": "3f",
@@ -29,7 +29,7 @@
         "OriginalFilename": "WinCryptSSHAgent.exe",
         "PrivateBuild": "",
         "ProductName": "WinCrypt SSH Agent",
-        "ProductVersion": "v1.1.1",
+        "ProductVersion": "v1.1.2",
         "SpecialBuild": ""
     },
     "VarFileInfo": {