Request can be affinitized to destinations (#174)

Request can be affinitized to destinations by an affinity key extracted from cookie or custom header. Affinity is active only when load balancing is enabled and it can be configured for each backend independently.
One request can have affinity to a single destination or to a set. It enables scenarios where a session needs to be affinitized to a pool of load-balanced destinations.
Affinity failures are handled by a policy configured per backend which has full access to `HttpContext` and decides whether the request processing can continue or not.

Fixes #45

samples/ReverseProxy.Sample/Startup.cs

@@ -62,7 +62,9 @@ public void Configure(IApplicationBuilder app)
 
                         return next();
                     });
+                    proxyPipeline.UseAffinitizedDestinationLookup();
                     proxyPipeline.UseProxyLoadBalancing();
+                    proxyPipeline.UseRequestAffinitizer();
                 });
             });
         }

samples/ReverseProxy.Sample/appsettings.json

@@ -20,6 +20,10 @@
         "LoadBalancing": {
           "Mode": "Random"
         },
+        "SessionAffinity": {
+          "Enabled": "true",
+          "Mode": "Cookie"
+        },
         "Metadata": {
           "CustomHealth": "false"
         },

samples/SampleClient/Program.cs

@@ -1,4 +1,4 @@
-// Copyright (c) Microsoft Corporation.
+// Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
 using System;
@@ -34,7 +34,8 @@ public static async Task<int> Main(string[] args)
                 {"Http1", () => new Http1Scenario()},
                 {"Http2", () => new Http2Scenario()},
                 {"RawUpgrade", () => new RawUpgradeScenario()},
-                {"WebSockets", () => new WebSocketsScenario()}
+                {"WebSockets", () => new WebSocketsScenario()},
+                {"SessionAffinity", () => new SessionAffinityScenario()}
             };
 
             if (string.IsNullOrEmpty(parsedArgs.Scenario))

samples/SampleClient/Scenarios/SessionAffinityScenario.cs

@@ -0,0 +1,84 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System;
+using System.Collections.Generic;
+using System.Diagnostics;
+using System.Linq;
+using System.Net;
+using System.Net.Http;
+using System.Text;
+using System.Text.Json;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace SampleClient.Scenarios
+{
+    internal class SessionAffinityScenario : IScenario
+    {
+        public async Task ExecuteAsync(CommandLineArgs args, CancellationToken cancellation)
+        {
+            using var handler = new HttpClientHandler
+            {
+                AllowAutoRedirect = false,
+                AutomaticDecompression = DecompressionMethods.None,
+                // Session affinity key will be stored in a cookie
+                UseCookies = true,
+                UseProxy = false
+            };
+            using var client = new HttpMessageInvoker(handler);
+            var targetUri = new Uri(new Uri(args.Target, UriKind.Absolute), "api/dump");
+            var stopwatch = Stopwatch.StartNew();
+
+            var request0 = new HttpRequestMessage(HttpMethod.Get, targetUri) { Version = new Version(1, 1) };
+            Console.WriteLine($"Sending first request to {targetUri} with HTTP/1.1");
+            var response0 = await client.SendAsync(request0, cancellation);
+
+            PrintDuration(stopwatch, response0);
+            PrintAffinityCookie(handler, targetUri, response0);
+            await ReadAndPrintBody(response0, cancellation);
+
+            stopwatch.Reset();
+
+            var request1 = new HttpRequestMessage(HttpMethod.Get, targetUri) { Version = new Version(1, 1) };
+            Console.WriteLine($"Sending second request to {targetUri} with HTTP/1.1");
+            var response1 = await client.SendAsync(request1, cancellation);
+
+            PrintDuration(stopwatch, response1);
+            PrintAffinityCookie(handler, targetUri, response1);
+            await ReadAndPrintBody(response1, cancellation);
+        }
+
+        private static void PrintDuration(Stopwatch stopwatch, HttpResponseMessage response)
+        {
+            Console.WriteLine($"Received response: {(int)response.StatusCode} in {stopwatch.ElapsedMilliseconds} ms");
+            response.EnsureSuccessStatusCode();
+        }
+
+        private static void PrintAffinityCookie(HttpClientHandler handler, Uri targetUri, HttpResponseMessage response)
+        {
+            if (response.Headers.TryGetValues("Set-Cookie", out var setCookieValue))
+            {
+                Console.WriteLine($"Received header Set-Cookie: {setCookieValue.ToArray()[0]}");
+            }
+            else
+            {
+                Console.WriteLine($"Response doesn't have Set-Cookie header.");
+            }
+
+            var affinityCookie = handler.CookieContainer.GetCookies(targetUri)[".Microsoft.ReverseProxy.Affinity"];
+            Console.WriteLine($"Affinity key stored on a cookie {affinityCookie.Value}");
+        }
+
+        private static async Task ReadAndPrintBody(HttpResponseMessage response, CancellationToken cancellation)
+        {
+            var body = await response.Content.ReadAsStringAsync(cancellation);
+            var json = JsonDocument.Parse(body);
+            Console.WriteLine(
+                "Received response:" +
+                $"{Environment.NewLine}" +
+                $"{JsonSerializer.Serialize(json.RootElement, new JsonSerializerOptions { WriteIndented = true })}");
+            response.EnsureSuccessStatusCode();
+        }
+    }
+}

