add cancellation support to TcpSocketClient

rdavisau · Jan 10, 2016 · 0896e56 · 0896e56
1 parent e7b1482
commit 0896e56
Show file tree

Hide file tree

Showing 8 changed files with 68 additions and 18 deletions.
diff --git a/Sockets/Sockets.Implementation.NET/Extensions/NativeExceptionExtensions.cs b/Sockets/Sockets.Implementation.NET/Extensions/NativeExceptionExtensions.cs
@@ -16,10 +16,11 @@ public static Task WrapNativeSocketExceptions(this Task task)
             return task.ContinueWith(
                 t =>
                 {
+
                     if (!t.IsFaulted)
                         return t;
-
-                    var ex = t.Exception.InnerException;
+                  
+                    var ex = t.Exception?.InnerException ?? t.Exception;
 
                     throw (NativeExceptionExtensions.NativeSocketExceptions.Contains(ex.GetType())) 
                         ? new PclSocketException(ex)
@@ -35,7 +36,7 @@ public static Task<T> WrapNativeSocketExceptions<T>(this Task<T> task)
                     if (!t.IsFaulted)
                         return t.Result;
 
-                    var ex = t.Exception.InnerException;
+                    var ex = t.Exception?.InnerException ?? t.Exception;
 
                     throw (NativeExceptionExtensions.NativeSocketExceptions.Contains(ex.GetType()))
                         ? new PclSocketException(ex)

diff --git a/Sockets/Sockets.Implementation.NET/TcpSocketClient.cs b/Sockets/Sockets.Implementation.NET/TcpSocketClient.cs
@@ -6,7 +6,7 @@
 using Sockets.Plugin.Abstractions;
 using System.Net.Security;
 using System.Security.Cryptography.X509Certificates;
-
+using System.Threading;
 using PlatformSocketException = System.Net.Sockets.SocketException;
 using PclSocketException = Sockets.Plugin.Abstractions.SocketException;
 
@@ -55,11 +55,33 @@ internal TcpSocketClient(TcpClient backingClient, int bufferSize)
         /// <param name="address">The address of the endpoint to connect to.</param>
         /// <param name="port">The port of the endpoint to connect to.</param>
         /// <param name="secure">True to enable TLS on the socket.</param>
-        public async Task ConnectAsync(string address, int port, bool secure = false)
+        /// <param name="cancellationToken">The cancellation token to cancel the operation.</param>
+        public async Task ConnectAsync(string address, int port, bool secure = false, CancellationToken cancellationToken = default(CancellationToken))
         {
-            await _backingTcpClient
-                .ConnectAsync(address, port)
-                .WrapNativeSocketExceptions();
+            // standard connect
+            var connectTask =
+                _backingTcpClient
+                    .ConnectAsync(address, port)
+                    .WrapNativeSocketExceptions();
+
+            // set up cancellation trigger
+            var ret = new TaskCompletionSource<bool>();
+            var canceller = cancellationToken.Register(() => ret.SetCanceled());
+
+            // if cancellation comes before connect completes, we honour it
+            var okOrCancelled = await Task.WhenAny(connectTask, ret.Task);
+
+            if (okOrCancelled == ret.Task)
+            {
+                // reset the backing field.
+                // depending on the state of the socket this may throw ODE which it is appropriate to ignore
+                try { await DisconnectAsync(); } catch (ObjectDisposedException) { }
+
+                // notify that we did cancel
+                cancellationToken.ThrowIfCancellationRequested();
+            }
+            else
+                canceller.Dispose();
 
             InitializeWriteStream();
 

diff --git a/Sockets/Sockets.Implementation.WinRT/Extensions/NativeExceptionExtensions.cs b/Sockets/Sockets.Implementation.WinRT/Extensions/NativeExceptionExtensions.cs
@@ -1,6 +1,7 @@
 using System;
 using System.Collections.Generic;
 using System.Diagnostics;
+using System.Threading;
 using System.Threading.Tasks;
 using Windows.Foundation;
 using Windows.Networking.Sockets;
@@ -11,9 +12,11 @@ namespace Sockets.Plugin
 {
     public static class NativeExceptionExtensions
     {
-        public static Task WrapNativeSocketExceptionsAsTask(this IAsyncAction task)
+        public static Task WrapNativeSocketExceptionsAsTask(this IAsyncAction task, CancellationToken cancellationToken = default(CancellationToken))
         {
             var tcs = new TaskCompletionSource<bool>();
+            if (cancellationToken != default(CancellationToken))
+                cancellationToken.Register(task.Cancel);
 
             task.Completed = delegate(IAsyncAction info, AsyncStatus status)
             {
@@ -34,10 +37,13 @@ public static Task WrapNativeSocketExceptionsAsTask(this IAsyncAction task)
                         break;
                 }
             };
-
+            
             return tcs.Task.ContinueWith(
                 t =>
                 {
+                    if (t.IsCanceled)
+                        cancellationToken.ThrowIfCancellationRequested();
+
                     if (!t.IsFaulted)
                         return t;
                     

diff --git a/Sockets/Sockets.Implementation.WinRT/TcpSocketClient.cs b/Sockets/Sockets.Implementation.WinRT/TcpSocketClient.cs
@@ -1,5 +1,6 @@
 using System;
 using System.IO;
+using System.Threading;
 using System.Threading.Tasks;
 using Windows.Networking;
 using Windows.Networking.Sockets;
@@ -22,7 +23,7 @@ public class TcpSocketClient : ITcpSocketClient
 #else
         private SocketProtectionLevel _secureSocketProtectionLevel = SocketProtectionLevel.Tls10;
 #endif               
-        private readonly StreamSocket _backingStreamSocket;
+        private StreamSocket _backingStreamSocket;
         private readonly int _bufferSize;
 
         /// <summary>
@@ -54,15 +55,16 @@ internal TcpSocketClient(StreamSocket nativeSocket, int bufferSize)
         /// <param name="address">The address of the endpoint to connect to.</param>
         /// <param name="port">The port of the endpoint to connect to.</param>
         /// <param name="secure">True to enable TLS on the socket.</param>
-        public Task ConnectAsync(string address, int port, bool secure = false)
+        /// <param name="cancellationToken">The cancellation token to cancel the operation.</param>
+        public Task ConnectAsync(string address, int port, bool secure = false, CancellationToken cancellationToken = default(CancellationToken))
         {
             var hn = new HostName(address);
             var sn = port.ToString();
             var spl = secure ? _secureSocketProtectionLevel : SocketProtectionLevel.PlainSocket;
 
             return _backingStreamSocket
                 .ConnectAsync(hn, sn, spl)
-                .WrapNativeSocketExceptionsAsTask();
+                .WrapNativeSocketExceptionsAsTask(cancellationToken);
         }
 
         /// <summary>
@@ -71,7 +73,11 @@ public Task ConnectAsync(string address, int port, bool secure = false)
         /// </summary>
         public Task DisconnectAsync()
         {
-            return Task.Run(() => _backingStreamSocket.Dispose());
+            return Task.Run(() =>
+            {
+                _backingStreamSocket.Dispose();
+                _backingStreamSocket = new StreamSocket();
+            });
         }
 
         /// <summary>

diff --git a/Sockets/Sockets.Plugin.Abstractions/ITcpSocketClient.cs b/Sockets/Sockets.Plugin.Abstractions/ITcpSocketClient.cs
@@ -1,5 +1,6 @@
 using System;
 using System.IO;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace Sockets.Plugin.Abstractions
@@ -18,7 +19,8 @@ public interface ITcpSocketClient : IDisposable
         /// <param name="address">The address of the endpoint to connect to.</param>
         /// <param name="port">The port of the endpoint to connect to.</param>
         /// <param name="secure">Is this socket secure?</param>
-        Task ConnectAsync(string address, int port, bool secure = false);
+        /// <param name="cancellationToken">The cancellation token to cancel the operation.</param>
+        Task ConnectAsync(string address, int port, bool secure = false, CancellationToken cancellationToken = default(CancellationToken));
 
         /// <summary>
         ///     Disconnects from an endpoint previously connected to using <code>ConnectAsync</code>.

diff --git a/Sockets/Sockets.Plugin/TcpSocketClient.cs b/Sockets/Sockets.Plugin/TcpSocketClient.cs
@@ -1,5 +1,6 @@
 using System;
 using System.IO;
+using System.Threading;
 using System.Threading.Tasks;
 using Sockets.Plugin.Abstractions;
 
@@ -35,8 +36,8 @@ public TcpSocketClient(int bufferSize) : this()
         /// </summary>
         /// <param name="address">The address of the endpoint to connect to.</param>
         /// <param name="port">The port of the endpoint to connect to.</param>
-        /// <param name="secure">True to enable TLS on the socket.</param>
-        public Task ConnectAsync(string address, int port, bool secure = false)
+        /// <param name="cancellationToken">The cancellation token to cancel the operation.</param>
+        public Task ConnectAsync(string address, int port, bool secure = false, CancellationToken cancellationToken = default(CancellationToken))
         {
             throw new NotImplementedException(PCL.BaitWithoutSwitchMessage);
         }

diff --git a/Sockets/Tests/Sockets.Tests.DroidRunner/Resources/Resource.Designer.cs b/Sockets/Tests/Sockets.Tests.DroidRunner/Resources/Resource.Designer.cs
diff --git a/Sockets/Tests/Sockets.Tests/TcpSocketClientTests.cs b/Sockets/Tests/Sockets.Tests/TcpSocketClientTests.cs
@@ -236,5 +236,17 @@ public async Task TcpSocketClient_ShouldBeAbleToDisconnectThenReconnect(int buff
             await listener.StopListeningAsync();
         }
 
+        [Fact]
+        public Task TcpSocketClient_Connect_ShouldCancelByCancellationToken()
+        {
+            var sut = new TcpSocketClient();
+
+            var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+            var ct = cts.Token;
+
+            // let's just hope no one's home :)
+            return Assert.ThrowsAsync<OperationCanceledException>(()=> sut.ConnectAsync("99.99.99.99", 51234, cancellationToken: cts.Token));
+        }
+
     }
 }