From c181d1db9621aa0e24666dee29d931f8061b7010 Mon Sep 17 00:00:00 2001 From: Yanbing Shi Date: Fri, 23 Mar 2018 17:07:16 -0700 Subject: [PATCH] Graceful close handshake --- samples/ANCMStressTestSample/Program.cs | 6 +- samples/ANCMStressTestSample/Startup.cs | 66 ++++++++++++------- .../WebSockets/Constants.cs | 3 - 3 files changed, 45 insertions(+), 30 deletions(-) diff --git a/samples/ANCMStressTestSample/Program.cs b/samples/ANCMStressTestSample/Program.cs index 1591952577..d3b60867b0 100644 --- a/samples/ANCMStressTestSample/Program.cs +++ b/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(); } ); diff --git a/samples/ANCMStressTestSample/Startup.cs b/samples/ANCMStressTestSample/Startup.cs index cfabe95866..0dd43be50a 100644 --- a/samples/ANCMStressTestSample/Startup.cs +++ b/samples/ANCMStressTestSample/Startup.cs @@ -158,7 +158,7 @@ namespace ANCMStressTestApp var upgradeFeature = context.Features.Get(); // 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]; - var result = await webSocket.ReceiveAsync(new ArraySegment(buffer), CancellationToken.None); - bool closeFromServer = false; - string closeFromServerCmd = "CloseFromServer"; - int closeFromServerLength = closeFromServerCmd.Length; - - while (!result.CloseStatus.HasValue) + try { - if ((result.Count == closeFromServerLength && System.Text.Encoding.ASCII.GetString(buffer).Substring(0, result.Count) == closeFromServerCmd) - || Program.AppLifetimeStopping == true) + var buffer = new byte[1024 * 4]; + var result = await webSocket.ReceiveAsync(new ArraySegment(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 - // or when any message is sent from client during the graceful shutdown. - closeFromServer = true; - await webSocket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, closeFromServerCmd, CancellationToken.None); + if (result.Count == closeFromServerLength && + System.Text.Encoding.ASCII.GetString(buffer).Substring(0, result.Count) == closeFromServerCmd) + { + // The client sent "CloseFromServer" text message to request the server to close (a test scenario). + closeFromServer = true; + } + else + { + await webSocket.SendAsync(new ArraySegment(buffer, 0, result.Count), result.MessageType, result.EndOfMessage, token); + result = await webSocket.ReceiveAsync(new ArraySegment(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(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(buffer), CancellationToken.None); + } } - - result = await webSocket.ReceiveAsync(new ArraySegment(buffer), CancellationToken.None); } - - if (!closeFromServer) + catch (Exception e) { - await webSocket.CloseAsync(result.CloseStatus.Value, result.CloseStatusDescription, CancellationToken.None); - + Console.WriteLine("{0} Exception caught!", e); } - - webSocket.Dispose(); } } } diff --git a/samples/ANCMStressTestSample/WebSockets/Constants.cs b/samples/ANCMStressTestSample/WebSockets/Constants.cs index 43b1822ca9..e002def8f1 100644 --- a/samples/ANCMStressTestSample/WebSockets/Constants.cs +++ b/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"; } } }