Abstract away socket logic and make able to handle both named pipes f…

…or Win10 SSH and AF_UNIX sockets for WSL

PageantHandler.cs

@@ -81,6 +81,7 @@ struct COPYDATASTRUCT
         [DllImport("kernel32.dll", SetLastError = true)]
         static extern bool CloseHandle(IntPtr hHandle);
 
+        internal static ArraySegment<byte> AGENT_EMPTY_RESPONSE = new ArraySegment<byte>(new byte[] { 0x00, 0x00, 0x00, 0x05, 0x0c, 0x00, 0x00, 0x00, 0x00 });
         internal const uint AGENT_MAX_MSGLEN = 8192;
         static readonly IntPtr AGENT_COPYDATA_ID = new IntPtr(0x804e50ba);
 

Program.cs

@@ -1,66 +1,54 @@
-using System;
-using System.IO;
-using System.Net.Sockets;
-using System.Threading;
+using Microsoft.Extensions.CommandLineUtils;
+using System.Collections.Generic;
 using System.Threading.Tasks;
 
 namespace WslSSHPageant
 {
     class Program
     {
-        static Mutex mutex;
-
-        static async Task Main(string[] args)
+        static void Main(string[] args)
         {
-            var socketPath = @".\ssh-agent.sock";
+            CommandLineApplication commandLineApplication = new CommandLineApplication(throwOnUnexpectedArg: false);
 
-            if (args.Length == 1)
-            {
-                socketPath = args[0];
-            }
-            else if (args.Length != 0)
-            {
-                Console.WriteLine(@"wsl-ssh-agent.exe <path: .\ssh-agent.sock>");
-                return;
-            }
+            CommandOption wslSocketPath = commandLineApplication.Option(
+                "--wsl <path>",
+                "Which path to listen on with the AF_UNIX socket for WSL",
+                CommandOptionType.SingleValue);
 
-            socketPath = Path.GetFullPath(socketPath);
+            CommandOption winsshPipeName = commandLineApplication.Option(
+                "--winssh <name>",
+                "Which pipe to listen on for Windows 10 OpenSSH Client",
+                CommandOptionType.SingleValue);
 
-            var mutexName = socketPath + "-{642b3e23-f0f5-4cc1-8a41-bf95e9a438ad}";
-            mutexName = mutexName.Replace(Path.DirectorySeparatorChar, '_');
-            mutex = new Mutex(true, mutexName);
+            commandLineApplication.HelpOption("-? | -h | --help");
 
-            if (!mutex.WaitOne(TimeSpan.Zero, true))
-            {
-                Console.Error.WriteLine("Already running on that AF_UNIX path");
-                Console.In.ReadLine();
-                return;
-            }
+            List<Task> runningServers = new List<Task>();
 
-            try
+            commandLineApplication.OnExecute(() =>
             {
-                File.Delete(socketPath);
-                var server = new Socket(AddressFamily.Unix, SocketType.Stream, ProtocolType.IP);
-                server.Bind(new UnixEndPoint(socketPath));
-                server.Listen(5);
-
-                Console.WriteLine(@"Listening on {0}", socketPath);
-
-                // Enter the listening loop.
-                while (true)
+                if (wslSocketPath.HasValue())
+                {
+                    WSLSocket wslSocket = new WSLSocket(wslSocketPath.Value());
+                    runningServers.Add(wslSocket.Listen());
+                }
+                if (winsshPipeName.HasValue())
                 {
-                    WSLClient client = new WSLClient(await server.AcceptAsync());
+                    WinSSHSocket winsshSocket = new WinSSHSocket(winsshPipeName.Value());
+                    runningServers.Add(winsshSocket.Listen());
+                }
 
-                    // Don't await this, we want to service other sockets
-#pragma warning disable CS4014
-                    client.WorkSocket();
-#pragma warning restore CS4014
+                if (runningServers.Count < 1)
+                {
+                    commandLineApplication.ShowHelp();
+                    return 1;
                 }
-            }
-            finally
-            {
-                mutex.ReleaseMutex();
-            }
+
+                Task.WaitAny(runningServers.ToArray());
+
+                return 0;
+            });
+
+            commandLineApplication.Execute(args);
         }
     }
 }

Readme.md

@@ -1,12 +1,10 @@
 # wsl-ssh-pageant
 
-**Now uses freshly baked AF_UNIX support in Windows 10 insider**
-
-## How to use
+## How to use with WSL
 
 1. On the Windows side run Pageant (or compatible agent such as gpg4win).
 
-2. Run wsl-ssh-pageant.exe on windows in a short path (max ~100 characters total!)
+2. Run `wsl-ssh-pageant.exe --wsl C:\wsl-ssh-pageant` (or any other path) on windows in a short path (max ~100 characters total!)
 
 3. In WSL run the following
 
