From c8f9364e3e88b15b313d4bc0cf28bdc184e01345 Mon Sep 17 00:00:00 2001
From: Cesar Blum Silveira <cesars@microsoft.com>
Date: Mon, 18 Sep 2017 14:14:31 -0700
Subject: [PATCH] Expose connection features to HTTP/2 connections.

---
 .../Internal/Http2/Http2Connection.cs         |  4 +++
 .../Internal/Http2/Http2ConnectionContext.cs  |  2 ++
 src/Kestrel.Core/Internal/HttpConnection.cs   | 29 ++++++++++++-------
 3 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/src/Kestrel.Core/Internal/Http2/Http2Connection.cs b/src/Kestrel.Core/Internal/Http2/Http2Connection.cs
index 494379a518..5796b73a91 100644
--- a/src/Kestrel.Core/Internal/Http2/Http2Connection.cs
+++ b/src/Kestrel.Core/Internal/Http2/Http2Connection.cs
@@ -8,6 +8,7 @@ using System.IO.Pipelines;
 using System.Text;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Hosting.Server;
+using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Protocols;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2.HPack;
@@ -50,6 +51,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2
 
         public IKestrelTrace Log => _context.ServiceContext.Log;
 
+        public IFeatureCollection ConnectionFeatures => _context.ConnectionFeatures;
+
         public void Abort(Exception ex)
         {
             _stopping = true;
@@ -266,6 +269,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2
                 ConnectionId = ConnectionId,
                 StreamId =  _incomingFrame.StreamId,
                 ServiceContext = _context.ServiceContext,
+                ConnectionFeatures = _context.ConnectionFeatures,
                 PipeFactory = _context.PipeFactory,
                 LocalEndPoint = _context.LocalEndPoint,
                 RemoteEndPoint = _context.RemoteEndPoint,
diff --git a/src/Kestrel.Core/Internal/Http2/Http2ConnectionContext.cs b/src/Kestrel.Core/Internal/Http2/Http2ConnectionContext.cs
index 2956040b16..1c089cc5b0 100644
--- a/src/Kestrel.Core/Internal/Http2/Http2ConnectionContext.cs
+++ b/src/Kestrel.Core/Internal/Http2/Http2ConnectionContext.cs
@@ -3,6 +3,7 @@
 
 using System.IO.Pipelines;
 using System.Net;
+using Microsoft.AspNetCore.Http.Features;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2
 {
@@ -10,6 +11,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2
     {
         public string ConnectionId { get; set; }
         public ServiceContext ServiceContext { get; set; }
+        public IFeatureCollection ConnectionFeatures { get; set; }
         public PipeFactory PipeFactory { get; set; }
         public IPEndPoint LocalEndPoint { get; set; }
         public IPEndPoint RemoteEndPoint { get; set; }
diff --git a/src/Kestrel.Core/Internal/HttpConnection.cs b/src/Kestrel.Core/Internal/HttpConnection.cs
index e74b46323b..9ccca7d182 100644
--- a/src/Kestrel.Core/Internal/HttpConnection.cs
+++ b/src/Kestrel.Core/Internal/HttpConnection.cs
@@ -115,19 +115,10 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
                 // _http1Connection must be initialized before adding the connection to the connection manager
                 CreateHttp1Connection(httpApplication, transport, application);
 
-                // _http2Connection must be initialized before yield control to the transport thread,
+                // _http2Connection must be initialized before yielding control to the transport thread,
                 // to prevent a race condition where _http2Connection.Abort() is called just as
                 // _http2Connection is about to be initialized.
-                _http2Connection = new Http2Connection(new Http2ConnectionContext
-                {
-                    ConnectionId = _context.ConnectionId,
-                    ServiceContext = _context.ServiceContext,
-                    PipeFactory = PipeFactory,
-                    LocalEndPoint = LocalEndPoint,
-                    RemoteEndPoint = RemoteEndPoint,
-                    Application = application,
-                    Transport = transport
-                });
+                CreateHttp2Connection(httpApplication, transport, application);
 
                 // Do this before the first await so we don't yield control to the transport until we've
                 // added the connection to the connection manager
@@ -135,6 +126,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
                 _lastTimestamp = _context.ServiceContext.SystemClock.UtcNow.Ticks;
 
                 _http1Connection.ConnectionFeatures.Set<IConnectionTimeoutFeature>(this);
+                _http2Connection.ConnectionFeatures.Set<IConnectionTimeoutFeature>(this);
 
                 if (adaptedPipeline != null)
                 {
@@ -190,6 +182,21 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
             });
         }
 
+        internal void CreateHttp2Connection<TContext>(IHttpApplication<TContext> httpApplication, IPipeConnection transport, IPipeConnection application)
+        {
+            _http2Connection = new Http2Connection(new Http2ConnectionContext
+            {
+                ConnectionId = _context.ConnectionId,
+                ServiceContext = _context.ServiceContext,
+                ConnectionFeatures = _context.ConnectionFeatures,
+                PipeFactory = PipeFactory,
+                LocalEndPoint = LocalEndPoint,
+                RemoteEndPoint = RemoteEndPoint,
+                Application = application,
+                Transport = transport
+            });
+        }
+
         public void OnConnectionClosed(Exception ex)
         {
             Abort(ex);