From 7c626431a10fee5aa6d6f47f97e0b88c5812311f Mon Sep 17 00:00:00 2001
From: Louis DeJardin <louis.dejardin@gmail.com>
Date: Fri, 27 Jun 2014 22:25:38 -0700
Subject: [PATCH] Adding Upgrade support, using WebSocket middleware from
 sample app

---
 src/Kestrel/ServerRequest.cs                  | 35 ++++++++++++++-
 .../Http/Frame.cs                             |  6 ++-
 src/SampleApp/Startup.cs                      | 45 +++++++++++++++++--
 src/SampleApp/project.json                    |  3 +-
 4 files changed, 82 insertions(+), 7 deletions(-)

diff --git a/src/Kestrel/ServerRequest.cs b/src/Kestrel/ServerRequest.cs
index 643e464d5f..d558308a52 100644
--- a/src/Kestrel/ServerRequest.cs
+++ b/src/Kestrel/ServerRequest.cs
@@ -3,10 +3,12 @@ using Microsoft.AspNet.Server.Kestrel.Http;
 using System;
 using System.Collections.Generic;
 using System.IO;
+using System.Linq;
+using System.Threading.Tasks;
 
 namespace Kestrel
 {
-    public class ServerRequest : IHttpRequestFeature, IHttpResponseFeature
+    public class ServerRequest : IHttpRequestFeature, IHttpResponseFeature, IHttpOpaqueUpgradeFeature
     {
         Frame _frame;
         string _scheme;
@@ -172,9 +174,40 @@ namespace Kestrel
                 _frame.ResponseBody = value;
             }
         }
+
         void IHttpResponseFeature.OnSendingHeaders(Action<object> callback, object state)
         {
             _frame.OnSendingHeaders(callback, state);
         }
+
+        bool IHttpOpaqueUpgradeFeature.IsUpgradableRequest
+        {
+            get
+            {
+                string[] values;
+                if (_frame.RequestHeaders.TryGetValue("Connection", out values))
+                {
+                    return values.Any(value => value.IndexOf("upgrade", StringComparison.OrdinalIgnoreCase) != -1);
+                }
+                return false;
+            }
+        }
+
+        async Task<Stream> IHttpOpaqueUpgradeFeature.UpgradeAsync()
+        {
+            _frame.StatusCode = 101;
+            _frame.ReasonPhrase = "Switching Protocols";
+            _frame.ResponseHeaders["Connection"] = new string[] { "Upgrade" };
+            if (!_frame.ResponseHeaders.ContainsKey("Upgrade"))
+            {
+                string[] values;
+                if (_frame.RequestHeaders.TryGetValue("Upgrade", out values))
+                {
+                    _frame.ResponseHeaders["Upgrade"] = values;
+                }
+            }
+            _frame.ProduceStart();
+            return _frame.DuplexStream;
+        }
     }
 }
diff --git a/src/Microsoft.AspNet.Server.Kestrel/Http/Frame.cs b/src/Microsoft.AspNet.Server.Kestrel/Http/Frame.cs
index 428b0ca6b4..fd0696d180 100644
--- a/src/Microsoft.AspNet.Server.Kestrel/Http/Frame.cs
+++ b/src/Microsoft.AspNet.Server.Kestrel/Http/Frame.cs
@@ -63,7 +63,6 @@ namespace Microsoft.AspNet.Server.Kestrel.Http
 
         List<KeyValuePair<Action<object>, object>> _onSendingHeaders;
         object _onSendingHeadersSync = new Object();
-        Stream _duplexStream;
 
         public Frame(ConnectionContext context) : base(context)
         {
@@ -87,6 +86,9 @@ namespace Microsoft.AspNet.Server.Kestrel.Http
         public IDictionary<string, string[]> ResponseHeaders { get; set; }
         public Stream ResponseBody { get; set; }
 
+        public Stream DuplexStream { get; set; }
+
+
 
         /*
         public bool LocalIntakeFin
@@ -175,7 +177,7 @@ namespace Microsoft.AspNet.Server.Kestrel.Http
             _keepAlive = MessageBody.RequestKeepAlive;
             RequestBody = new FrameRequestStream(MessageBody);
             ResponseBody = new FrameResponseStream(this);
-            _duplexStream = new FrameDuplexStream(RequestBody, ResponseBody);
+            DuplexStream = new FrameDuplexStream(RequestBody, ResponseBody);
             SocketInput.Free();
             Task.Run(ExecuteAsync);
         }
diff --git a/src/SampleApp/Startup.cs b/src/SampleApp/Startup.cs
index 9eba33a386..65462519d6 100644
--- a/src/SampleApp/Startup.cs
+++ b/src/SampleApp/Startup.cs
@@ -1,5 +1,8 @@
 using Microsoft.AspNet.Builder;
 using System;
+using System.Net.WebSockets;
+using System.Threading;
+using System.Threading.Tasks;
 
 namespace SampleApp
 {
@@ -7,6 +10,8 @@ namespace SampleApp
     {
         public void Configure(IBuilder app)
         {
+            app.UseWebSockets();
+
             app.Run(async context =>
             {
                 Console.WriteLine("{0} {1}{2}{3}",
@@ -15,10 +20,44 @@ namespace SampleApp
                     context.Request.Path,
                     context.Request.QueryString);
 
-                context.Response.ContentLength = 11;
-                context.Response.ContentType = "text/plain";
-                await context.Response.WriteAsync("Hello world");
+                if (context.IsWebSocketRequest)
+                {
+                    var webSocket = await context.AcceptWebSocketAsync();
+                    await EchoAsync(webSocket);
+                }
+                else
+                {
+                    context.Response.ContentLength = 11;
+                    context.Response.ContentType = "text/plain";
+                    await context.Response.WriteAsync("Hello world");
+                }
             });
         }
+
+        public async Task EchoAsync(WebSocket webSocket)
+        {
+            var buffer = new ArraySegment<byte>(new byte[8192]);
+            for (; ;)
+            {
+                var result = await webSocket.ReceiveAsync(
+                    buffer,
+                    CancellationToken.None);
+
+                if (result.MessageType == WebSocketMessageType.Close)
+                {
+                    return;
+                }
+                else if (result.MessageType == WebSocketMessageType.Text)
+                {
+                    Console.WriteLine("{0}", System.Text.Encoding.UTF8.GetString(buffer.Array, 0, result.Count));
+                }
+
+                await webSocket.SendAsync(
+                    new ArraySegment<byte>(buffer.Array, 0, result.Count),
+                    result.MessageType,
+                    result.EndOfMessage,
+                    CancellationToken.None);
+            }
+        }
     }
 }
\ No newline at end of file
diff --git a/src/SampleApp/project.json b/src/SampleApp/project.json
index 19174804d7..f34a5acc74 100644
--- a/src/SampleApp/project.json
+++ b/src/SampleApp/project.json
@@ -2,7 +2,8 @@
     "version": "1.0.0-*",
     "dependencies": {
         "Kestrel": "1.0.0-*",
-        "Microsoft.AspNet.Server.WebListener": "1.0.0-*"
+        "Microsoft.AspNet.Server.WebListener": "1.0.0-*",
+        "Microsoft.AspNet.WebSockets.Middleware": "1.0.0-*"
     },
     "configurations": {
         "net45": { },