@@ -20,6 +18,25 @@ $ export SSH_AUTH_SOCK=/mnt/c/wsl-ssh-pageant/ssh-agent.sock
 
 4. The SSH keys from Pageant should now be usable by `ssh`!
 
+## How to use with Windows 10 native OpenSSH client
+
+1. On the Windows side run Pageant (or compatible agent such as gpg4win).
+
+2. Run `wsl-ssh-pageant.exe --winssh ssh-pageant` (or any other name) on windows in a short path (max ~100 characters total!)
+
+3. In cmd.exe run the following (or define it in your Environment Variables on windows)
+
+```
+$ set SSH_AUTH_SOCK=\\.\pipe\ssh-pageant
+```
+(or whichever name you gave the pipe)
+
+4. The SSH keys from Pageant should now be usable by `ssh`!
+
+## Note
+
+You can use both `--winssh` and `--wsl` parameters at the same time with the same process to proxy for both
+
 ## Credit
 
 Thanks to [John Starks](https://github.com/jstarks/) for [npiperelay](https://github.com/jstarks/npiperelay/), showing a more secure way to create a stream between WSL and Linux.

SSHAgentClient.cs

@@ -0,0 +1,138 @@
+using System;
+using System.Net.Sockets;
+using System.Threading.Tasks;
+
+namespace WslSSHPageant
+{
+    internal abstract class SSHAgentClientPartialRead : SSHAgentClient
+    {
+        internal SSHAgentClientPartialRead()
+        {
+        }
+
+        protected override async Task<bool> ReceiveArraySegment(ArraySegment<byte> buf)
+        {
+            int i;
+            while ((i = await ReceivePartialArraySegment(buf)) != 0)
+            {
+                buf = buf.Slice(i);
+                if (buf.Count <= 0)
+                {
+                    return true;
+                }
+            }
+            return false;
+        }
+
+        protected abstract Task<int> ReceivePartialArraySegment(ArraySegment<byte> buf);
+    }
+
+    internal abstract class SSHAgentClient
+    {
+        internal SSHAgentClient()
+        {
+        }
+
+        protected virtual void Initialize()
+        {
+            return;
+        }
+
+        protected abstract bool IsConnected();
+
+        protected abstract void Close();
+
+        protected abstract Task<bool> ReceiveArraySegment(ArraySegment<byte> buf);
+
+        protected abstract Task<bool> SendArraySegment(ArraySegment<byte> buf);
+
+        internal async Task WorkSocket()
+        {
+            Initialize();
+
+            bool clientWasSuccess = false;
+
+            try
+            {
+                clientWasSuccess = await ServiceSocket();
+            }
+            catch (TimeoutException)
+            {
+                // Ignore timeouts, those should not explode our stuff
+                Console.Error.WriteLine("Socket timeout");
+            }
+            // These two just mean the remote end closed the socket, we don't care, same for TaskCanceledException
+            catch (ObjectDisposedException) { }
+            catch (InvalidOperationException) { }
+            catch (TaskCanceledException) { }
+            catch (SocketException e)
+            {
+                // Other socket errors can happen and shouldn't kill the app
+                Console.Error.WriteLine(e);
+            }
+            catch (PageantException e)
+            {
+                // Pageant errors can happen, too
+                Console.Error.WriteLine(e);
+            }
+            catch (Exception e)
+            {
+                Console.Error.WriteLine(e);
+                throw e;
+            }
+            finally
+            {
+                if (IsConnected() && !clientWasSuccess)
+                {
+                    try
+                    {
+                        await SendArraySegment(PageantHandler.AGENT_EMPTY_RESPONSE);
+                    }
+                    catch { }
+                }
+
+                Close();
+            }
+        }
+
+        private async Task<bool> ServiceSocket()
+        {
+            var bytes = new byte[PageantHandler.AGENT_MAX_MSGLEN];
+
+            bool lastWasSuccess = true;
+
+            while (IsConnected())
+            {
+                // Read length as uint32 (4 bytes)
+                if (!await ReceiveArraySegment(new ArraySegment<byte>(bytes, 0, 4)))
+                {
+                    break;
+                }
+
+                lastWasSuccess = false;
+
+                var len = (bytes[0] << 24) |
+                          (bytes[1] << 16) |
+                          (bytes[2] << 8) |
+                          (bytes[3]);
+
+                if (len + 4 > PageantHandler.AGENT_MAX_MSGLEN)
+                {
+                    break;
+                }
+
+                // Read actual data in the part after len
+                if (!await ReceiveArraySegment(new ArraySegment<byte>(bytes, 4, len)))
+                {
+                    break;
+                }
+
+                var msg = PageantHandler.Query(new ArraySegment<byte>(bytes, 0, len + 4));
+                await SendArraySegment(new ArraySegment<byte>(msg, 0, msg.Length));
+                lastWasSuccess = true;
+            }
+
+            return lastWasSuccess;
+        }
+    }
+}