Pull request for issue #225

src/Fleck.Tests/SocketWrapperTests.cs

@@ -29,7 +29,7 @@ public void Setup()
         [Test]
         public void ShouldCompleteAcceptTaskOnDispose()
         {
-            Task task = _wrapper.Accept(socket => { }, exception => { });
+            Task task = _wrapper.Accept(socket => { }, () => { }, exception => { });
             _wrapper.Dispose();
 
             Assert.DoesNotThrow(task.Wait);

src/Fleck.Tests/WebSocketServerTests.cs

@@ -43,7 +43,7 @@ public void ShouldStart()
             _server.Start(connection => { });
 
             socketMock.Verify(s => s.Bind(It.Is<IPEndPoint>(i => i.Port == 8000)));
-            socketMock.Verify(s => s.Accept(It.IsAny<Action<ISocket>>(), It.IsAny<Action<Exception>>()));
+            socketMock.Verify(s => s.Accept(It.IsAny<Action<ISocket>>(), It.IsAny<Action>(), It.IsAny<Action<Exception>>()));
         }
 
         [Test]

src/Fleck/Interfaces/ISocket.cs

@@ -17,7 +17,7 @@ public interface ISocket
         bool NoDelay { get; set; }
         EndPoint LocalEndPoint { get; }
 
-        Task<ISocket> Accept(Action<ISocket> callback, Action<Exception> error);
+        Task<ISocket> Accept(Action<ISocket> callback, Action setHandle, Action<Exception> error);
         Task Send(byte[] buffer, Action callback, Action<Exception> error);
         Task<int> Receive(byte[] buffer, Action<int> callback, Action<Exception> error, int offset = 0);
         Task Authenticate(X509Certificate2 certificate, SslProtocols enabledSslProtocols, Action callback, Action<Exception> error);

src/Fleck/Interfaces/IWebSocketServer.cs

@@ -5,5 +5,6 @@ namespace Fleck
     public interface IWebSocketServer : IDisposable
     {
         void Start(Action<IWebSocketConnection> config);
+        void Stop();
     }
-}
+}

src/Fleck/SocketWrapper.cs

@@ -130,11 +130,11 @@ public Task<int> Receive(byte[] buffer, Action<int> callback, Action<Exception>
             }
         }
 
-        public Task<ISocket> Accept(Action<ISocket> callback, Action<Exception> error)
+        public Task<ISocket> Accept(Action<ISocket> callback, Action setHandle,  Action<Exception> error)
         {
             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)
+            task.ContinueWith(t => { setHandle(); callback(t.Result); }, TaskContinuationOptions.OnlyOnRanToCompletion)
                 .ContinueWith(t => error(t.Exception), TaskContinuationOptions.OnlyOnFaulted);
             task.ContinueWith(t => error(t.Exception), TaskContinuationOptions.OnlyOnFaulted);
             return task;

src/Fleck/WebSocketServer.cs

@@ -4,6 +4,7 @@
 using System.Security.Cryptography.X509Certificates;
 using System.Collections.Generic;
 using System.Security.Authentication;
+using System.Threading.Tasks;
 using Fleck.Helpers;
 
 namespace Fleck
@@ -13,6 +14,7 @@ public class WebSocketServer : IWebSocketServer
         private readonly string _scheme;
         private readonly IPAddress _locationIP;
         private Action<IWebSocketConnection> _config;
+        private bool _running;
 
         public WebSocketServer(string location)
         {
@@ -38,7 +40,6 @@ public WebSocketServer(string location)
         public X509Certificate2 Certificate { get; set; }
         public SslProtocols EnabledSslProtocols { get; set; }
         public IEnumerable<string> SupportedSubProtocols { get; set; }
-        public bool RestartAfterListenError {get; set; }
 
         public bool IsSecure
         {
@@ -47,6 +48,7 @@ public bool IsSecure
 
         public void Dispose()
         {
+	    Stop();
             ListenerSocket.Dispose();
         }
 
@@ -89,38 +91,42 @@ public void Start(Action<IWebSocketConnection> config)
                     FleckLog.Debug("Using default TLS 1.0 security protocol.");
                 }
             }
+
+            _running = true;
             ListenForClients();
             _config = config;
         }
+
+         public void Stop() 
+         {
+             _running = false;
+         }
 
         private void ListenForClients()
         {
-            ListenerSocket.Accept(OnClientConnect, e => {
-                FleckLog.Error("Listener socket is closed", e);
-                if(RestartAfterListenError){
-                    FleckLog.Info("Listener socket restarting");
-                    try
-                    {
-                        ListenerSocket.Dispose();
-                        var socket = new Socket(_locationIP.AddressFamily, SocketType.Stream, ProtocolType.IP);
-                        ListenerSocket = new SocketWrapper(socket);
-                        Start(_config);
-                        FleckLog.Info("Listener socket restarted");
-                    }
-                    catch (Exception ex)
-                    {
-                        FleckLog.Error("Listener could not be restarted", ex);
-                    }
-                }
+	    System.Threading.ManualResetEvent acceptDone = new System.Threading.ManualResetEvent (false);
+
+            Task.Run( () => {
+
+                while (_running) {
+                    acceptDone.Reset ();
+
+                    ListenerSocket.Accept( OnClientConnect, () => acceptDone.Set(),
+                        e =>  FleckLog.Error ("An error occurred while accepting a client connection", e)  );
+
+                    acceptDone.WaitOne ();
+               }
+
             });
+
         }
 
         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();
+            //ListenForClients();
 
             WebSocketConnection connection = null;