Graceful close handshake

samples/ANCMStressTestSample/Program.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Linq;
+using System.Threading;
 using Microsoft.AspNetCore.Authentication;
 using Microsoft.AspNetCore.Builder;
 using Microsoft.AspNetCore.Hosting;
@@ -17,7 +18,7 @@ namespace ANCMStressTestApp
     public class Program
     {
         public static IApplicationLifetime AppLifetime;
-        public static bool AppLifetimeStopping = false;
+        public static CancellationTokenSource Cts = new CancellationTokenSource();
 
         public static void Main(string[] args)
         {
@@ -33,7 +34,8 @@ namespace ANCMStressTestApp
             AppLifetime = (IApplicationLifetime)host.Services.GetService(typeof(IApplicationLifetime));
             AppLifetime.ApplicationStopping.Register(
                 () => {
-                    AppLifetimeStopping = true;
+                    Cts.Cancel();
+                    Cts.Dispose();
                 }
             );
 

samples/ANCMStressTestSample/Startup.cs

@@ -158,7 +158,7 @@ namespace ANCMStressTestApp
                 var upgradeFeature = context.Features.Get<IHttpUpgradeFeature>();
 
                 // Generate WebSocket response headers
-                string key = string.Join(", ", context.Request.Headers[Constants.Headers.SecWebSocketKey]);
+                string key = context.Request.Headers[Constants.Headers.SecWebSocketKey].ToString();
                 var responseHeaders = HandshakeHelpers.GenerateResponseHeaders(key);
                 foreach (var headerPair in responseHeaders)
                 {
@@ -171,43 +171,59 @@ namespace ANCMStressTestApp
                 // Get the WebSocket object
                 var ws = WebSocketProtocol.CreateFromStream(opaqueTransport, isServer: true, subProtocol: null, keepAliveInterval: TimeSpan.FromMinutes(2));
 
-                await Echo(ws);
+                await Echo(ws, Program.Cts.Token);
             });
         }
 
-        private async Task Echo(WebSocket webSocket)
+        private async Task Echo(WebSocket webSocket, CancellationToken token)
         {
-            var buffer = new byte[1024 * 4];
+            try
-            var result = await webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), CancellationToken.None);
-            bool closeFromServer = false;
-            string closeFromServerCmd = "CloseFromServer";
-            int closeFromServerLength = closeFromServerCmd.Length;
-
-            while (!result.CloseStatus.HasValue)
             {
-                if ((result.Count == closeFromServerLength && System.Text.Encoding.ASCII.GetString(buffer).Substring(0, result.Count) == closeFromServerCmd)
+                var buffer = new byte[1024 * 4];
-                    || Program.AppLifetimeStopping == true)
+                var result = await webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), token);
+                bool closeFromServer = false;
+                string closeFromServerCmd = "CloseFromServer";
+                int closeFromServerLength = closeFromServerCmd.Length;
+
+                while (!result.CloseStatus.HasValue && !token.IsCancellationRequested && !closeFromServer)
                 {
-                    // start closing handshake from backend process when client send "CloseFromServer" text message 
+                    if (result.Count == closeFromServerLength &&
-                    // or when any message is sent from client during the graceful shutdown.
+                        System.Text.Encoding.ASCII.GetString(buffer).Substring(0, result.Count) == closeFromServerCmd)
-                    closeFromServer = true;
+                    {
-                    await webSocket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, closeFromServerCmd, CancellationToken.None);
+                        // The client sent "CloseFromServer" text message to request the server to close (a test scenario).
+                        closeFromServer = true;
+                    }
+                    else
+                    {
+                        await webSocket.SendAsync(new ArraySegment<byte>(buffer, 0, result.Count), result.MessageType, result.EndOfMessage, token);
+                        result = await webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), token);
+                    }
+                }
+
+                if (result.CloseStatus.HasValue)
+                {
+                    // Client-initiated close handshake
+                    await webSocket.CloseAsync(result.CloseStatus.Value, result.CloseStatusDescription, CancellationToken.None);
                 }
                 else
                 {
-                    await webSocket.SendAsync(new ArraySegment<byte>(buffer, 0, result.Count), result.MessageType, result.EndOfMessage, CancellationToken.None);
+                    // Server-initiated close handshake due to either of the two conditions:
+                    // (1) The applicaton host is performing a graceful shutdown.
+                    // (2) The client sent "CloseFromServer" text message to request the server to close (a test scenario).
+                    await webSocket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, closeFromServerCmd, CancellationToken.None);
+
+                    // The server has sent the Close frame.
+                    // Stop sending but keep receiving until we get the Close frame from the client.
+                    while (!result.CloseStatus.HasValue)
+                    {
+                        result = await webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), CancellationToken.None);
+                    }
                 }
-
-                result = await webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), CancellationToken.None);
             }
-
+            catch (Exception e)
-            if (!closeFromServer)
             {
-                await webSocket.CloseAsync(result.CloseStatus.Value, result.CloseStatusDescription, CancellationToken.None);
+                Console.WriteLine("{0} Exception caught!", e);
-                
             }
-
-            webSocket.Dispose();
         }
     }
 }

samples/ANCMStressTestSample/WebSockets/Constants.cs

@@ -12,10 +12,7 @@ namespace ANCMStressTestApp
             public const string Connection = "Connection";
             public const string ConnectionUpgrade = "Upgrade";
             public const string SecWebSocketKey = "Sec-WebSocket-Key";
-            public const string SecWebSocketVersion = "Sec-WebSocket-Version";
-            public const string SecWebSocketProtocol = "Sec-WebSocket-Protocol";
             public const string SecWebSocketAccept = "Sec-WebSocket-Accept";
-            public const string SupportedVersion = "13";
         }
     }
 }