From 534ccfdaaed1a61fb0f741e879e70d1d09abd6c6 Mon Sep 17 00:00:00 2001
From: Samuel Andersson <smulster@gmail.com>
Date: Tue, 26 Aug 2014 10:21:09 +0200
Subject: [PATCH 1/2] Don't throw exception when socket is closed.

---
 src/Fleck/SocketWrapper.cs   | 5 +----
 src/Fleck/WebSocketServer.cs | 2 ++
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/src/Fleck/SocketWrapper.cs b/src/Fleck/SocketWrapper.cs
index feeae668..ff907f8b 100644
--- a/src/Fleck/SocketWrapper.cs
+++ b/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);
diff --git a/src/Fleck/WebSocketServer.cs b/src/Fleck/WebSocketServer.cs
index 013e3d1d..f71b25ac 100644
--- a/src/Fleck/WebSocketServer.cs
+++ b/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();
 

From d025c1cc4401fa863fb179c50e40d6d55e416bb3 Mon Sep 17 00:00:00 2001
From: Samuel Andersson <smulster@gmail.com>
Date: Mon, 1 Sep 2014 10:01:57 +0200
Subject: [PATCH 2/2] Add test to ensure task completes correctly when wrapper
 is disposed.

---
 src/Fleck.Tests/SocketWrapperTests.cs | 30 +++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/src/Fleck.Tests/SocketWrapperTests.cs b/src/Fleck.Tests/SocketWrapperTests.cs
index 2adead77..990e01b5 100644
--- a/src/Fleck.Tests/SocketWrapperTests.cs
+++ b/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
     {