diff --git a/src/Owin.WebSocket/Handlers/NetWebSocket.cs b/src/Owin.WebSocket/Handlers/NetWebSocket.cs index e0ca2b3..9c86162 100644 --- a/src/Owin.WebSocket/Handlers/NetWebSocket.cs +++ b/src/Owin.WebSocket/Handlers/NetWebSocket.cs @@ -1,4 +1,5 @@ using System; +using System.IO; using System.Net.WebSockets; using System.Threading; using System.Threading.Tasks; @@ -56,7 +57,10 @@ public Task Send(ArraySegment data, WebSocketMessageType messageType, bool public Task Close(WebSocketCloseStatus closeStatus, string closeDescription, CancellationToken cancelToken) { - return mWebSocket.CloseAsync(closeStatus, closeDescription, cancelToken); + using (mWebSocket) + { + return mWebSocket.CloseAsync(closeStatus, closeDescription, cancelToken); + } } public async Task, WebSocketMessageType>> ReceiveMessage(byte[] buffer, CancellationToken cancelToken) @@ -65,6 +69,10 @@ public async Task, WebSocketMessageType>> ReceiveMessag WebSocketReceiveResult result; do { + if (count == buffer.Length) + throw new InternalBufferOverflowException( + "The Buffer is to small to get the Websocket Message! Increase in the Constructor!"); + var segment = new ArraySegment(buffer, count, buffer.Length - count); result = await mWebSocket.ReceiveAsync(segment, cancelToken);