diff --git a/src/Renci.SshNet/Abstractions/SocketAbstraction.cs b/src/Renci.SshNet/Abstractions/SocketAbstraction.cs index 080d2231..dfa2c993 100644 --- a/src/Renci.SshNet/Abstractions/SocketAbstraction.cs +++ b/src/Renci.SshNet/Abstractions/SocketAbstraction.cs @@ -3,7 +3,6 @@ using System.Globalization; 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 @@ namespace Renci.SshNet.Abstractions 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 @@ namespace Renci.SshNet.Abstractions return totalBytesRead; } -#if !NET - public static Task 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); diff --git a/src/Renci.SshNet/Abstractions/SocketExtensions.cs b/src/Renci.SshNet/Abstractions/SocketExtensions.cs index a7560506..85a595ce 100644 --- a/src/Renci.SshNet/Abstractions/SocketExtensions.cs +++ b/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() + if (cancellationToken.IsCancellationRequested) { - Completed += (sender, e) => SetCompleted(); + return Task.FromCanceled(cancellationToken); } - public AwaitableSocketAsyncEventArgs ExecuteAsync(Func func) + 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); + + return connectTask.IsCompleted || !cancellationToken.CanBeCanceled + ? connectTask + : WaitWithCancellation(connectTask, socket, cancellationToken); + + static async Task WaitWithCancellation(Task connectTask, Socket socket, CancellationToken cancellationToken) { - if (!func(this)) + using (cancellationToken.Register(static s => ((Socket)s).Dispose(), socket, useSynchronizationContext: false)) { - SetCompleted(); - } - - return this; - } - - private void SetCompleted() - { - IsCompleted = true; - var continuation = Interlocked.Exchange(ref _continuationAction, SENTINEL); - if (continuation is not null) - { - continuation(); - } - } - - public void SetCancelled() - { - _isCancelled = true; - SetCompleted(); - } - - public AwaitableSocketAsyncEventArgs GetAwaiter() - { - return this; - } - - public bool IsCompleted { get; private set; } - - void INotifyCompletion.OnCompleted(Action continuation) - { - if (_continuationAction == SENTINEL || Interlocked.CompareExchange(ref _continuationAction, continuation, comparand: null) == SENTINEL) - { - // 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 ReceiveAsync(this Socket socket, byte[] buffer, CancellationToken cancellationToken) { - cancellationToken.ThrowIfCancellationRequested(); - - using (var args = new AwaitableSocketAsyncEventArgs()) - { - args.RemoteEndPoint = remoteEndpoint; - - using (cancellationToken.Register(o => ((AwaitableSocketAsyncEventArgs)o).SetCancelled(), args, useSynchronizationContext: false)) - { - await args.ExecuteAsync(socket.ConnectAsync); - } - } + return ReceiveAsync(socket, buffer, 0, buffer.Length, cancellationToken); } - public static async Task ReceiveAsync(this Socket socket, byte[] buffer, int offset, int length, CancellationToken cancellationToken) + public static Task ReceiveAsync(this Socket socket, byte[] buffer, int offset, int length, CancellationToken cancellationToken) { - cancellationToken.ThrowIfCancellationRequested(); - - using (var args = new AwaitableSocketAsyncEventArgs()) + if (cancellationToken.IsCancellationRequested) { - args.SetBuffer(buffer, offset, length); + return Task.FromCanceled(cancellationToken); + } - using (cancellationToken.Register(o => ((AwaitableSocketAsyncEventArgs)o).SetCancelled(), args, useSynchronizationContext: false)) + 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(buffer, offset, length), + state: socket); + + return receiveTask.IsCompleted || !cancellationToken.CanBeCanceled + ? receiveTask + : WaitWithCancellation(receiveTask, socket, cancellationToken); + + static async Task WaitWithCancellation(Task 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; } } } diff --git a/src/Renci.SshNet/Connection/ConnectorBase.cs b/src/Renci.SshNet/Connection/ConnectorBase.cs index 7b1da720..44aa89e1 100644 --- a/src/Renci.SshNet/Connection/ConnectorBase.cs +++ b/src/Renci.SshNet/Connection/ConnectorBase.cs @@ -72,7 +72,7 @@ namespace Renci.SshNet.Connection 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,13 +83,6 @@ namespace Renci.SshNet.Connection } } - 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]; @@ -97,23 +90,6 @@ namespace Renci.SshNet.Connection return buffer[0]; } - /// - /// Performs a blocking read on the socket until bytes are received. - /// - /// The to read from. - /// An array of type that is the storage location for the received data. - /// The position in parameter to store the received data. - /// The number of bytes to read. - /// - /// The number of bytes read. - /// - /// The socket is closed. - /// The read failed. - protected static int SocketRead(Socket socket, byte[] buffer, int offset, int length) - { - return SocketRead(socket, buffer, offset, length, Timeout.InfiniteTimeSpan); - } - /// /// Performs a blocking read on the socket until bytes are received. /// diff --git a/src/Renci.SshNet/Connection/ProtocolVersionExchange.cs b/src/Renci.SshNet/Connection/ProtocolVersionExchange.cs index 9464977e..c0ff064d 100644 --- a/src/Renci.SshNet/Connection/ProtocolVersionExchange.cs +++ b/src/Renci.SshNet/Connection/ProtocolVersionExchange.cs @@ -181,7 +181,7 @@ namespace Renci.SshNet.Connection 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();