src/ReverseProxy/Abstractions/BackendDiscovery/Contract/Backend.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using Microsoft.ReverseProxy.Abstractions.BackendDiscovery.Contract;
 
 namespace Microsoft.ReverseProxy.Abstractions
 {
@@ -39,6 +40,11 @@ public sealed class Backend : IDeepCloneable<Backend>
         /// </summary>
         public LoadBalancingOptions LoadBalancing { get; set; }
 
+        /// <summary>
+        /// Session affinity options.
+        /// </summary>
+        public SessionAffinityOptions SessionAffinity { get; set; }
+
         /// <summary>
         /// Active health checking options.
         /// </summary>
@@ -64,6 +70,7 @@ Backend IDeepCloneable<Backend>.DeepClone()
                 QuotaOptions = QuotaOptions?.DeepClone(),
                 PartitioningOptions = PartitioningOptions?.DeepClone(),
                 LoadBalancing = LoadBalancing?.DeepClone(),
+                SessionAffinity = SessionAffinity?.DeepClone(),
                 HealthCheckOptions = HealthCheckOptions?.DeepClone(),
                 Destinations = Destinations.DeepClone(StringComparer.Ordinal),
                 Metadata = Metadata?.DeepClone(StringComparer.Ordinal),

src/ReverseProxy/Abstractions/BackendDiscovery/Contract/CookieSessionAffinityProviderOptions.cs

@@ -0,0 +1,42 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System;
+using Microsoft.AspNetCore.Http;
+
+namespace Microsoft.ReverseProxy.Abstractions.BackendDiscovery.Contract
+{
+    /// <summary>
+    /// Defines cookie-specific affinity provider options.
+    /// </summary>
+    public class CookieSessionAffinityProviderOptions
+    {
+        private CookieBuilder _cookieBuilder = new AffinityCookieBuilder();
+
+        public static readonly string DefaultCookieName = ".Microsoft.ReverseProxy.Affinity";
+
+        public CookieBuilder Cookie
+        {
+            get => _cookieBuilder;
+            set => _cookieBuilder = value ?? throw new ArgumentNullException(nameof(value));
+        }
+
+        private class AffinityCookieBuilder : CookieBuilder
+        {
+            public AffinityCookieBuilder()
+            {
+                Name = DefaultCookieName;
+                SecurePolicy = CookieSecurePolicy.None;
+                SameSite = SameSiteMode.Unspecified;
+                HttpOnly = true;
+                IsEssential = false;
+            }
+
+            public override TimeSpan? Expiration
+            {
+                get => null;
+                set => throw new InvalidOperationException(nameof(Expiration) + " cannot be set for the cookie defined by " + nameof(CookieSessionAffinityProviderOptions));
+            }
+        }
+    }
+}

src/ReverseProxy/Abstractions/BackendDiscovery/Contract/SessionAffinityConstants.cs

@@ -0,0 +1,25 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+namespace Microsoft.ReverseProxy.Abstractions.BackendDiscovery.Contract
+{
+    /// <summary>
+    /// Names of built-in session affinity services.
+    /// </summary>
+    public static class SessionAffinityConstants
+    {
+        public static class Modes
+        {
+            public static string Cookie => "Cookie";
+
+            public static string CustomHeader => "CustomHeader";
+        }
+
+        public static class AffinityFailurePolicies
+        {
+            public static string Redistribute => "Redistribute";
+
+            public static string Return503Error => "Return503Error";
+        }
+    }
+}

src/ReverseProxy/Abstractions/BackendDiscovery/Contract/SessionAffinityOptions.cs

@@ -0,0 +1,45 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System;
+using System.Collections.Generic;
+
+namespace Microsoft.ReverseProxy.Abstractions.BackendDiscovery.Contract
+{
+    /// <summary>
+    /// Session affinitity options.
+    /// </summary>
+    public sealed class SessionAffinityOptions
+    {
+        /// <summary>
+        /// Indicates whether session affinity is enabled.
+        /// </summary>
+        public bool Enabled { get; set; }
+
+        /// <summary>
+        /// Session affinity mode which is implemented by one of providers.
+        /// </summary>
+        public string Mode { get; set; }
+
+        /// <summary>
+        /// Strategy handling missing destination for an affinitized request.
+        /// </summary>
+        public string AffinityFailurePolicy { get; set; }
+
+        /// <summary>
+        /// Key-value pair collection holding extra settings specific to different affinity modes.
+        /// </summary>
+        public IDictionary<string, string> Settings { get; set; }
+
+        internal SessionAffinityOptions DeepClone()
+        {
+            return new SessionAffinityOptions
+            {
+                Enabled = Enabled,
+                Mode = Mode,
+                AffinityFailurePolicy = AffinityFailurePolicy,
+                Settings = Settings?.DeepClone(StringComparer.Ordinal)
+            };
+        }
+    }
+}

src/ReverseProxy/Configuration/DependencyInjection/BuilderExtensions/IReverseProxyBuilderExtensions.cs

@@ -1,15 +1,20 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+using System;
+using Microsoft.AspNetCore.DataProtection;
+using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.DependencyInjection.Extensions;
 using Microsoft.ReverseProxy.Abstractions;
+using Microsoft.ReverseProxy.Abstractions.BackendDiscovery.Contract;
 using Microsoft.ReverseProxy.Abstractions.Telemetry;
 using Microsoft.ReverseProxy.Abstractions.Time;
 using Microsoft.ReverseProxy.Service;
 using Microsoft.ReverseProxy.Service.Management;
 using Microsoft.ReverseProxy.Service.Metrics;
 using Microsoft.ReverseProxy.Service.Proxy;
 using Microsoft.ReverseProxy.Service.Proxy.Infrastructure;
+using Microsoft.ReverseProxy.Service.SessionAffinity;
 using Microsoft.ReverseProxy.Telemetry;
 using Microsoft.ReverseProxy.Utilities;
 
@@ -82,5 +87,19 @@ public static IReverseProxyBuilder AddBackgroundWorkers(this IReverseProxyBuilde
 
             return builder;
         }
+
+        public static IReverseProxyBuilder AddSessionAffinityProvider(this IReverseProxyBuilder builder)
+        {
+            builder.Services.TryAddEnumerable(new[] {
+                new ServiceDescriptor(typeof(IAffinityFailurePolicy), typeof(RedistributeAffinityFailurePolicy), ServiceLifetime.Singleton),
+                new ServiceDescriptor(typeof(IAffinityFailurePolicy), typeof(Return503ErrorAffinityFailurePolicy), ServiceLifetime.Singleton)
+            });
+            builder.Services.TryAddEnumerable(new[] {
+                new ServiceDescriptor(typeof(ISessionAffinityProvider), typeof(CookieSessionAffinityProvider), ServiceLifetime.Singleton),
+                new ServiceDescriptor(typeof(ISessionAffinityProvider), typeof(CustomHeaderSessionAffinityProvider), ServiceLifetime.Singleton)
+            });
+
+            return builder;
+        }
     }
 }

