Fix DelegationRule to work after receiver restarts (#40420)

7 files changed, 108 insertions, 51 deletions

diff --git a/src/Servers/HttpSys/HttpSysServer.slnf b/src/Servers/HttpSys/HttpSysServer.slnf
index 3990e33925..c546d5797a 100644
--- a/src/Servers/HttpSys/HttpSysServer.slnf
+++ b/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",
diff --git a/src/Servers/HttpSys/src/DelegationRule.cs b/src/Servers/HttpSys/src/DelegationRule.cs
index 7df8c7e9b7..d3286b4df8 100644
--- a/src/Servers/HttpSys/src/DelegationRule.cs
+++ b/src/Servers/HttpSys/src/DelegationRule.cs
@@ -12,17 +12,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)
@@ -31,8 +33,7 @@ public class DelegationRule : IDisposable
         _logger = logger;
         QueueName = queueName;
         UrlPrefix = urlPrefix;
-        Queue = new RequestQueue(queueName, UrlPrefix, _logger, receiver: true);
-        _urlGroup = Queue.UrlGroup;
+        Queue = new RequestQueue(queueName, _logger);
     }
 
     /// <inheritdoc />
@@ -50,7 +51,6 @@ public class DelegationRule : IDisposable
             _sourceQueueUrlGroup.UnSetDelegationProperty(Queue, throwOnError: false);
         }
         catch (ObjectDisposedException) { /* Server may have been shutdown */ }
-        _urlGroup.Dispose();
         Queue.Dispose();
     }
 }
diff --git a/src/Servers/HttpSys/src/NativeInterop/RequestQueue.cs b/src/Servers/HttpSys/src/NativeInterop/RequestQueue.cs
index 353ea0ebf8..eb235706b4 100644
--- a/src/Servers/HttpSys/src/NativeInterop/RequestQueue.cs
+++ b/src/Servers/HttpSys/src/NativeInterop/RequestQueue.cs
@@ -17,25 +17,16 @@ internal sealed partial 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;
@@ -115,10 +106,15 @@ internal sealed partial class RequestQueue
     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
@@ -136,6 +132,11 @@ internal sealed partial class RequestQueue
 
     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
diff --git a/src/Servers/HttpSys/src/NativeInterop/UrlGroup.cs b/src/Servers/HttpSys/src/NativeInterop/UrlGroup.cs
index dbe3ca96f5..06ebd23a8b 100644
--- a/src/Servers/HttpSys/src/NativeInterop/UrlGroup.cs
+++ b/src/Servers/HttpSys/src/NativeInterop/UrlGroup.cs
@@ -40,24 +40,6 @@ internal partial class UrlGroup : IDisposable
         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)
diff --git a/src/Servers/HttpSys/src/RequestProcessing/RequestContext.cs b/src/Servers/HttpSys/src/RequestProcessing/RequestContext.cs
index af329d2571..6b39296569 100644
--- a/src/Servers/HttpSys/src/RequestProcessing/RequestContext.cs
+++ b/src/Servers/HttpSys/src/RequestProcessing/RequestContext.cs
@@ -285,10 +285,13 @@ internal partial class RequestContext : NativeRequestContext, IThreadPoolWorkIte
                 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);
         }
diff --git a/src/Servers/HttpSys/src/ServerDelegationPropertyFeature.cs b/src/Servers/HttpSys/src/ServerDelegationPropertyFeature.cs
index c83bff2536..ccdb71c174 100644
--- a/src/Servers/HttpSys/src/ServerDelegationPropertyFeature.cs
+++ b/src/Servers/HttpSys/src/ServerDelegationPropertyFeature.cs
@@ -8,18 +8,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;
     }
 }
diff --git a/src/Servers/HttpSys/test/FunctionalTests/DelegateTests.cs b/src/Servers/HttpSys/test/FunctionalTests/DelegateTests.cs
index 9c5b64b5f0..b3db093023 100644
--- a/src/Servers/HttpSys/test/FunctionalTests/DelegateTests.cs
+++ b/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;
 
@@ -21,13 +19,13 @@ public class DelegateTests
     {
         var queueName = Guid.NewGuid().ToString();
         using var receiver = Utilities.CreateHttpServer(out var receiverAddress, async httpContext =>
-       {
-           await httpContext.Response.WriteAsync(_expectedResponseString);
-       },
-       options =>
-       {
-           options.RequestQueueName = queueName;
-       });
+        {
+            await httpContext.Response.WriteAsync(_expectedResponseString);
+        },
+        options =>
+        {
+            options.RequestQueueName = queueName;
+        });
 
         DelegationRule destination = default;
 
@@ -198,6 +196,72 @@ public class DelegateTests
         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();