Fix for issue #3

websocket-sharp/Ext.cs

@@ -112,24 +112,24 @@ namespace WebSocketSharp
     public static byte[] ReadBytes<TStream>(this TStream stream, ulong length, int bufferLength)
       where TStream : System.IO.Stream
     {
-      List<byte> readData = new List<byte>();
-
       ulong count     = length / (ulong)bufferLength;
       int   remainder = (int)(length % (ulong)bufferLength);
 
-      byte[] buffer1 = new byte[bufferLength];
+      List<byte> readData = new List<byte>();
+      byte[]     buffer1  = new byte[bufferLength];
+      int        readLen  = 0;
 
       count.Times(() =>
       {
-        stream.Read(buffer1, 0, bufferLength);
-        readData.AddRange(buffer1);
+        readLen = stream.Read(buffer1, 0, bufferLength);
+        if (readLen > 0) readData.AddRange(buffer1.SubArray(0, readLen));
       });
 
       if (remainder > 0)
       {
         byte[] buffer2 = new byte[remainder];
-        stream.Read(buffer2, 0, remainder);
-        readData.AddRange(buffer2);
+        readLen = stream.Read(buffer2, 0, remainder);
+        if (readLen > 0) readData.AddRange(buffer2.SubArray(0, readLen));
       }
 
       return readData.ToArray();

websocket-sharp/Frame/WsFrame.cs

@@ -179,9 +179,14 @@ namespace WebSocketSharp.Frame
       int    buffer2Len    = 0;
       ulong  buffer3Len    = 0;
       int    maskingKeyLen = 4;
+      int    readLen       = 0;
 
       buffer1 = new byte[buffer1Len];
-      stream.Read(buffer1, 0, buffer1Len);
+      readLen = stream.Read(buffer1, 0, buffer1Len);
+      if (readLen < buffer1Len)
+      {
+        return null;
+      }
 
       // FIN
       fin = (buffer1[0] & 0x80) == 0x80
@@ -224,7 +229,13 @@ namespace WebSocketSharp.Frame
       if (buffer2Len > 0)
       {
         buffer2 = new byte[buffer2Len];
-        stream.Read(buffer2, 0, buffer2Len);
+        readLen = stream.Read(buffer2, 0, buffer2Len);
+
+        if (readLen < buffer2Len)
+        {
+          return null;
+        }
+
         extPayloadLen = buffer2;
         switch (buffer2Len)
         {
@@ -245,17 +256,32 @@ namespace WebSocketSharp.Frame
       if (masked == Mask.MASK)
       {
         maskingKey = new byte[maskingKeyLen];
-        stream.Read(maskingKey, 0, maskingKeyLen);
+        readLen    = stream.Read(maskingKey, 0, maskingKeyLen);
+
+        if (readLen < maskingKeyLen)
+        {
+          return null;
+        }
       }
       // Payload Data
       if (buffer3Len <= (ulong)_readBufferLen)
       {
         buffer3 = new byte[buffer3Len];
-        stream.Read(buffer3, 0, (int)buffer3Len);
+        readLen = stream.Read(buffer3, 0, (int)buffer3Len);
+
+        if (readLen < (int)buffer3Len)
+        {
+          return null;
+        }
       }
       else
       {
         buffer3 = stream.ReadBytes(buffer3Len, _readBufferLen);
+
+        if ((ulong)buffer3.LongLength < buffer3Len)
+        {
+          return null;
+        }
       }
 
       if (masked == Mask.MASK)

websocket-sharp/Server/WebSocketService.cs

@@ -64,10 +64,10 @@ namespace WebSocketSharp.Server
         _server.AddService(this);
       };
 
-      _socket.OnClose += (sender, e) =>
-      {
-        _server.RemoveService(this);
-      };
+//      _socket.OnClose += (sender, e) =>
+//      {
+//        _server.RemoveService(this);
+//      };
     }
 
     #endregion

websocket-sharp/WebSocket.cs

@@ -752,6 +752,10 @@ namespace WebSocketSharp
       {
         close(CloseStatusCode.TOO_BIG, ex.Message);
       }
+      catch (Exception ex)
+      {
+        close(CloseStatusCode.ABNORMAL, ex.Message);
+      }
     }
 
     private void messageLoop()
@@ -762,39 +766,16 @@ namespace WebSocketSharp
       }
     }
 