src/ReverseProxy/Configuration/DependencyInjection/ReverseProxyServiceCollectionExtensions.cs

@@ -29,9 +29,12 @@ public static IReverseProxyBuilder AddReverseProxy(this IServiceCollection servi
                 .AddRuntimeStateManagers()
                 .AddConfigManager()
                 .AddDynamicEndpointDataSource()
+                .AddSessionAffinityProvider()
                 .AddProxy()
                 .AddBackgroundWorkers();
 
+            services.AddDataProtection();
+
             return builder;
         }
 

src/ReverseProxy/Configuration/ReverseProxyIEndpointRouteBuilderExtensions.cs

@@ -22,7 +22,9 @@ public static void MapReverseProxy(this IEndpointRouteBuilder endpoints)
         {
             endpoints.MapReverseProxy(app =>
             {
+                app.UseAffinitizedDestinationLookup();
                 app.UseProxyLoadBalancing();
+                app.UseRequestAffinitizer();
             });
         }
 

src/ReverseProxy/EventIds.cs

@@ -40,5 +40,13 @@ internal static class EventIds
         public static readonly EventId OperationStarted = new EventId(31, "OperationStarted");
         public static readonly EventId OperationEnded = new EventId(32, "OperationEnded");
         public static readonly EventId OperationFailed = new EventId(33, "OperationFailed");
+        public static readonly EventId AffinityResolutionFailedForBackend = new EventId(34, "AffinityResolutionFailedForBackend");
+        public static readonly EventId MultipleDestinationsOnBackendToEstablishRequestAffinity = new EventId(35, "MultipleDestinationsOnBackendToEstablishRequestAffinity");
+        public static readonly EventId AffinityCannotBeEstablishedBecauseNoDestinationsFoundOnBackend = new EventId(36, "AffinityCannotBeEstablishedBecauseNoDestinationsFoundOnBackend");
+        public static readonly EventId NoDestinationOnBackendToEstablishRequestAffinity = new EventId(37, "NoDestinationOnBackendToEstablishRequestAffinity");
+        public static readonly EventId RequestAffinityKeyDecryptionFailed = new EventId(38, "RequestAffinityKeyDecryptionFailed");
+        public static readonly EventId DestinationMatchingToAffinityKeyNotFound = new EventId(39, "DestinationMatchingToAffinityKeyNotFound");
+        public static readonly EventId RequestAffinityHeaderHasMultipleValues = new EventId(40, "RequestAffinityHeaderHasMultipleValues");
+        public static readonly EventId AffinityResolutionFailureWasHandledProcessingWillBeContinued = new EventId(41, "AffinityResolutionFailureWasHandledProcessingWillBeContinued");
     }
 }