Fix DelegationRule to work after receiver restarts (#40967)

Author: NGloreous

src/Servers/HttpSys/HttpSysServer.slnf

@@ -4,9 +4,11 @@
     "projects": [
       "src\\DefaultBuilder\\src\\Microsoft.AspNetCore.csproj",
       "src\\Extensions\\Features\\src\\Microsoft.Extensions.Features.csproj",
+      "src\\FileProviders\\Embedded\\src\\Microsoft.Extensions.FileProviders.Embedded.csproj",
       "src\\Hosting\\Abstractions\\src\\Microsoft.AspNetCore.Hosting.Abstractions.csproj",
       "src\\Hosting\\Hosting\\src\\Microsoft.AspNetCore.Hosting.csproj",
       "src\\Hosting\\Server.Abstractions\\src\\Microsoft.AspNetCore.Hosting.Server.Abstractions.csproj",
+      "src\\Hosting\\Server.IntegrationTesting\\src\\Microsoft.AspNetCore.Server.IntegrationTesting.csproj",
       "src\\Http\\Authentication.Abstractions\\src\\Microsoft.AspNetCore.Authentication.Abstractions.csproj",
       "src\\Http\\Authentication.Core\\src\\Microsoft.AspNetCore.Authentication.Core.csproj",
       "src\\Http\\Headers\\src\\Microsoft.Net.Http.Headers.csproj",

src/Servers/HttpSys/src/DelegationRule.cs

@@ -13,17 +13,19 @@ namespace Microsoft.AspNetCore.Server.HttpSys
     public class DelegationRule : IDisposable
     {
         private readonly ILogger _logger;
-        private readonly UrlGroup _urlGroup;
         private readonly UrlGroup _sourceQueueUrlGroup;
         private bool _disposed;
+
         /// <summary>
         /// The name of the Http.Sys request queue
         /// </summary>
         public string QueueName { get; }
+
         /// <summary>
         /// The URL of the Http.Sys Url Prefix
         /// </summary>
         public string UrlPrefix { get; }
+
         internal RequestQueue Queue { get; }
 
         internal DelegationRule(UrlGroup sourceQueueUrlGroup, string queueName, string urlPrefix, ILogger logger)
@@ -32,8 +34,7 @@ internal DelegationRule(UrlGroup sourceQueueUrlGroup, string queueName, string u
             _logger = logger;
             QueueName = queueName;
             UrlPrefix = urlPrefix;
-            Queue = new RequestQueue(queueName, UrlPrefix, _logger, receiver: true);
-            _urlGroup = Queue.UrlGroup;
+            Queue = new RequestQueue(queueName, _logger);
         }
 
         /// <inheritdoc />
@@ -51,7 +52,6 @@ public void Dispose()
                 _sourceQueueUrlGroup.UnSetDelegationProperty(Queue, throwOnError: false);
             }
             catch (ObjectDisposedException) { /* Server may have been shutdown */ }
-            _urlGroup.Dispose();
             Queue.Dispose();
         }
     }

src/Servers/HttpSys/src/NativeInterop/RequestQueue.cs

@@ -19,25 +19,16 @@ internal class RequestQueue
         private readonly ILogger _logger;
         private bool _disposed;
 
-        internal RequestQueue(string requestQueueName, string urlPrefix, ILogger logger, bool receiver)
-            : this(urlGroup: null!, requestQueueName, RequestQueueMode.Attach, logger, receiver)
+        internal RequestQueue(string requestQueueName, ILogger logger)
+            : this(urlGroup: null, requestQueueName, RequestQueueMode.Attach, logger, receiver: true)
         {
-            try
-            {
-                UrlGroup = new UrlGroup(this, UrlPrefix.Create(urlPrefix), logger);
-            }
-            catch
-            {
-                Dispose();
-                throw;
-            }
         }
 
         internal RequestQueue(UrlGroup urlGroup, string? requestQueueName, RequestQueueMode mode, ILogger logger)
             : this(urlGroup, requestQueueName, mode, logger, false)
         { }
 
