Some graceful handling for client disconnects. Fixes statianzo#2

joeyespo · Feb 11, 2011 · 5f82140 · 5f82140
1 parent 774fedc
commit 5f82140
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 28 deletions.
diff --git a/src/Fleck/Receiver.cs b/src/Fleck/Receiver.cs
@@ -33,32 +33,50 @@ public void Receive(DataFrame frame = null)
 				return;
 			}
 
-			Socket.BeginReceive(buffer, 0, buffer.Length, SocketFlags.None,
-				r =>
-				{
-					int size = Socket.EndReceive(r);
-					var dataframe = frame;
-
-					if (size <= 0)
+			try
+			{
+				Socket.BeginReceive(buffer, 0, buffer.Length, SocketFlags.None,
+					r =>
 					{
-						_connection.Close();
-						return;
-					}
+						try
+						{
+							int size = Socket.EndReceive(r);
+							var dataframe = frame;
 
-					for (int i = 0; i < size; i++)
-						_queue.Enqueue(buffer[i]);
+							if (size <= 0)
+							{
+								_connection.Close();
+								return;
+							}
 
-					while (_queue.Count > 0)
-					{
-						dataframe.Append(_queue.Dequeue());
-						if (!dataframe.IsComplete) continue;
+							for (int i = 0; i < size; i++)
+								_queue.Enqueue(buffer[i]);
+
+							while (_queue.Count > 0)
+							{
+								dataframe.Append(_queue.Dequeue());
+								if (!dataframe.IsComplete) continue;
+
+								var data = dataframe.ToString();
+								_connection.OnMessage(data);
+								dataframe = new DataFrame();
+							}
+							Receive(dataframe);
 
-						var data = dataframe.ToString();
-						_connection.OnMessage(data);
-						dataframe = new DataFrame();
-					}
-					Receive(dataframe);
-				}, null);
+						}
+						catch (SocketException e)
+						{
+
+							Log.Error(e.Message);
+							_connection.Close();
+						}
+					}, null);
+			}
+			catch (SocketException e)
+			{
+				Log.Error(e.Message);
+				_connection.Close();
+			}
 		}
 	}
 }
diff --git a/src/Fleck/Sender.cs b/src/Fleck/Sender.cs
@@ -17,11 +17,27 @@ public void Send(string data)
 		{
 			if (!Socket.Connected) return;
 			var wrapped = DataFrame.Wrap(data);
-			Socket.BeginSend(wrapped, 0, wrapped.Length, SocketFlags.None, r =>
-				{
-					if (Socket.Connected)
-						Socket.EndSend(r);
-				}, null);
+			try
+			{
+				Socket.BeginSend(wrapped, 0, wrapped.Length, SocketFlags.None, r =>
+					{
+						try
+						{
+							if (Socket.Connected)
+								Socket.EndSend(r);
+						}
+						catch (SocketException e)
+						{
+							Log.Error(e.Message);
+							_connection.Close();
+						}
+					}, null);
+			}
+			catch (SocketException e)
+			{
+				Log.Error(e.Message);
+				_connection.Close();
+			}
 		}
 	}
 }
diff --git a/src/Fleck/WebSocketConnection.cs b/src/Fleck/WebSocketConnection.cs
@@ -37,7 +37,6 @@ public void StartReceiving()
 
 		public void Close()
 		{
-			if (!Socket.Connected) return;
 			OnClose();
 			Socket.Close();
 		}