Merge pull request #106 from trew/socket-exception

Fix exception in SocketWrapper when WebSocketServer.Dispose() is called.

src/Fleck.Tests/SocketWrapperTests.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Threading.Tasks;
 using NUnit.Framework;
 using System.Net.Sockets;
 using System.Net;
@@ -7,6 +8,35 @@
 
 namespace Fleck.Tests
 {
+    [TestFixture]
+    public class SocketWrapperDisposeTest
+    {
+        private Socket _client;
+        private EndPoint _endpoint;
+        private SocketWrapper _wrapper;
+
+        [SetUp]
+        public void Setup()
+        {
+            _endpoint = new IPEndPoint(IPAddress.Loopback, 45982);
+            _client = new Socket(_endpoint.AddressFamily, SocketType.Stream, ProtocolType.IP);
+
+            _wrapper = new SocketWrapper(_client);
+            _wrapper.Bind(_endpoint);
+            _wrapper.Listen(100);
+        }
+
+        [Test]
+        public void ShouldCompleteAcceptTaskOnDispose()
+        {
+            Task task = _wrapper.Accept(socket => { }, exception => { });
+            _wrapper.Dispose();
+
+            Assert.DoesNotThrow(task.Wait);
+            Assert.IsTrue(task.IsCompleted);
+        }
+    }
+
     [TestFixture]
     public class SocketWrapperTests
     {

src/Fleck/SocketWrapper.cs

@@ -108,10 +108,7 @@ 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 => {
-                _tokenSource.Token.ThrowIfCancellationRequested();
-                return new SocketWrapper(_socket.EndAccept(r));
-            };
+            Func<IAsyncResult, ISocket> end = r => _tokenSource.Token.IsCancellationRequested ? null : 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);

src/Fleck/WebSocketServer.cs

@@ -85,6 +85,8 @@ private void ListenForClients()
 
         private void OnClientConnect(ISocket clientSocket)
         {
+            if (clientSocket == null) return; // socket closed
+
             FleckLog.Debug(String.Format("Client connected from {0}:{1}", clientSocket.RemoteIpAddress, clientSocket.RemotePort.ToString()));
             ListenForClients();