-        private RequestQueue(UrlGroup urlGroup, string? requestQueueName, RequestQueueMode mode, ILogger logger, bool receiver)
+        private RequestQueue(UrlGroup? urlGroup, string? requestQueueName, RequestQueueMode mode, ILogger logger, bool receiver)
         {
             _mode = mode;
             UrlGroup = urlGroup;
@@ -117,10 +108,15 @@ private RequestQueue(UrlGroup urlGroup, string? requestQueueName, RequestQueueMo
         internal SafeHandle Handle { get; }
         internal ThreadPoolBoundHandle BoundHandle { get; }
 
-        internal UrlGroup UrlGroup { get; }
+        internal UrlGroup? UrlGroup { get; }
 
         internal unsafe void AttachToUrlGroup()
         {
+            if (UrlGroup == null)
+            {
+                throw new NotSupportedException("Can't attach when UrlGroup is null");
+            }
+
             Debug.Assert(Created);
             CheckDisposed();
             // Set the association between request queue and url group. After this, requests for registered urls will
@@ -138,6 +134,11 @@ internal unsafe void AttachToUrlGroup()
 
         internal unsafe void DetachFromUrlGroup()
         {
+            if (UrlGroup == null)
+            {
+                throw new NotSupportedException("Can't detach when UrlGroup is null");
+            }
+
             Debug.Assert(Created);
             CheckDisposed();
             // Break the association between request queue and url group. After this, requests for registered urls

src/Servers/HttpSys/src/NativeInterop/UrlGroup.cs

@@ -41,24 +41,6 @@ internal unsafe UrlGroup(ServerSession serverSession, ILogger logger)
             Id = urlGroupId;
         }
 
-        internal unsafe UrlGroup(RequestQueue requestQueue, UrlPrefix url, ILogger logger)
-        {
-            _logger = logger;
-
-            ulong urlGroupId = 0;
-            _created = false;
-            var statusCode = HttpApi.HttpFindUrlGroupId(
-                url.FullPrefix, requestQueue.Handle, &urlGroupId);
-
-            if (statusCode != UnsafeNclNativeMethods.ErrorCodes.ERROR_SUCCESS)
-            {
-                throw new HttpSysException((int)statusCode);
-            }
-
-            Debug.Assert(urlGroupId != 0, "Invalid id returned by HttpCreateUrlGroup");
-            Id = urlGroupId;
-        }
-
         internal ulong Id { get; private set; }
 
         internal unsafe void SetMaxConnections(long maxConnections)

src/Servers/HttpSys/src/RequestProcessing/RequestContext.cs

@@ -289,10 +289,13 @@ internal unsafe void Delegate(DelegationRule destination)
                     PropertyInfoLength = (uint)System.Text.Encoding.Unicode.GetByteCount(destination.UrlPrefix)
                 };
 
+                // Passing 0 for delegateUrlGroupId allows http.sys to find the right group for the
+                // URL passed in via the property above. If we passed in the receiver's URL group id
+                // instead of 0, then delegation would fail if the receiver restarted.
                 statusCode = HttpApi.HttpDelegateRequestEx(source.Handle,
                                                                destination.Queue.Handle,
                                                                Request.RequestId,
-                                                               destination.Queue.UrlGroup.Id,
+                                                               delegateUrlGroupId: 0,
                                                                propertyInfoSetSize: 1,
                                                                &property);
             }

src/Servers/HttpSys/src/ServerDelegationPropertyFeature.cs

@@ -14,18 +14,23 @@ namespace Microsoft.AspNetCore.Server.HttpSys
     internal class ServerDelegationPropertyFeature : IServerDelegationFeature
     {
         private readonly ILogger _logger;
-        private readonly RequestQueue _queue;
+        private readonly UrlGroup _urlGroup;
 
         public ServerDelegationPropertyFeature(RequestQueue queue, ILogger logger)
         {
-            _queue = queue;
+            if (queue.UrlGroup == null)
+            {
+                throw new ArgumentException($"{nameof(queue)}.UrlGroup can't be null");
+            }
+
+            _urlGroup = queue.UrlGroup;
             _logger = logger;
         }
 
         public DelegationRule CreateDelegationRule(string queueName, string uri)
         {
-            var rule = new DelegationRule(_queue.UrlGroup, queueName, uri, _logger);
-            _queue.UrlGroup.SetDelegationProperty(rule.Queue);
+            var rule = new DelegationRule(_urlGroup, queueName, uri, _logger);
+            _urlGroup.SetDelegationProperty(rule.Queue);
             return rule;
         }
     }

src/Servers/HttpSys/test/FunctionalTests/DelegateTests.cs

@@ -1,13 +1,11 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using System;
-using System.IO;
 using System.Net.Http;
-using System.Threading.Tasks;
+using System.Runtime.InteropServices;
 using Microsoft.AspNetCore.Http;
+using Microsoft.AspNetCore.HttpSys.Internal;
 using Microsoft.AspNetCore.Testing;
-using Xunit;
 
 namespace Microsoft.AspNetCore.Server.HttpSys.FunctionalTests
 {
@@ -198,6 +196,72 @@ public async Task UpdateDelegationRuleTest()
             destination?.Dispose();
         }
 
+        [ConditionalFact]
+        [DelegateSupportedCondition(true)]
+        public async Task DelegateAfterReceiverRestart()
+        {
+            var queueName = Guid.NewGuid().ToString();
+            using var receiver = Utilities.CreateHttpServer(out var receiverAddress, async httpContext =>
+            {
+                await httpContext.Response.WriteAsync(_expectedResponseString);
+            },
+            options =>
+            {
+                options.RequestQueueName = queueName;
+            });
+
+            DelegationRule destination = default;
+            using var delegator = Utilities.CreateHttpServer(out var delegatorAddress, httpContext =>
+            {
+                var delegateFeature = httpContext.Features.Get<IHttpSysRequestDelegationFeature>();
+                delegateFeature.DelegateRequest(destination);
+                return Task.CompletedTask;
+            });
+
+            var delegationProperty = delegator.Features.Get<IServerDelegationFeature>();
+            destination = delegationProperty.CreateDelegationRule(queueName, receiverAddress);
+
+            var responseString = await SendRequestAsync(delegatorAddress);
+            Assert.Equal(_expectedResponseString, responseString);
+
+            // Stop the receiver
+            receiver?.Dispose();
+
+            // Start the receiver again but this time we need to attach to the existing queue.
+            // Due to https://github.com/dotnet/aspnetcore/issues/40359, we have to manually
+            // register URL prefixes and attach the server's queue to them.
+            using var receiverRestarted = (MessagePump)Utilities.CreateHttpServer(out receiverAddress, async httpContext =>
+            {
+                await httpContext.Response.WriteAsync(_expectedResponseString);
+            },
+            options =>
+            {
+                options.RequestQueueName = queueName;
+                options.RequestQueueMode = RequestQueueMode.Attach;
+                options.UrlPrefixes.Clear();
+                options.UrlPrefixes.Add(receiverAddress);
+            });
+            AttachToUrlGroup(receiverRestarted.Listener.RequestQueue);
+            receiverRestarted.Listener.Options.UrlPrefixes.RegisterAllPrefixes(receiverRestarted.Listener.UrlGroup);
+
+            responseString = await SendRequestAsync(delegatorAddress);
+            Assert.Equal(_expectedResponseString, responseString);
+
+            destination?.Dispose();
+        }
+
+        private unsafe void AttachToUrlGroup(RequestQueue requestQueue)
+        {
+            var info = new HttpApiTypes.HTTP_BINDING_INFO();
+            info.Flags = HttpApiTypes.HTTP_FLAGS.HTTP_PROPERTY_FLAG_PRESENT;
+            info.RequestQueueHandle = requestQueue.Handle.DangerousGetHandle();
+
+            var infoptr = new IntPtr(&info);
+
+            requestQueue.UrlGroup.SetProperty(HttpApiTypes.HTTP_SERVER_PROPERTY.HttpServerBindingProperty,
+                infoptr, (uint)Marshal.SizeOf<HttpApiTypes.HTTP_BINDING_INFO>());
+        }
+
         private async Task<string> SendRequestAsync(string uri)
         {
             using var client = new HttpClient();