Added improved WebSocket support

XMAT/Engines/WebServiceProxy/Logic/Proxy/BaseWebSocketMessage.cs

@@ -0,0 +1,49 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System;
+using System.ComponentModel;
+using System.ComponentModel.DataAnnotations;
+using System.Reflection.PortableExecutable;
+using System.Text;
+
+namespace XMAT.WebServiceCapture.Proxy
+{
+    public abstract class BaseWebSocketMessage
+    {
+        internal StringBuilder _message = new();
+
+        [Description("WEB_SVC_SCRIPT_PROP_DESC_WS_REQNUM")]
+        public int RequestNumber;
+        [Description("WEB_SVC_SCRIPT_PROP_DESC_WS_BODY")]
+        public byte[] BodyBytes { get; set; }
+
+        public BaseWebSocketMessage()
+        {
+            BodyBytes = Array.Empty<byte>();
+        }
+        public override string ToString()
+        {
+            string final = "";
+
+            if (BodyBytes != null && BodyBytes.Length > 0)
+            {
+                final = Encoding.ASCII.GetString(BodyBytes);
+            }
+
+            return final;
+        }
+        public byte[] ToByteArray()
+        {
+            byte[] final = new byte[BodyBytes.Length];
+            Array.Copy(BodyBytes, 0, final, 0, BodyBytes.Length);
+            return final;
+        }
+    }
+
+    public class WebSocketMessage : BaseWebSocketMessage
+    {
+        [Description("WEB_SVC_SCRIPT_PROP_DESC_WS_FROM_HOST")]
+        public bool FromHost { get; set; }
+    }
+}

XMAT/Engines/WebServiceProxy/Logic/Proxy/IWebServiceProxy.cs

