Cancel Accept on Dispose(). Fixes #36

statianzo · Jan 31, 2012 · 2d90886 · 2d90886
1 parent 535e51b
commit 2d90886
Showing 1 changed file with 14 additions and 4 deletions.
diff --git a/src/Fleck/SocketWrapper.cs b/src/Fleck/SocketWrapper.cs
@@ -6,14 +6,17 @@
 using System.Security.Authentication;
 using System.Security.Cryptography.X509Certificates;
 using System.Threading.Tasks;
+using System.Threading;
 
 namespace Fleck
 {
     public class SocketWrapper : ISocket
     {
         private readonly Socket _socket;
         private Stream _stream;
-
+        private CancellationTokenSource _tokenSource;
+        private TaskFactory _taskFactory;
+
         public string RemoteIpAddress
         {
             get
@@ -25,6 +28,8 @@ public string RemoteIpAddress
 
         public SocketWrapper(Socket socket)
         {
+            _tokenSource = new CancellationTokenSource();
+            _taskFactory = new TaskFactory(_tokenSource.Token);
             _socket = socket;
             if (_socket.Connected)
                 _stream = new NetworkStream(_socket);
@@ -79,22 +84,27 @@ public Task<int> Receive(byte[] buffer, Action<int> callback, Action<Exception>
 
         public Task<ISocket> Accept(Action<ISocket> callback, Action<Exception> error)
         {
-            Func<IAsyncResult, ISocket> end = r => new SocketWrapper(_socket.EndAccept(r));
-            var task = Task.Factory.FromAsync(_socket.BeginAccept, end, null);
-            task.ContinueWith(t => callback(t.Result), TaskContinuationOptions.NotOnFaulted)
+            Func<IAsyncResult, ISocket> end = r => {
+                _tokenSource.Token.ThrowIfCancellationRequested();
+                return new SocketWrapper(_socket.EndAccept(r));
+            };
+            var task = _taskFactory.FromAsync(_socket.BeginAccept, end, null);
+            task.ContinueWith(t => callback(t.Result), TaskContinuationOptions.OnlyOnRanToCompletion)
                 .ContinueWith(t => error(t.Exception), TaskContinuationOptions.OnlyOnFaulted);
             task.ContinueWith(t => error(t.Exception), TaskContinuationOptions.OnlyOnFaulted);
             return task;
         }
 
         public void Dispose()
         {
+            _tokenSource.Cancel();
             if (_stream != null) _stream.Dispose();
             if (_socket != null) _socket.Dispose();
         }
 
         public void Close()
         {
+            _tokenSource.Cancel();
             if (_stream != null) _stream.Close();
             if (_socket != null) _socket.Close();
         }