fix

masahide · Aug 25, 2024 · 6c04425 · 6c04425
1 parent cb61fa8
commit 6c04425
Show file tree

Hide file tree

Showing 4 changed files with 118 additions and 33 deletions.
diff --git a/cmd/wsl2-ssh-agent-proxy/main.go b/cmd/wsl2-ssh-agent-proxy/main.go
@@ -5,14 +5,17 @@ import (
 	"bytes"
 	_ "embed"
 	"encoding/binary"
+	"errors"
 	"io"
 	"log"
 	"net"
 	"os"
 	"os/exec"
+	"os/signal"
 	"path/filepath"
 	"strings"
 	"sync"
+	"syscall"
 )
 
 //go:embed pwsh.ps1
@@ -62,6 +65,10 @@ func (mux *Multiplexer) readLoop() {
 		header := make([]byte, HeaderSize)
 		_, err := io.ReadFull(mux.reader, header)
 		if err != nil {
+			if errors.Is(err, io.EOF) {
+				log.Println("Connection closed")
+				return
+			}
 			log.Println("Error reading header:", err)
 			return
 		}
@@ -155,6 +162,7 @@ func handleConnection(conn net.Conn, mux *Multiplexer, channelID uint32) {
 		writer.Flush()
 	}
 }
+
 func checkStartAgent(r io.Reader) {
 	buf := make([]byte, 1024)
 	n, err := r.Read(buf)
@@ -193,6 +201,11 @@ func main() {
 	if err != nil {
 		exePath = ("powershell.exe")
 	}
+	socketPath := os.Getenv("SSH_AUTH_SOCK")
+	if len(socketPath) == 0 {
+		log.Fatal("env SSH_AUTH_SOCK is not set")
+	}
+
 	cmd := exec.Command(exePath, "-NoProfile", "-Command", "-")
 
 	psIn, err := cmd.StdinPipe()
@@ -221,17 +234,37 @@ func main() {
 
 	mux := NewMultiplexer(psIn, psOut)
 
-	socketPath := os.Getenv("SSH_AUTH_SOCK")
-	if len(socketPath) == 0 {
-		log.Fatal("env SSH_AUTH_SOCK is not set")
+	// Capture kill signals
+	sigChan := make(chan os.Signal, 1)
+	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
+
+	killPwsh := func() int {
+		defer psIn.Close()
+		defer psOut.Close()
+		if err := cmd.Process.Signal(syscall.SIGTERM); err != nil {
+			log.Printf("send signal err:%s", err)
+			return 1
+		}
+		log.Printf("PowerShell process exited")
+		return 0
 	}
+	// Handle kill signals
+	go func() {
+		sig := <-sigChan
+		log.Printf("Received signal: %s. Shutting down PowerShell process...", sig)
+		os.Exit(killPwsh())
+	}()
+
 	log.Printf("get env SSH_AUTH_SOCK:%s", socketPath)
 	if _, err := os.Stat(socketPath); err == nil {
 		err := os.Remove(socketPath)
 		if err != nil {
 			log.Fatal("remove old socket err:", err)
 		}
 	}
+
+	log.Printf("Started PowerShell process with PID: %d", cmd.Process.Pid)
+
 	log.Printf("listen socket:%s", socketPath)
 	listener, err := net.Listen("unix", socketPath)
 	if err != nil {
@@ -244,11 +277,16 @@ func main() {
 	for {
 		conn, err := listener.Accept()
 		if err != nil {
+			if errors.Is(err, net.ErrClosed) {
+				log.Println("Listener has been closed.")
+				break
+			}
 			log.Println("Error accepting connection:", err)
 			continue
 		}
 		//log.Printf("accept: %v", conn.LocalAddr())
 		go handleConnection(conn, mux, channelID)
 		channelID++
 	}
+	os.Exit(killPwsh())
 }
diff --git a/cmd/wsl2-ssh-agent-proxy/pwsh.ps1 b/cmd/wsl2-ssh-agent-proxy/pwsh.ps1
@@ -165,9 +165,17 @@ class PacketReader {
         [Console]::Error.WriteLine("PacketReader started.")
         while ($true) {
             #[Console]::Error.WriteLine("PacketReader: Waiting for packets.")
-            $Packet = $this.ReadPacket($this.InputStreamReader)
+            $Packet = $null
+            try {
+                $Packet = $this.ReadPacket($this.InputStreamReader)
+            }
+            catch {
+                [Console]::Error.WriteLine("InputStreamRead error: [$error]")
+                return
+            }
             if ($null -ne $Packet.Error) {
                 [Console]::Error.WriteLine($Packet.Error)
+                Start-Sleep -Seconds 1.0
                 continue
             }
             #[Console]::Error.WriteLine("PacketReader: Packet received. Type: $($Packet.TypeNum), Channel ID: $($Packet.ChannelID), Length: $($Packet.Length).")

diff --git a/hack/ubuntu.setup.sh b/hack/ubuntu.setup.sh
@@ -1,41 +1,40 @@
 #!/bin/sh
 
-set -ex
-NAME=wsl2-ssh-agent-proxy
-SSH_AUTH_SOCK="${HOME}/.ssh/${NAME}/${NAME}.sock"
-PROXYCMD_DIR="${HOME}/${NAME}"
-CMD="${PROXYCMD_DIR}/${NAME}"
-
-RELEASE_NAME=$1
-REPO_URL=https://github.com/masahide/OmniSSHAgent
-if [ -z "${RELEASE_NAME}" ]; then
-  VER_PATH="releases/latest"
-else
-  VER_PATH="download/${RELEASE_NAME}"
-fi
-
-__get_proxy() {
-  echo "Downloading ${NAME}.gz"
-  mkdir -p "${PROXYCMD_DIR}"
-  curl "${REPO_URL}/releases/${VER_PATH}/${NAME}.gz" -sL | gunzip > "${CMD}"
-  chmod +x "${CMD}"
+OMNISOCATCMD="$HOME/omni-socat/omni-socat.exe"
+export SSH_AUTH_SOCK="$HOME/.ssh/agent.sock"
+
+__get_omnisocat() {
+  echo "Downloading omni-socat.exe"
+  sudo apt -y install unzip
+  curl https://github.com/masahide/OmniSSHAgent/releases/latest/download/omni-socat.zip \
+    -sLo omni-socat.zip
+  unzip -o omni-socat.zip -d "$(dirname "$OMNISOCATCMD")"
+  chmod +x "$OMNISOCATCMD"
+  rm -f omni-socat.zip
 }
 
-setup_proxy() {
-  [ -f "${CMD}" ] || __get_proxy
+__get_socat() {
+  echo "Installing socat"
+  sudo apt -y install socat
+}
+
+setup_omnisocat() {
+  [ -f "$OMNISOCATCMD" ] || __get_omnisocat
+  command -v socat >/dev/null 2>&1 || __get_socat
 
   # Checks wether $SSH_AUTH_SOCK is a socket or not
-  (ps aux | grep "${CMD}" | grep -qv "grep") && [ -S "${SSH_AUTH_SOCK}" ] && return
+  (ss -a | grep -q "$SSH_AUTH_SOCK") && [ -S "$SSH_AUTH_SOCK" ] && return
 
   # Create directory for the socket, if it is missing
-  SSH_AUTH_SOCK_DIR="$(dirname "${SSH_AUTH_SOCK}")"
-  mkdir -p "${SSH_AUTH_SOCK_DIR}"
-  # Applying best-practice permissions if we are creating ${HOME}/.ssh
-  if [ "${SSH_AUTH_SOCK_DIR}" = "${HOME}/.ssh" ]; then
-    chmod 700 "${SSH_AUTH_SOCK_DIR}"
+  SSH_AUTH_SOCK_DIR="$(dirname "$SSH_AUTH_SOCK")"
+  mkdir -p "$SSH_AUTH_SOCK_DIR"
+  # Applying best-practice permissions if we are creating $HOME/.ssh
+  if [ "$SSH_AUTH_SOCK_DIR" = "$HOME/.ssh" ]; then
+    chmod 700 "$SSH_AUTH_SOCK_DIR"
   fi
 
-  ${CMD} >> "${PROXYCMD_DIR}/${NAME}.log" 2>&1 &
+  rm -f "$SSH_AUTH_SOCK"
+  (setsid socat UNIX-LISTEN:"$SSH_AUTH_SOCK",fork EXEC:"$OMNISOCATCMD",nofork &) >/dev/null 2>&1
 }
 
-setup_proxy
+setup_omnisocat
diff --git a/hack/ubuntu.wsl2-ssh-agent-proxy.sh b/hack/ubuntu.wsl2-ssh-agent-proxy.sh
@@ -0,0 +1,40 @@
+#!/bin/sh
+
+NAME=wsl2-ssh-agent-proxy
+export SSH_AUTH_SOCK="${HOME}/.ssh/${NAME}/${NAME}.sock"
+PROXYCMD_DIR="${HOME}/${NAME}"
+CMD="${PROXYCMD_DIR}/${NAME}"
+
+RELEASE_NAME=$1
+REPO_URL=https://github.com/masahide/OmniSSHAgent
+if [ -z "${RELEASE_NAME}" ]; then
+  VER_PATH="releases/latest"
+else
+  VER_PATH="download/${RELEASE_NAME}"
+fi
+
+__get_proxy() {
+  echo "Downloading ${NAME}.gz"
+  mkdir -p "${PROXYCMD_DIR}"
+  curl "${REPO_URL}/releases/${VER_PATH}/${NAME}.gz" -sL | gunzip >"${CMD}"
+  chmod +x "${CMD}"
+}
+
+setup_proxy() {
+  [ -f "${CMD}" ] || __get_proxy
+
+  # Checks wether $SSH_AUTH_SOCK is a socket or not
+  (ps aux | grep "${CMD}" | grep -qv "grep") && [ -S "${SSH_AUTH_SOCK}" ] && return
+
+  # Create directory for the socket, if it is missing
+  SSH_AUTH_SOCK_DIR="$(dirname "${SSH_AUTH_SOCK}")"
+  mkdir -p "${SSH_AUTH_SOCK_DIR}"
+  # Applying best-practice permissions if we are creating ${HOME}/.ssh
+  if [ "${SSH_AUTH_SOCK_DIR}" = "${HOME}/.ssh" ]; then
+    chmod 700 "${SSH_AUTH_SOCK_DIR}"
+  fi
+
+  nohup "${CMD}" >>"${PROXYCMD_DIR}/${NAME}.log" 2>&1 &
+}
+
+setup_proxy