@@ -6,7 +6,7 @@
 namespace XMAT.WebServiceCapture.Proxy
 {
     public interface IWebServiceProxy
-    {          
+    {
         event EventHandler<InitialConnectionEventArgs> ReceivedInitialConnection;
         event EventHandler<SslConnectionRequestEventArgs> ReceivedSslConnectionRequest;
         event EventHandler<SslConnectionCompletionEventArgs> CompletedSslConnectionRequest;

XMAT/Engines/WebServiceProxy/Logic/Proxy/IWebSocketProxy.cs

@@ -0,0 +1,15 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System;
+using Windows.Networking.Sockets;
+
+namespace XMAT.WebServiceCapture.Proxy
+{
+    public interface IWebSocketProxy
+    {
+        event EventHandler<WebSocketOpenedEventArgs> WebSocketOpened;
+        event EventHandler<WebSocketMessageEventArgs> WebSocketMessage;
+        event EventHandler<WebSocketClosedEventArgs> WebSocketClosed;
+    }
+}

XMAT/Engines/WebServiceProxy/Logic/Proxy/WebServiceProxy.cs

@@ -397,14 +397,17 @@ private async Task<bool> ForwardRequestAsync(ClientState clientState, ClientRequ
 
         private async Task HandleWebSocketRequest(Uri uri, ClientState clientState, ClientRequest clientRequest)
         {
-            // TODO: this needs to be fleshed out into something bigger,
-            // but since it only partially works and needs a native implementation anyway,
-            // I'm leaving it as-is for now.
-
-            // this will spawn read/write threads that will run until the websocket disconnects,
+            // This will spawn read/write threads that will run until the websocket disconnects,
             // after which the threads will terminate and all will be cleaned up
+
             var wspc = new WebSocketProxy();
-            await wspc.StartWebSocketProxy(uri, clientState, clientRequest);
+
+            // TODO: Add events for driving UI
+            wspc.WebSocketOpened += null;
+            wspc.WebSocketClosed += null;
+            wspc.WebSocketMessage += null;
+
+            await wspc.StartWebSocketProxy(uri, clientState, clientRequest, _logger);
         }
 
         private async Task<bool> HandleWebRequest(Uri uri, ClientState clientState, ClientRequest clientRequest)
@@ -508,39 +511,6 @@ private ClientRequest ParseRequestAndHeaders(int clientId, List<string> lines)
             return clientRequest;
         }
 
-        private int PrefetchContentLengthFromHeaders(int clientId, List<string> lines)
-        {
-            // Return 0 if not found or malformed
-            int contentLength = 0;
-
-            if (lines == null || lines.Count == 0)
-            {
-                _logger.Log(clientId, LogLevel.ERROR, "Connect request has no data.");
-                return contentLength;
-            }
-
-            // Isolate the content-length header only
-            for (int i = 1; i < lines.Count; i++)
-            {
-                string line = lines[i];
-                string[] header = line.Split(':', 2, StringSplitOptions.RemoveEmptyEntries | StringSplitOptions.TrimEntries);
-                if (header.Length < 2)
-                {
-                    _logger.Log(clientId, LogLevel.ERROR, $"Malformed header: {line}");
-                }
-                else
-                {
-                    if (header[0].Trim().ToLower() == "content-length")
-                    {
-                        contentLength = int.Parse(header[1].Trim());
-                        _logger.Log(clientId, LogLevel.DEBUG, $"Prefetch of Content-Length: {contentLength}");
-                    }
-                }
-            }
-
-            return contentLength;
-        }
-
         private async Task<bool> ReturnResponseAsync(ClientRequest clientRequest, ClientState clientState, HttpResponseMessage responseMessage)
         {
             ServerResponse serverResponse = await ParseServerResponseAsync(clientState.ID, responseMessage).ConfigureAwait(false);
@@ -628,6 +598,8 @@ private async Task<Tuple<List<string>, byte[]>> ReadRequestAndHeadersAsync(Clien
             {
                 do
                 {
+
+
                     readThisFrame = await stream.ReadAsync(buffer, readTotal, buffer.Length - readTotal);
                     readTotal += readThisFrame;
 
@@ -642,9 +614,9 @@ private async Task<Tuple<List<string>, byte[]>> ReadRequestAndHeadersAsync(Clien
                 }
                 while (clientState.TcpClient.Available > 0);
             }
-            catch (Exception ex)
+            catch
             {
-                _logger.Log(clientState.ID, LogLevel.ERROR, $"Failed reading client request: {ex}");
+
             }
 
             //File.WriteAllBytes($"XMAT_{clientState.ID}.dat", buffer);
@@ -685,38 +657,13 @@ private async Task<Tuple<List<string>, byte[]>> ReadRequestAndHeadersAsync(Clien
                 }
             }
 
-            // Get the actual body length from the content-length header
-            int expectedBodyLength = 0;
-            expectedBodyLength = PrefetchContentLengthFromHeaders(clientState.ID, lines);
-
-            // Read the rest of the delayed body if needed
-            if (expectedBodyLength > readTotal - startIndex)
-            {
-                try
-                {
-                    do
-                    {
-                        _logger.Log(clientState.ID, LogLevel.DEBUG, $"Delayed body bytes, reading {expectedBodyLength - (readTotal - startIndex)} at {readTotal - startIndex}");
-
-                        readThisFrame = await stream.ReadAsync(buffer, expectedBodyLength - (readTotal - startIndex), buffer.Length - readTotal);
-                        readTotal += readThisFrame;
 
-                        // Do we need to resize?
-                        if (readTotal >= buffer.Length)
-                        {
-                            _logger.Log(clientState.ID, LogLevel.DEBUG, $"Resizing buffer from {buffer.Length} bytes, to {buffer.Length * 2} bytes");
-                            Array.Resize<byte>(ref buffer, buffer.Length * 2);
-                        }
-                    }
-                    while (expectedBodyLength > readTotal - startIndex);
-                }
-                catch (Exception ex)
-                {
-                    _logger.Log(clientState.ID, LogLevel.ERROR, $"Failed reading client request for delayed body: {ex}");
-                }
-            }
-
-            _logger.Log(clientState.ID, LogLevel.DEBUG, $"Done with reading, read {readTotal} bytes total, and processed {startIndex} as headers, {readTotal-startIndex} as body");
+            //string strHeadersTotal = strAllData.Substring(0, startIndex);
+            //string strBody = strAllData.Substring(startIndex);
+            //_logger.Log(clientState.ID, LogLevel.DEBUG, $"Whole buffer was {strAllData}");
+            //_logger.Log(clientState.ID, LogLevel.DEBUG, $"Headers buffer was {strHeadersTotal}");
+            //_logger.Log(clientState.ID, LogLevel.DEBUG, $"Body buffer was {strBody}");
+            //_logger.Log(clientState.ID, LogLevel.DEBUG, $"Done with read, read {readTotal} bytes total, and processed {startIndex} as headers, body remainder is {readTotal-startIndex}");
 
             // resize the buffer, otherwise we're taking up unnecessary memory
             Array.Resize<byte>(ref buffer, readTotal);
@@ -750,8 +697,6 @@ private void CloseClientState(ClientState clientState)
                     Timestamp = DateTime.Now
                 }
             );
-
-            _logger.Log(clientState.ID, LogLevel.DEBUG, $"ConnectionClosed Event: {clientState.ID} | {DateTime.Now.ToString()}");
         }
 
         private bool RaiseReceivedInitialConnection(int connectionID, TcpClient client)