-    private void startMessageThread()
+    private void pong(PayloadData data)
     {
-      if (_isClient)
-      {
-        _msgThread = new Thread(new ThreadStart(messageLoop)); 
-        _msgThread.IsBackground = true;
-        _msgThread.Start();
-      }
-      else
-      {
-        _autoEvent = new AutoResetEvent(false);
-        Action act = () =>
-        {
-          if (_readyState == WsState.OPEN)
-          {
-            message();
-          }
-        };
-        AsyncCallback callback = (ar) =>
-        {
-          act.EndInvoke(ar);
+      var frame = createFrame(Fin.FINAL, Opcode.PONG, data);
+      send(frame);
+    }
 
-          if (_readyState == WsState.OPEN)
-          {
-            act.BeginInvoke(callback, null);
-          }
-          else
-          {
-            _autoEvent.Set();
-          }
-        };
-        act.BeginInvoke(callback, null);
-      }
+    private void pong(string data)
+    {
+      var payloadData = new PayloadData(data);
+      pong(payloadData);
     }
 
     private MessageEventArgs receive()
@@ -806,7 +787,15 @@ namespace WebSocketSharp
       Opcode           opcode;
       PayloadData      payloadData;
 
+      Action act = () =>
+      {
+        var msg = "WebSocket data frame can not be read from network stream.";
+        close(CloseStatusCode.ABNORMAL, msg);
+      };
+
       frame = _wsStream.ReadFrame();
+      if (frame.IsNullDo(act)) return null;
+
       if ((frame.Fin == Fin.FINAL && frame.Opcode == Opcode.CONT) ||
           (frame.Fin == Fin.MORE  && frame.Opcode == Opcode.CONT))
       {
@@ -825,6 +814,7 @@ namespace WebSocketSharp
           while (true)
           {
             frame = _wsStream.ReadFrame();
+            if (frame.IsNullDo(act)) return null;
 
             if (frame.Fin == Fin.MORE)
             {
@@ -856,8 +846,10 @@ namespace WebSocketSharp
             }
             else if (frame.Opcode == Opcode.PING)
             {// FINAL & PING
+              #if DEBUG
+              Console.WriteLine("WS: Info@receive: Return Pong.");
+              #endif
               pong(frame.PayloadData);
-              OnMessage.Emit(this, new MessageEventArgs(frame.Opcode, frame.PayloadData));
             }
             else if (frame.Opcode == Opcode.PONG)
             {// FINAL & PONG
@@ -889,8 +881,11 @@ namespace WebSocketSharp
               close(payloadData);
               break;
             case Opcode.PING:
+              #if DEBUG
+              Console.WriteLine("WS: Info@receive: Return Pong.");
+              #endif
               pong(payloadData);
-              goto default;
+              break;
             default:
               eventArgs = new MessageEventArgs(opcode, payloadData);
               break;
@@ -902,18 +897,6 @@ namespace WebSocketSharp
       return eventArgs;
     }
 
-    private void pong(PayloadData data)
-    {
-      var frame = createFrame(Fin.FINAL, Opcode.PONG, data);
-      send(frame);
-    }
-
-    private void pong(string data)
-    {
-      var payloadData = new PayloadData(data);
-      pong(payloadData);
-    }
-
     private string[] receiveOpeningHandshake()
     {
       var readData = new List<byte>();
@@ -1096,10 +1079,60 @@ namespace WebSocketSharp
       _wsStream.Write(buffer, 0, buffer.Length);
     }
 
+    private void startMessageThread()
+    {
+      if (_isClient)
+      {
+        _msgThread = new Thread(new ThreadStart(messageLoop)); 
+        _msgThread.IsBackground = true;
+        _msgThread.Start();
+      }
+      else
+      {
+        _autoEvent = new AutoResetEvent(false);
+        Action act = () =>
+        {
+          if (_readyState == WsState.OPEN)
+          {
+            message();
+          }
+        };
+        AsyncCallback callback = (ar) =>
+        {
+          act.EndInvoke(ar);
+
+          if (_readyState == WsState.OPEN)
+          {
+            act.BeginInvoke(callback, null);
+          }
+          else
+          {
+            _autoEvent.Set();
+          }
+        };
+        act.BeginInvoke(callback, null);
+      }
+    }
+
     #endregion
 
     #region Public Methods
 
+    public void Close()
+    {
+      Close(CloseStatusCode.NORMAL);
+    }
+
+    public void Close(CloseStatusCode code)
+    {
+      Close(code, String.Empty);
+    }
+
+    public void Close(CloseStatusCode code, string reason)
+    {
+      close(code, reason);
+    }
+
     public void Connect()
     {
       if (_readyState == WsState.OPEN)
@@ -1128,21 +1161,6 @@ namespace WebSocketSharp
       }
     }
 
-    public void Close()
-    {
-      Close(CloseStatusCode.NORMAL);
-    }
-
-    public void Close(CloseStatusCode code)
-    {
-      Close(code, String.Empty);
-    }
-
-    public void Close(CloseStatusCode code, string reason)
-    {
-      close(code, reason);
-    }
-
     public void Dispose()
     {
       Close(CloseStatusCode.AWAY);