Remove AwaitableSocketAsyncEventArgs

src/Renci.SshNet/Abstractions/SocketAbstraction.cs

@@ -3,7 +3,6 @@
 using System.Net;
 using System.Net.Sockets;
 using System.Threading;
-using System.Threading.Tasks;
 
 using Renci.SshNet.Common;
 using Renci.SshNet.Messages.Transport;
@@ -26,11 +25,6 @@ public static void Connect(Socket socket, EndPoint remoteEndpoint, TimeSpan conn
             ConnectCore(socket, remoteEndpoint, connectTimeout, ownsSocket: false);
         }
 
-        public static async Task ConnectAsync(Socket socket, EndPoint remoteEndpoint, CancellationToken cancellationToken)
-        {
-            await socket.ConnectAsync(remoteEndpoint, cancellationToken).ConfigureAwait(false);
-        }
-
         private static void ConnectCore(Socket socket, EndPoint remoteEndpoint, TimeSpan connectTimeout, bool ownsSocket)
         {
             var connectCompleted = new ManualResetEvent(initialState: false);
@@ -251,13 +245,6 @@ public static int Read(Socket socket, byte[] buffer, int offset, int size, TimeS
             return totalBytesRead;
         }
 
-#if !NET
-        public static Task<int> ReadAsync(Socket socket, byte[] buffer, CancellationToken cancellationToken)
-        {
-            return socket.ReceiveAsync(buffer, 0, buffer.Length, cancellationToken);
-        }
-#endif
-
         public static void Send(Socket socket, byte[] data)
         {
             Send(socket, data, 0, data.Length);

src/Renci.SshNet/Abstractions/SocketExtensions.cs

@@ -2,118 +2,69 @@
 using System;
 using System.Net;
 using System.Net.Sockets;
-using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
 
 namespace Renci.SshNet.Abstractions
 {
-    // Async helpers based on https://devblogs.microsoft.com/pfxteam/awaiting-socket-operations/
     internal static class SocketExtensions
     {
-        private sealed class AwaitableSocketAsyncEventArgs : SocketAsyncEventArgs, INotifyCompletion
+        public static Task ConnectAsync(this Socket socket, EndPoint remoteEndpoint, CancellationToken cancellationToken)
         {
-            private static readonly Action SENTINEL = () => { };
-
-            private bool _isCancelled;
-            private Action _continuationAction;
-
-            public AwaitableSocketAsyncEventArgs()
-            {
-                Completed += (sender, e) => SetCompleted();
-            }
-
-            public AwaitableSocketAsyncEventArgs ExecuteAsync(Func<SocketAsyncEventArgs, bool> func)
-            {
-                if (!func(this))
-                {
-                    SetCompleted();
-                }
-
-                return this;
-            }
-
-            private void SetCompleted()
-            {
-                IsCompleted = true;
-                var continuation = Interlocked.Exchange(ref _continuationAction, SENTINEL);
-                if (continuation is not null)
-                {
-                    continuation();
-                }
-            }
-
-            public void SetCancelled()
+            if (cancellationToken.IsCancellationRequested)
             {
-                _isCancelled = true;
-                SetCompleted();
+                return Task.FromCanceled(cancellationToken);
             }
 
-            public AwaitableSocketAsyncEventArgs GetAwaiter()
-            {
-                return this;
-            }
+            var connectTask = Task.Factory.FromAsync(
+                    static (s, remoteEndpoint, callback, state) => s.BeginConnect(remoteEndpoint, callback, state),
+                    static result => ((Socket)result.AsyncState).EndConnect(result),
+                    socket,
+                    remoteEndpoint,
+                    state: socket);
 
-            public bool IsCompleted { get; private set; }
+            return connectTask.IsCompleted || !cancellationToken.CanBeCanceled
+                ? connectTask
+                : WaitWithCancellation(connectTask, socket, cancellationToken);
 
-            void INotifyCompletion.OnCompleted(Action continuation)
+            static async Task WaitWithCancellation(Task connectTask, Socket socket, CancellationToken cancellationToken)
             {
-                if (_continuationAction == SENTINEL || Interlocked.CompareExchange(ref _continuationAction, continuation, comparand: null) == SENTINEL)
+                using (cancellationToken.Register(static s => ((Socket)s).Dispose(), socket, useSynchronizationContext: false))
                 {
-                    // We have already completed; run continuation asynchronously
-                    _ = Task.Run(continuation);
-                }
-            }
-
-            public void GetResult()
-            {
-                if (_isCancelled)
-                {
-                    throw new TaskCanceledException();
-                }
-
-                if (!IsCompleted)
-                {
-                    // We don't support sync/async
-                    throw new InvalidOperationException("The asynchronous operation has not yet completed.");
-                }
-
-                if (SocketError != SocketError.Success)
-                {
-                    throw new SocketException((int)SocketError);
+                    await connectTask.ConfigureAwait(false);
                 }
             }
         }
 
-        public static async Task ConnectAsync(this Socket socket, EndPoint remoteEndpoint, CancellationToken cancellationToken)
+        public static Task<int> ReceiveAsync(this Socket socket, byte[] buffer, CancellationToken cancellationToken)
         {
-            cancellationToken.ThrowIfCancellationRequested();
+            return ReceiveAsync(socket, buffer, 0, buffer.Length, cancellationToken);
+        }
 
-            using (var args = new AwaitableSocketAsyncEventArgs())
+        public static Task<int> ReceiveAsync(this Socket socket, byte[] buffer, int offset, int length, CancellationToken cancellationToken)
+        {
+            if (cancellationToken.IsCancellationRequested)
             {
-                args.RemoteEndPoint = remoteEndpoint;
-
-                using (cancellationToken.Register(o => ((AwaitableSocketAsyncEventArgs)o).SetCancelled(), args, useSynchronizationContext: false))
-                {
-                    await args.ExecuteAsync(socket.ConnectAsync);
-                }
+                return Task.FromCanceled<int>(cancellationToken);
             }
-        }
 
-        public static async Task<int> ReceiveAsync(this Socket socket, byte[] buffer, int offset, int length, CancellationToken cancellationToken)
-        {
-            cancellationToken.ThrowIfCancellationRequested();
+            var receiveTask = Task.Factory.FromAsync(
+                    static (s, buffer, callback, state) => s.BeginReceive(buffer.Array, buffer.Offset, buffer.Count, SocketFlags.None, callback, state),
+                    static result => ((Socket)result.AsyncState).EndReceive(result),
+                    socket,
+                    new ArraySegment<byte>(buffer, offset, length),
+                    state: socket);
 
-            using (var args = new AwaitableSocketAsyncEventArgs())
-            {
-                args.SetBuffer(buffer, offset, length);
+            return receiveTask.IsCompleted || !cancellationToken.CanBeCanceled
+                ? receiveTask
+                : WaitWithCancellation(receiveTask, socket, cancellationToken);
 
-                using (cancellationToken.Register(o => ((AwaitableSocketAsyncEventArgs)o).SetCancelled(), args, useSynchronizationContext: false))
+            static async Task<int> WaitWithCancellation(Task<int> receiveTask, Socket socket, CancellationToken cancellationToken)
+            {
+                using (cancellationToken.Register(static s => ((Socket)s).Dispose(), socket, useSynchronizationContext: false))
                 {
-                    await args.ExecuteAsync(socket.ReceiveAsync);
+                    return await receiveTask.ConfigureAwait(false);
                 }
-
-                return args.BytesTransferred;
             }
         }
     }

src/Renci.SshNet/Connection/ConnectorBase.cs

@@ -72,7 +72,7 @@ protected async Task<Socket> SocketConnectAsync(EndPoint endPoint, CancellationT
             var socket = SocketFactory.Create(SocketType.Stream, ProtocolType.Tcp);
             try
             {
-                await SocketAbstraction.ConnectAsync(socket, endPoint, cancellationToken).ConfigureAwait(false);
+                await socket.ConnectAsync(endPoint, cancellationToken).ConfigureAwait(false);
 
                 return socket;
             }
@@ -83,37 +83,13 @@ protected async Task<Socket> SocketConnectAsync(EndPoint endPoint, CancellationT
             }
         }
 
-        protected static byte SocketReadByte(Socket socket)
-        {
-            var buffer = new byte[1];
-            _ = SocketRead(socket, buffer, 0, 1, Timeout.InfiniteTimeSpan);
-            return buffer[0];
-        }
-
         protected static byte SocketReadByte(Socket socket, TimeSpan readTimeout)
         {
             var buffer = new byte[1];
             _ = SocketRead(socket, buffer, 0, 1, readTimeout);
             return buffer[0];
         }
 
-        /// <summary>
-        /// Performs a blocking read on the socket until <paramref name="length"/> bytes are received.
-        /// </summary>
-        /// <param name="socket">The <see cref="Socket"/> to read from.</param>
-        /// <param name="buffer">An array of type <see cref="byte"/> that is the storage location for the received data.</param>
-        /// <param name="offset">The position in <paramref name="buffer"/> parameter to store the received data.</param>
-        /// <param name="length">The number of bytes to read.</param>
-        /// <returns>
-        /// The number of bytes read.
-        /// </returns>
-        /// <exception cref="SshConnectionException">The socket is closed.</exception>
-        /// <exception cref="SocketException">The read failed.</exception>
-        protected static int SocketRead(Socket socket, byte[] buffer, int offset, int length)
-        {
-            return SocketRead(socket, buffer, offset, length, Timeout.InfiniteTimeSpan);
-        }
-
         /// <summary>
         /// Performs a blocking read on the socket until <paramref name="length"/> bytes are received.
         /// </summary>

src/Renci.SshNet/Connection/ProtocolVersionExchange.cs

@@ -181,7 +181,7 @@ private static async Task<string> SocketReadLineAsync(Socket socket, List<byte>
 
             while (buffer.Count < MaximumBannerLineLength)
             {
-                var bytesRead = await SocketAbstraction.ReadAsync(socket, data, cancellationToken).ConfigureAwait(false);
+                var bytesRead = await socket.ReceiveAsync(data, cancellationToken).ConfigureAwait(false);
                 if (bytesRead == 0)
                 {
                     throw CreateConnectionLostException();