Modified processFragments method in WebSocket.cs

websocket-sharp/Ext.cs

@@ -71,10 +71,8 @@ namespace WebSocketSharp {
         return value;
 
       using (var input = new MemoryStream(value))
-      using (var output = input.compress())
       {
-        output.Close();
-        return output.ToArray();
+        return input.compressToArray();
       }
     }
 
@@ -84,19 +82,10 @@ namespace WebSocketSharp {
       if (stream.Length == 0)
         return output;
 
+      stream.Position = 0;
       using (var ds = new DeflateStream(output, CompressionMode.Compress, true))
       {
-        int readLen = 0;
-        var buffer = new byte[256];
-        while (true)
-        {
-          readLen = stream.Read(buffer, 0, buffer.Length);
-          if (readLen == 0)
-            break;
-
-          ds.Write(buffer, 0, readLen);
-        }
-
+        stream.CopyTo(ds);
         ds.Close(); // "BFINAL" set to 1.
         output.Position = 0;
 
@@ -104,16 +93,23 @@ namespace WebSocketSharp {
       }
     }
 
+    private static byte[] compressToArray(this Stream stream)
+    {
+      using (var comp = stream.compress())
+      {
+        comp.Close();
+        return comp.ToArray();
+      }
+    }
+
     private static byte[] decompress(this byte[] value)
     {
       if (value.LongLength == 0)
         return value;
 
       using (var input = new MemoryStream(value))
-      using (var output = input.decompress())
       {
-        output.Close();
-        return output.ToArray();
+        return input.decompressToArray();
       }
     }
 
@@ -123,24 +119,23 @@ namespace WebSocketSharp {
       if (stream.Length == 0)
         return output;
 
+      stream.Position = 0;
       using (var ds = new DeflateStream(stream, CompressionMode.Decompress, true))
       {
-        int readLen = 0;
-        var buffer = new byte[256];
-        while (true)
-        {
-          readLen = ds.Read(buffer, 0, buffer.Length);
-          if (readLen == 0)
-            break;
-
-          output.Write(buffer, 0, readLen);
-        }
-
-        output.Position = 0;
+        ds.CopyTo(output, true);
         return output;
       }
     }
 
+    private static byte[] decompressToArray(this Stream stream)
+    {
+      using (var decomp = stream.decompress())
+      {
+        decomp.Close();
+        return decomp.ToArray();
+      }
+    }
+
     private static void times(this ulong n, Action act)
     {
       for (ulong i = 0; i < n; i++)
@@ -165,6 +160,32 @@ namespace WebSocketSharp {
              : stream;
     }
 
+    internal static byte[] CompressToArray(this Stream stream, CompressionMethod method)
+    {
+      return method == CompressionMethod.DEFLATE
+             ? stream.compressToArray()
+             : stream.ToByteArray();
+    }
+
+    internal static void CopyTo(this Stream src, Stream dest)
+    {
+      src.CopyTo(dest, false);
+    }
+
+    internal static void CopyTo(this Stream src, Stream dest, bool setDefaultPosition)
+    {
+      int readLen;
+      int bufferLen = 256;
+      var buffer = new byte[bufferLen];
+      while ((readLen = src.Read(buffer, 0, bufferLen)) > 0)
+      {
+        dest.Write(buffer, 0, readLen);
+      }
+
+      if (setDefaultPosition)
+        dest.Position = 0;
+    }
+
     internal static byte[] Decompress(this byte[] value, CompressionMethod method)
     {
       return method == CompressionMethod.DEFLATE
@@ -179,6 +200,13 @@ namespace WebSocketSharp {
              : stream;
     }
 
+    internal static byte[] DecompressToArray(this Stream stream, CompressionMethod method)
+    {
+      return method == CompressionMethod.DEFLATE
+             ? stream.decompressToArray()
+             : stream.ToByteArray();
+    }
+
     internal static string GetNameInternal(this string nameAndValue, string separator)
     {
       int i = nameAndValue.IndexOf(separator);
@@ -234,6 +262,23 @@ namespace WebSocketSharp {
              : String.Format("\"{0}\"", value.Replace("\"", "\\\""));
     }
 
+    internal static string RemovePrefix(this string value, params string[] prefixes)
+    {
+      int i = 0;
+      foreach (var prefix in prefixes)
+      {
+        if (value.StartsWith(prefix))
+        {
+          i = prefix.Length;
+          break;
+        }
+      }
+
+      return i > 0
+             ? value.Substring(i)
+             : value;
+    }
+
     internal static IEnumerable<string> SplitHeaderValue(this string value, params char[] separator)
     {
       var separators = new string(separator);
@@ -275,6 +320,26 @@ namespace WebSocketSharp {
         yield return buffer.ToString();
     }
 
+    internal static byte[] ToByteArray(this Stream stream)
+    {
+      using (var output = new MemoryStream())
+      {
+        stream.Position = 0;
+        stream.CopyTo(output);
+        output.Close();
+
+        return output.ToArray();
+      }
+    }
+
+    internal static void WriteBytes(this Stream stream, byte[] value)
+    {
+      using (var input = new MemoryStream(value))
+      {
+        input.CopyTo(stream);
+      }
+    }
+
     #endregion
 
     #region Public Methods

websocket-sharp/WebSocket.cs

@@ -36,7 +36,6 @@ using System.Collections.Generic;
 using System.Collections.Specialized;
 using System.Diagnostics;
 using System.IO;
-using System.IO.Compression;
 using System.Linq;
 using System.Net.Sockets;
 using System.Security.Cryptography;
@@ -770,6 +769,59 @@ namespace WebSocketSharp {
                  response.HeaderExists("Sec-WebSocket-Version", _version);
     }
 
+    private bool mergeFragments(Stream dest)
+    {
+      Func<WsFrame, bool> processContinuation = contFrame =>
+      {
+        if (!contFrame.IsContinuation)
+          return false;
+
+        if (contFrame.IsCompressed)
+          contFrame.Decompress(_compression);
+
+        dest.WriteBytes(contFrame.PayloadData.ToByteArray());
+        return true;
+      };
+
+      while (true)
+      {
+        var frame = readFrame();
+        if (processAbnormal(frame))
+          return false;
+
+        if (!frame.IsFinal)
+        {
+          // MORE & CONT
+          if (processContinuation(frame))
+            continue;
+        }
+        else
+        {
+          // FINAL & CONT
+          if (processContinuation(frame))
+            break;
+
+          // FINAL & PING
+          if (processPing(frame))
+            continue;
+
+          // FINAL & PONG
+          if (processPong(frame))
+            continue;
+
+          // FINAL & CLOSE
+          if (processClose(frame))
+            return false;
+        }
+
+        // ?
+        processIncorrectFrame();
+        return false;
+      }
+
+      return true;
+    }
+
     private void onClose(CloseEventArgs eventArgs)
     {
       if (!Thread.CurrentThread.IsBackground)
@@ -897,64 +949,28 @@ namespace WebSocketSharp {
 
     private void processFragments(WsFrame first)
     {
-      bool compressed = first.IsCompressed;
-      if (compressed && _perFrameCompress)
-        first.Decompress(_compression);
-
-      var buffer = new List<byte>(first.PayloadData.ToByteArray());
-      Func<WsFrame, bool> processContinuation = contFrame =>
+      using (var merge = new MemoryStream())
       {
-        if (!contFrame.IsContinuation)
-          return false;
-
-        if (contFrame.IsCompressed)
-          contFrame.Decompress(_compression);
+        if (first.IsCompressed && _perFrameCompress)
+          first.Decompress(_compression);
 
-        buffer.AddRange(contFrame.PayloadData.ToByteArray());
-        return true;
-      };
- 
-      while (true)
-      {
-        var frame = readFrame();
-        if (processAbnormal(frame))
+        merge.WriteBytes(first.PayloadData.ToByteArray());
+        if (!mergeFragments(merge))
           return;
 
-        if (!frame.IsFinal)
+        byte[] data;
+        if (_compression != CompressionMethod.NONE && !_perFrameCompress)
         {
-          // MORE & CONT
-          if (processContinuation(frame))
-            continue;
+          data = merge.DecompressToArray(_compression);
         }
         else
         {
-          // FINAL & CONT
-          if (processContinuation(frame))
-            break;
-
-          // FINAL & PING
-          if (processPing(frame))
-            continue;
-
-          // FINAL & PONG
-          if (processPong(frame))
-            continue;
-
-          // FINAL & CLOSE
-          if (processClose(frame))
-            return;
+          merge.Close();
+          data = merge.ToArray();
         }
 
-        // ?
-        processIncorrectFrame();
-        return;
+        onMessage(new MessageEventArgs(first.Opcode, new PayloadData(data)));
       }
-
-      var data = compressed && !_perFrameCompress
-               ? buffer.ToArray().Decompress(_compression)
-               : buffer.ToArray();
-
-      onMessage(new MessageEventArgs(first.Opcode, new PayloadData(data)));
     }
 
     private void processFrame(WsFrame frame)
@@ -1015,10 +1031,7 @@ namespace WebSocketSharp {
       foreach (var extension in extensions.SplitHeaderValue(','))
       {
         var e = extension.Trim();
-        var tmp = e.StartsWith("x-webkit-")
-                ? e.Substring(9)
-                : e;
-
+        var tmp = e.RemovePrefix("x-webkit-");
         if (!compress)
         {
           if (isPerFrameCompressExtension(tmp))