better websocker unmask

Author: honfika

examples/Titanium.Web.Proxy.Examples.Basic/ProxyTestController.cs

@@ -152,22 +152,34 @@ private async Task onBeforeTunnelConnectRequest(object sender, TunnelConnectSess
             }
         }
 
+        private void WebSocket_DataSent(object sender, DataEventArgs e)
+        {
+            var args = (SessionEventArgs)sender;
+            WebSocketDataSentReceived(args, e, true);
+        }
+
         private void WebSocket_DataReceived(object sender, DataEventArgs e)
         {
             var args = (SessionEventArgs)sender;
+            WebSocketDataSentReceived(args, e, false);
+        }
+
+        private void WebSocketDataSentReceived(SessionEventArgs args, DataEventArgs e, bool sent)
+        {
+            var color = sent ? ConsoleColor.Green : ConsoleColor.Blue;
 
             foreach (var frame in args.WebSocketDecoder.Decode(e.Buffer, e.Offset, e.Count))
             {
                 if (frame.OpCode == WebsocketOpCode.Binary)
                 {
                     var data = frame.Data.ToArray();
                     string str = string.Join(",", data.ToArray().Select(x => x.ToString("X2")));
-                    writeToConsole(str, ConsoleColor.Blue).Wait();
+                    writeToConsole(str, color).Wait();
                 }
 
                 if (frame.OpCode == WebsocketOpCode.Text)
                 {
-                    writeToConsole(frame.GetText(), ConsoleColor.Blue).Wait();
+                    writeToConsole(frame.GetText(), color).Wait();
                 }
             }
         }
@@ -233,6 +245,7 @@ private async Task onResponse(object sender, SessionEventArgs e)
         {
             if (e.HttpClient.ConnectRequest?.TunnelType == TunnelType.Websocket)
             {
+                e.DataSent += WebSocket_DataSent;
                 e.DataReceived += WebSocket_DataReceived;
             }
 

src/Titanium.Web.Proxy/WebSocketDecoder.cs

@@ -1,6 +1,7 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Runtime.InteropServices;
 using Titanium.Web.Proxy.StreamExtended.BufferPool;
 
 namespace Titanium.Web.Proxy
@@ -60,25 +61,51 @@ public IEnumerable<WebSocketFrame> Decode(byte[] data, int offset, int count)
                     }
                 }
 
-                uint mask = 0;
-                if (masked)
+                if (size < 0)
                 {
-                    //mask = (uint)(((long)data1[idx++] << 24) + (data1[idx++] << 16) + (data1[idx++] << 8) + data1[idx++]);
-                    mask = (uint)(data1[idx++] + (data1[idx++] << 8) + (data1[idx++] << 16) + ((long)data1[idx++] << 24));
+                    ;
+                }
+
+                if (data1.Length < idx + size)
+                {
+                    break;
                 }
 
                 if (masked)
                 {
-                    uint m = mask;
-                    for (int i = 0; i < size; i++)
-                    {
-                        data[i + idx] = (byte)(data1[i + idx] ^ (byte)mask);
+                    //mask = (uint)(((long)data1[idx++] << 24) + (data1[idx++] << 16) + (data1[idx++] << 8) + data1[idx++]);
+                    //mask = (uint)(data1[idx++] + (data1[idx++] << 8) + (data1[idx++] << 16) + ((long)data1[idx++] << 24));
+                    var uData = MemoryMarshal.Cast<byte, uint>(data1.Slice(idx, (int)size + 4));
+                    idx += 4;
 
-                        m >>= 8;
+                    uint mask = uData[0];
+                    long size1 = size;
+                    if (size > 4)
+                    {
+                        uData = uData.Slice(1);
+                        for (int i = 0; i < uData.Length; i++)
+                        {
+                            uData[i] = uData[i] ^ mask;
+                        }
 
-                        if (m == 0)
-                            m = mask;
+                        size1 -= uData.Length * 4;
                     }
+
+                    if (size1 > 0)
+                    {
+                        int pos = (int)(idx + size - size1);
+                        data1[pos] ^= (byte)mask;
+
+                        if (size1 > 1)
+                        {
+                            data1[pos + 1] ^= (byte)(mask >> 8);
+                        }
+
+                        if (size1 > 2)
+                        {
+                            data1[pos + 2] ^= (byte)(mask >> 16);
+                        }
+;                    }
                 }
 
                 var frameData = buffer.Slice(idx, (int)size);