diff --git a/samples/ReverseProxy.Sample/Startup.cs b/samples/ReverseProxy.Sample/Startup.cs index a74a197b9..42f9aa178 100644 --- a/samples/ReverseProxy.Sample/Startup.cs +++ b/samples/ReverseProxy.Sample/Startup.cs @@ -62,7 +62,9 @@ public void Configure(IApplicationBuilder app) return next(); }); + proxyPipeline.UseAffinitizedDestinationLookup(); proxyPipeline.UseProxyLoadBalancing(); + proxyPipeline.UseRequestAffinitizer(); }); }); } diff --git a/samples/ReverseProxy.Sample/appsettings.json b/samples/ReverseProxy.Sample/appsettings.json index c42715bbb..05e80defd 100644 --- a/samples/ReverseProxy.Sample/appsettings.json +++ b/samples/ReverseProxy.Sample/appsettings.json @@ -20,6 +20,10 @@ "LoadBalancing": { "Mode": "Random" }, + "SessionAffinity": { + "Enabled": "true", + "Mode": "Cookie" + }, "Metadata": { "CustomHealth": "false" }, diff --git a/samples/SampleClient/Program.cs b/samples/SampleClient/Program.cs index cb9690e10..f79d20941 100644 --- a/samples/SampleClient/Program.cs +++ b/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 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)) diff --git a/samples/SampleClient/Scenarios/SessionAffinityScenario.cs b/samples/SampleClient/Scenarios/SessionAffinityScenario.cs new file mode 100644 index 000000000..77ba4e2fe --- /dev/null +++ b/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(); + } + } +} diff --git a/src/ReverseProxy/Abstractions/BackendDiscovery/Contract/Backend.cs b/src/ReverseProxy/Abstractions/BackendDiscovery/Contract/Backend.cs index b2cf47d74..a74b59ef2 100644 --- a/src/ReverseProxy/Abstractions/BackendDiscovery/Contract/Backend.cs +++ b/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 /// public LoadBalancingOptions LoadBalancing { get; set; } + /// + /// Session affinity options. + /// + public SessionAffinityOptions SessionAffinity { get; set; } + /// /// Active health checking options. /// @@ -64,6 +70,7 @@ Backend IDeepCloneable.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), diff --git a/src/ReverseProxy/Abstractions/BackendDiscovery/Contract/CookieSessionAffinityProviderOptions.cs b/src/ReverseProxy/Abstractions/BackendDiscovery/Contract/CookieSessionAffinityProviderOptions.cs new file mode 100644 index 000000000..5d89c27f1 --- /dev/null +++ b/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 +{ + /// + /// Defines cookie-specific affinity provider options. + /// + 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)); + } + } + } +} diff --git a/src/ReverseProxy/Abstractions/BackendDiscovery/Contract/SessionAffinityConstants.cs b/src/ReverseProxy/Abstractions/BackendDiscovery/Contract/SessionAffinityConstants.cs new file mode 100644 index 000000000..2bff0e643 --- /dev/null +++ b/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 +{ + /// + /// Names of built-in session affinity services. + /// + 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"; + } + } +} diff --git a/src/ReverseProxy/Abstractions/BackendDiscovery/Contract/SessionAffinityOptions.cs b/src/ReverseProxy/Abstractions/BackendDiscovery/Contract/SessionAffinityOptions.cs new file mode 100644 index 000000000..11212c7cc --- /dev/null +++ b/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 +{ + /// + /// Session affinitity options. + /// + public sealed class SessionAffinityOptions + { + /// + /// Indicates whether session affinity is enabled. + /// + public bool Enabled { get; set; } + + /// + /// Session affinity mode which is implemented by one of providers. + /// + public string Mode { get; set; } + + /// + /// Strategy handling missing destination for an affinitized request. + /// + public string AffinityFailurePolicy { get; set; } + + /// + /// Key-value pair collection holding extra settings specific to different affinity modes. + /// + public IDictionary Settings { get; set; } + + internal SessionAffinityOptions DeepClone() + { + return new SessionAffinityOptions + { + Enabled = Enabled, + Mode = Mode, + AffinityFailurePolicy = AffinityFailurePolicy, + Settings = Settings?.DeepClone(StringComparer.Ordinal) + }; + } + } +} diff --git a/src/ReverseProxy/Configuration/DependencyInjection/BuilderExtensions/IReverseProxyBuilderExtensions.cs b/src/ReverseProxy/Configuration/DependencyInjection/BuilderExtensions/IReverseProxyBuilderExtensions.cs index 8f20ddfa6..c50b0b1e7 100644 --- a/src/ReverseProxy/Configuration/DependencyInjection/BuilderExtensions/IReverseProxyBuilderExtensions.cs +++ b/src/ReverseProxy/Configuration/DependencyInjection/BuilderExtensions/IReverseProxyBuilderExtensions.cs @@ -1,8 +1,12 @@ // 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; @@ -10,6 +14,7 @@ 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; + } } } diff --git a/src/ReverseProxy/Configuration/DependencyInjection/ReverseProxyServiceCollectionExtensions.cs b/src/ReverseProxy/Configuration/DependencyInjection/ReverseProxyServiceCollectionExtensions.cs index d45d9b9db..96cd6b64b 100644 --- a/src/ReverseProxy/Configuration/DependencyInjection/ReverseProxyServiceCollectionExtensions.cs +++ b/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; } diff --git a/src/ReverseProxy/Configuration/ReverseProxyIEndpointRouteBuilderExtensions.cs b/src/ReverseProxy/Configuration/ReverseProxyIEndpointRouteBuilderExtensions.cs index beffa6a6d..d79081adc 100644 --- a/src/ReverseProxy/Configuration/ReverseProxyIEndpointRouteBuilderExtensions.cs +++ b/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(); }); } diff --git a/src/ReverseProxy/EventIds.cs b/src/ReverseProxy/EventIds.cs index 683177ffc..c9645858b 100644 --- a/src/ReverseProxy/EventIds.cs +++ b/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"); } } diff --git a/src/ReverseProxy/Middleware/AffinitizeRequestMiddleware.cs b/src/ReverseProxy/Middleware/AffinitizeRequestMiddleware.cs new file mode 100644 index 000000000..9b18dae56 --- /dev/null +++ b/src/ReverseProxy/Middleware/AffinitizeRequestMiddleware.cs @@ -0,0 +1,101 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using System.Collections.Generic; +using System.Threading.Tasks; +using Microsoft.AspNetCore.Http; +using Microsoft.Extensions.Logging; +using Microsoft.ReverseProxy.Abstractions.Telemetry; +using Microsoft.ReverseProxy.RuntimeModel; +using Microsoft.ReverseProxy.Service.SessionAffinity; + +namespace Microsoft.ReverseProxy.Middleware +{ + /// + /// Affinitizes request to a chosen . + /// + internal class AffinitizeRequestMiddleware + { + private readonly Random _random = new Random(); + private readonly RequestDelegate _next; + private readonly IDictionary _sessionAffinityProviders; + private readonly IOperationLogger _operationLogger; + private readonly ILogger _logger; + + public AffinitizeRequestMiddleware( + RequestDelegate next, + IEnumerable sessionAffinityProviders, + IOperationLogger operationLogger, + ILogger logger) + { + _next = next ?? throw new ArgumentNullException(nameof(next)); + _logger = logger ?? throw new ArgumentNullException(nameof(logger)); + _sessionAffinityProviders = sessionAffinityProviders.ToProviderDictionary(); + _operationLogger = operationLogger ?? throw new ArgumentNullException(nameof(operationLogger)); + } + + public Task Invoke(HttpContext context) + { + var backend = context.GetRequiredBackend(); + var options = backend.Config.Value?.SessionAffinityOptions ?? default; + + if (options.Enabled) + { + var destinationsFeature = context.GetRequiredDestinationFeature(); + var candidateDestinations = destinationsFeature.Destinations; + + if (candidateDestinations.Count == 0) + { + // Only log the warning about missing destinations here, but allow the request to proceed further. + // The final check for selected destination is to be done at the pipeline end. + Log.NoDestinationOnBackendToEstablishRequestAffinity(_logger, backend.BackendId); + } + else + { + var chosenDestination = candidateDestinations[0]; + if (candidateDestinations.Count > 1) + { + Log.MultipleDestinationsOnBackendToEstablishRequestAffinity(_logger, backend.BackendId); + // It's assumed that all of them match to the request's affinity key. + chosenDestination = candidateDestinations[_random.Next(candidateDestinations.Count)]; + destinationsFeature.Destinations = chosenDestination; + } + + _operationLogger.Execute("ReverseProxy.AffinitizeRequest", () => AffinitizeRequest(context, options, chosenDestination)); + } + } + + return _next(context); + } + + private void AffinitizeRequest(HttpContext context, BackendConfig.BackendSessionAffinityOptions options, DestinationInfo destination) + { + var currentProvider = _sessionAffinityProviders.GetRequiredServiceById(options.Mode); + currentProvider.AffinitizeRequest(context, options, destination); + } + + private static class Log + { + private static readonly Action _multipleDestinationsOnBackendToEstablishRequestAffinity = LoggerMessage.Define( + LogLevel.Warning, + EventIds.MultipleDestinationsOnBackendToEstablishRequestAffinity, + "The request still has multiple destinations on the backend `{backendId}` to choose from when establishing affinity, load balancing may not be properly configured. A random destination will be used."); + + private static readonly Action _noDestinationOnBackendToEstablishRequestAffinity = LoggerMessage.Define( + LogLevel.Warning, + EventIds.NoDestinationOnBackendToEstablishRequestAffinity, + "The request doesn't have any destinations on the backend `{backendId}` to choose from when establishing affinity, load balancing may not be properly configured."); + + public static void MultipleDestinationsOnBackendToEstablishRequestAffinity(ILogger logger, string backendId) + { + _multipleDestinationsOnBackendToEstablishRequestAffinity(logger, backendId, null); + } + + public static void NoDestinationOnBackendToEstablishRequestAffinity(ILogger logger, string backendId) + { + _noDestinationOnBackendToEstablishRequestAffinity(logger, backendId, null); + } + } + } +} diff --git a/src/ReverseProxy/Middleware/AffinitizedDestinationLookupMiddleware.cs b/src/ReverseProxy/Middleware/AffinitizedDestinationLookupMiddleware.cs new file mode 100644 index 000000000..ec5776614 --- /dev/null +++ b/src/ReverseProxy/Middleware/AffinitizedDestinationLookupMiddleware.cs @@ -0,0 +1,125 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using System.Collections.Generic; +using System.Threading.Tasks; +using Microsoft.AspNetCore.Authorization; +using Microsoft.AspNetCore.Http; +using Microsoft.Extensions.Logging; +using Microsoft.ReverseProxy.Abstractions; +using Microsoft.ReverseProxy.Abstractions.Telemetry; +using Microsoft.ReverseProxy.RuntimeModel; +using Microsoft.ReverseProxy.Service.SessionAffinity; + +namespace Microsoft.ReverseProxy.Middleware +{ + /// + /// Looks up an affinitized matching the request's affinity key if any is set + /// + internal class AffinitizedDestinationLookupMiddleware + { + private readonly RequestDelegate _next; + private readonly IDictionary _sessionAffinityProviders; + private readonly IDictionary _affinityFailurePolicies; + private readonly IOperationLogger _operationLogger; + private readonly ILogger _logger; + + public AffinitizedDestinationLookupMiddleware( + RequestDelegate next, + IEnumerable sessionAffinityProviders, + IEnumerable affinityFailurePolicies, + IOperationLogger operationLogger, + ILogger logger) + { + _next = next ?? throw new ArgumentNullException(nameof(next)); + _logger = logger ?? throw new ArgumentNullException(nameof(logger)); + _operationLogger = operationLogger ?? throw new ArgumentNullException(nameof(logger)); + _sessionAffinityProviders = sessionAffinityProviders.ToProviderDictionary(); + _affinityFailurePolicies = affinityFailurePolicies?.ToPolicyDictionary() ?? throw new ArgumentNullException(nameof(affinityFailurePolicies)); + } + + public Task Invoke(HttpContext context) + { + var backend = context.GetRequiredBackend(); + + var options = backend.Config.Value?.SessionAffinityOptions ?? default; + + if (!options.Enabled) + { + return _next(context); + } + + return InvokeInternal(context, options, backend); + } + + private async Task InvokeInternal(HttpContext context, BackendConfig.BackendSessionAffinityOptions options, BackendInfo backend) + { + var destinationsFeature = context.GetRequiredDestinationFeature(); + var destinations = destinationsFeature.Destinations; + + var affinityResult = _operationLogger.Execute( + "ReverseProxy.FindAffinitizedDestinations", + () => + { + var currentProvider = _sessionAffinityProviders.GetRequiredServiceById(options.Mode); + return currentProvider.FindAffinitizedDestinations(context, destinations, backend.BackendId, options); + }); + switch (affinityResult.Status) + { + case AffinityStatus.OK: + destinationsFeature.Destinations = affinityResult.Destinations; + break; + case AffinityStatus.AffinityKeyNotSet: + // Nothing to do so just continue processing + break; + case AffinityStatus.AffinityKeyExtractionFailed: + case AffinityStatus.DestinationNotFound: + var keepProcessing = await _operationLogger.ExecuteAsync("ReverseProxy.HandleAffinityFailure", () => + { + var failurePolicy = _affinityFailurePolicies.GetRequiredServiceById(options.AffinityFailurePolicy); + return failurePolicy.Handle(context, options, affinityResult.Status); + }); + + if (!keepProcessing) + { + // Policy reported the failure is unrecoverable and took the full responsibility for its handling, + // so we simply stop processing. + Log.AffinityResolutionFailedForBackend(_logger, backend.BackendId); + return; + } + + Log.AffinityResolutionFailureWasHandledProcessingWillBeContinued(_logger, backend.BackendId, options.AffinityFailurePolicy); + + break; + default: + throw new NotSupportedException($"Affinity status '{affinityResult.Status}' is not supported."); + } + + await _next(context); + } + + private static class Log + { + private static readonly Action _affinityResolutioFailedForBackend = LoggerMessage.Define( + LogLevel.Warning, + EventIds.AffinityResolutionFailedForBackend, + "Affinity resolution failed for backend `{backendId}`."); + + private static readonly Action _affinityResolutionFailureWasHandledProcessingWillBeContinued = LoggerMessage.Define( + LogLevel.Debug, + EventIds.AffinityResolutionFailureWasHandledProcessingWillBeContinued, + "Affinity resolution failure for backend `{backendId}` was handled successfully by the policy `{policyName}`. Request processing will be continued."); + + public static void AffinityResolutionFailedForBackend(ILogger logger, string backendId) + { + _affinityResolutioFailedForBackend(logger, backendId, null); + } + + public static void AffinityResolutionFailureWasHandledProcessingWillBeContinued(ILogger logger, string backendId, string policyName) + { + _affinityResolutionFailureWasHandledProcessingWillBeContinued(logger, backendId, policyName, null); + } + } + } +} diff --git a/src/ReverseProxy/Middleware/HttpContextFeaturesExtensions.cs b/src/ReverseProxy/Middleware/HttpContextFeaturesExtensions.cs new file mode 100644 index 000000000..a3fef3609 --- /dev/null +++ b/src/ReverseProxy/Middleware/HttpContextFeaturesExtensions.cs @@ -0,0 +1,27 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using Microsoft.AspNetCore.Http; +using Microsoft.ReverseProxy.RuntimeModel; + +namespace Microsoft.ReverseProxy.Middleware +{ + internal static class HttpContextFeaturesExtensions + { + public static BackendInfo GetRequiredBackend(this HttpContext context) + { + return context.Features.Get() ?? throw new InvalidOperationException("Backend unspecified."); + } + + public static IAvailableDestinationsFeature GetRequiredDestinationFeature(this HttpContext context) + { + var destinationsFeature = context.Features.Get(); + if (destinationsFeature?.Destinations == null) + { + throw new InvalidOperationException("The IAvailableDestinationsFeature Destinations collection was not set."); + } + return destinationsFeature; + } + } +} diff --git a/src/ReverseProxy/Middleware/LoadBalancingMiddleware.cs b/src/ReverseProxy/Middleware/LoadBalancingMiddleware.cs index 2b343458b..7b5353d6f 100644 --- a/src/ReverseProxy/Middleware/LoadBalancingMiddleware.cs +++ b/src/ReverseProxy/Middleware/LoadBalancingMiddleware.cs @@ -35,10 +35,9 @@ public LoadBalancingMiddleware( public Task Invoke(HttpContext context) { - var backend = context.Features.Get() ?? throw new InvalidOperationException("Backend unspecified."); - var destinationsFeature = context.Features.Get(); - var destinations = destinationsFeature?.Destinations - ?? throw new InvalidOperationException("The IAvailableDestinationsFeature Destinations collection was not set."); + var backend = context.GetRequiredBackend(); + var destinationsFeature = context.GetRequiredDestinationFeature(); + var destinations = destinationsFeature.Destinations; var loadBalancingOptions = backend.Config.Value?.LoadBalancingOptions ?? new BackendConfig.BackendLoadBalancingOptions(default); diff --git a/src/ReverseProxy/Middleware/ProxyMiddlewareAppBuilderExtensions.cs b/src/ReverseProxy/Middleware/ProxyMiddlewareAppBuilderExtensions.cs index 7fbadf5fe..5c5285a7c 100644 --- a/src/ReverseProxy/Middleware/ProxyMiddlewareAppBuilderExtensions.cs +++ b/src/ReverseProxy/Middleware/ProxyMiddlewareAppBuilderExtensions.cs @@ -17,5 +17,26 @@ public static IApplicationBuilder UseProxyLoadBalancing(this IApplicationBuilder { return builder.UseMiddleware(); } + + /// + /// Checks if a request has an established affinity relationship and if the associated destination is available. + /// This should be placed before load balancing and other destination selection components. + /// Requests without an affinity relationship will be a processed normally and have the affinity relationship + /// established by a later component. See . + /// + public static IApplicationBuilder UseAffinitizedDestinationLookup(this IApplicationBuilder builder) + { + return builder.UseMiddleware(); + } + + /// + /// Establishes the affinity relationship to the selected destination. + /// If there are multiple affinitized destinations found for the request, it randomly picks one of them. + /// This should be placed after load balancing and other destination selection processes. + /// + public static IApplicationBuilder UseRequestAffinitizer(this IApplicationBuilder builder) + { + return builder.UseMiddleware(); + } } } diff --git a/src/ReverseProxy/Service/Config/ConfigErrors.cs b/src/ReverseProxy/Service/Config/ConfigErrors.cs index d98cb30e3..f75ccd9d2 100644 --- a/src/ReverseProxy/Service/Config/ConfigErrors.cs +++ b/src/ReverseProxy/Service/Config/ConfigErrors.cs @@ -20,6 +20,8 @@ internal static class ConfigErrors internal const string ParsedRouteRuleInvalidMatcher = "ParsedRoute_RuleInvalidMatcher"; internal const string ConfigBuilderBackendIdMismatch = "ConfigBuilder_BackendIdMismatch"; + internal const string ConfigBuilderBackendNoProviderFoundForSessionAffinityMode = "ConfigBuilder_BackendNoProviderFoundForSessionAffinityMode"; + internal const string ConfigBuilderBackendNoAffinityFailurePolicyFoundForSpecifiedName = "ConfigBuilder_NoAffinityFailurePolicyFoundForSpecifiedName"; internal const string ConfigBuilderBackendException = "ConfigBuilder_BackendException"; internal const string ConfigBuilderRouteException = "ConfigBuilder_RouteException"; } diff --git a/src/ReverseProxy/Service/Config/DynamicConfigBuilder.cs b/src/ReverseProxy/Service/Config/DynamicConfigBuilder.cs index 7c7c0a4fb..3213c67a1 100644 --- a/src/ReverseProxy/Service/Config/DynamicConfigBuilder.cs +++ b/src/ReverseProxy/Service/Config/DynamicConfigBuilder.cs @@ -3,12 +3,13 @@ using System; using System.Collections.Generic; -using System.Diagnostics; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.Options; using Microsoft.ReverseProxy.Abstractions; +using Microsoft.ReverseProxy.Abstractions.BackendDiscovery.Contract; using Microsoft.ReverseProxy.ConfigModel; +using Microsoft.ReverseProxy.Service.SessionAffinity; using Microsoft.ReverseProxy.Utilities; namespace Microsoft.ReverseProxy.Service @@ -19,21 +20,29 @@ internal class DynamicConfigBuilder : IDynamicConfigBuilder private readonly IBackendsRepo _backendsRepo; private readonly IRoutesRepo _routesRepo; private readonly IRouteValidator _parsedRouteValidator; + private readonly IDictionary _sessionAffinityProviders; + private readonly IDictionary _affinityFailurePolicies; public DynamicConfigBuilder( IEnumerable filters, IBackendsRepo backendsRepo, IRoutesRepo routesRepo, - IRouteValidator parsedRouteValidator) + IRouteValidator parsedRouteValidator, + IEnumerable sessionAffinityProviders, + IEnumerable affinityFailurePolicies) { Contracts.CheckValue(filters, nameof(filters)); Contracts.CheckValue(backendsRepo, nameof(backendsRepo)); Contracts.CheckValue(routesRepo, nameof(routesRepo)); Contracts.CheckValue(parsedRouteValidator, nameof(parsedRouteValidator)); + Contracts.CheckValue(sessionAffinityProviders, nameof(sessionAffinityProviders)); + Contracts.CheckValue(affinityFailurePolicies, nameof(affinityFailurePolicies)); _filters = filters; _backendsRepo = backendsRepo; _routesRepo = routesRepo; _parsedRouteValidator = parsedRouteValidator; + _sessionAffinityProviders = sessionAffinityProviders.ToProviderDictionary(); + _affinityFailurePolicies = affinityFailurePolicies.ToPolicyDictionary(); } public async Task> BuildConfigAsync(IConfigErrorReporter errorReporter, CancellationToken cancellation) @@ -73,6 +82,8 @@ public async Task> GetBackendsAsync(IConfigErrorRep await filter.ConfigureBackendAsync(backend, cancellation); } + ValidateSessionAffinity(errorReporter, id, backend); + configuredBackends[id] = backend; } catch (Exception ex) @@ -84,6 +95,37 @@ public async Task> GetBackendsAsync(IConfigErrorRep return configuredBackends; } + private void ValidateSessionAffinity(IConfigErrorReporter errorReporter, string id, Backend backend) + { + if (backend.SessionAffinity == null || !backend.SessionAffinity.Enabled) + { + // Session affinity is disabled + return; + } + + if (string.IsNullOrEmpty(backend.SessionAffinity.Mode)) + { + backend.SessionAffinity.Mode = SessionAffinityConstants.Modes.Cookie; + } + + var affinityMode = backend.SessionAffinity.Mode; + if (!_sessionAffinityProviders.ContainsKey(affinityMode)) + { + errorReporter.ReportError(ConfigErrors.ConfigBuilderBackendNoProviderFoundForSessionAffinityMode, id, $"No matching {nameof(ISessionAffinityProvider)} found for the session affinity mode {affinityMode} set on the backend {backend.Id}."); + } + + if (string.IsNullOrEmpty(backend.SessionAffinity.AffinityFailurePolicy)) + { + backend.SessionAffinity.AffinityFailurePolicy = SessionAffinityConstants.AffinityFailurePolicies.Redistribute; + } + + var affinityFailurePolicy = backend.SessionAffinity.AffinityFailurePolicy; + if (!_affinityFailurePolicies.ContainsKey(affinityFailurePolicy)) + { + errorReporter.ReportError(ConfigErrors.ConfigBuilderBackendNoAffinityFailurePolicyFoundForSpecifiedName, id, $"No matching {nameof(IAffinityFailurePolicy)} found for the affinity failure policy name {affinityFailurePolicy} set on the backend {backend.Id}."); + } + } + private async Task> GetRoutesAsync(IConfigErrorReporter errorReporter, CancellationToken cancellation) { var routes = await _routesRepo.GetRoutesAsync(cancellation); diff --git a/src/ReverseProxy/Service/Management/ReverseProxyConfigManager.cs b/src/ReverseProxy/Service/Management/ReverseProxyConfigManager.cs index 1ae6cd0cc..30999d992 100644 --- a/src/ReverseProxy/Service/Management/ReverseProxyConfigManager.cs +++ b/src/ReverseProxy/Service/Management/ReverseProxyConfigManager.cs @@ -94,7 +94,12 @@ private void UpdateRuntimeBackends(DynamicConfigRoot config) port: configBackend.HealthCheckOptions?.Port ?? 0, path: configBackend.HealthCheckOptions?.Path ?? string.Empty), new BackendConfig.BackendLoadBalancingOptions( - mode: configBackend.LoadBalancing?.Mode ?? default)); + mode: configBackend.LoadBalancing?.Mode ?? default), + new BackendConfig.BackendSessionAffinityOptions( + enabled: configBackend.SessionAffinity?.Enabled ?? false, + mode: configBackend.SessionAffinity?.Mode, + affinityFailurePolicy: configBackend.SessionAffinity?.AffinityFailurePolicy, + settings: configBackend.SessionAffinity?.Settings as IReadOnlyDictionary)); var currentBackendConfig = backend.Config.Value; if (currentBackendConfig == null || diff --git a/src/ReverseProxy/Service/Proxy/HttpProxy.cs b/src/ReverseProxy/Service/Proxy/HttpProxy.cs index 8fb94b7b7..387cd2ea2 100644 --- a/src/ReverseProxy/Service/Proxy/HttpProxy.cs +++ b/src/ReverseProxy/Service/Proxy/HttpProxy.cs @@ -19,6 +19,7 @@ using Microsoft.Extensions.Primitives; using Microsoft.ReverseProxy.Service.Metrics; using Microsoft.ReverseProxy.Service.Proxy.Infrastructure; +using Microsoft.ReverseProxy.Service.SessionAffinity; using Microsoft.ReverseProxy.Utilities; namespace Microsoft.ReverseProxy.Service.Proxy @@ -450,7 +451,7 @@ static void CopyHeaders(HttpHeaders source, IHeaderDictionary destination) } ////this.logger.LogInformation($" Copying downstream <-- Proxy <-- upstream response header {header.Key}: {string.Join(",", header.Value)}"); - destination.TryAdd(header.Key, new StringValues(header.Value.ToArray())); + destination.Append(header.Key, new StringValues(header.Value.ToArray())); } } } diff --git a/src/ReverseProxy/Service/RuntimeModel/BackendConfig.cs b/src/ReverseProxy/Service/RuntimeModel/BackendConfig.cs index 387f4c900..b7e87fb72 100644 --- a/src/ReverseProxy/Service/RuntimeModel/BackendConfig.cs +++ b/src/ReverseProxy/Service/RuntimeModel/BackendConfig.cs @@ -2,6 +2,7 @@ // Licensed under the MIT License. using System; +using System.Collections.Generic; using Microsoft.ReverseProxy.Abstractions; using Microsoft.ReverseProxy.Utilities; @@ -17,20 +18,24 @@ namespace Microsoft.ReverseProxy.RuntimeModel /// Instead, instances of are replaced /// in ther entirety when values need to change. /// - internal sealed class BackendConfig + public sealed class BackendConfig { public BackendConfig( BackendHealthCheckOptions healthCheckOptions, - BackendLoadBalancingOptions loadBalancingOptions) + BackendLoadBalancingOptions loadBalancingOptions, + BackendSessionAffinityOptions sessionAffinityOptions) { HealthCheckOptions = healthCheckOptions; LoadBalancingOptions = loadBalancingOptions; + SessionAffinityOptions = sessionAffinityOptions; } public BackendHealthCheckOptions HealthCheckOptions { get; } public BackendLoadBalancingOptions LoadBalancingOptions { get; } + public BackendSessionAffinityOptions SessionAffinityOptions { get; } + /// /// Active health probing options for a backend. /// @@ -38,7 +43,7 @@ public BackendConfig( /// Struct used only to keep things organized as we add more configuration options inside of `BackendConfig`. /// Each "feature" can have its own struct. /// - internal readonly struct BackendHealthCheckOptions + public readonly struct BackendHealthCheckOptions { public BackendHealthCheckOptions(bool enabled, TimeSpan interval, TimeSpan timeout, int port, string path) { @@ -75,7 +80,7 @@ public BackendHealthCheckOptions(bool enabled, TimeSpan interval, TimeSpan timeo public string Path { get; } } - internal readonly struct BackendLoadBalancingOptions + public readonly struct BackendLoadBalancingOptions { public BackendLoadBalancingOptions(LoadBalancingMode mode) { @@ -88,5 +93,24 @@ public BackendLoadBalancingOptions(LoadBalancingMode mode) internal AtomicCounter RoundRobinState { get; } } + + public readonly struct BackendSessionAffinityOptions + { + public BackendSessionAffinityOptions(bool enabled, string mode, string affinityFailurePolicy, IReadOnlyDictionary settings) + { + Mode = mode; + AffinityFailurePolicy = affinityFailurePolicy; + Settings = settings; + Enabled = enabled; + } + + public bool Enabled { get; } + + public string Mode { get; } + + public string AffinityFailurePolicy { get; } + + public IReadOnlyDictionary Settings { get; } + } } } diff --git a/src/ReverseProxy/Service/SessionAffinity/AffinityResult.cs b/src/ReverseProxy/Service/SessionAffinity/AffinityResult.cs new file mode 100644 index 000000000..360d07a5b --- /dev/null +++ b/src/ReverseProxy/Service/SessionAffinity/AffinityResult.cs @@ -0,0 +1,24 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.Collections.Generic; +using Microsoft.ReverseProxy.RuntimeModel; + +namespace Microsoft.ReverseProxy.Service.SessionAffinity +{ + /// + /// Affinity resolution result. + /// + public readonly struct AffinityResult + { + public IReadOnlyList Destinations { get; } + + public AffinityStatus Status { get; } + + public AffinityResult(IReadOnlyList destinations, AffinityStatus status) + { + Destinations = destinations; + Status = status; + } + } +} diff --git a/src/ReverseProxy/Service/SessionAffinity/AffinityStatus.cs b/src/ReverseProxy/Service/SessionAffinity/AffinityStatus.cs new file mode 100644 index 000000000..b2d2a0529 --- /dev/null +++ b/src/ReverseProxy/Service/SessionAffinity/AffinityStatus.cs @@ -0,0 +1,16 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +namespace Microsoft.ReverseProxy.Service.SessionAffinity +{ + /// + /// Affinity resolution status. + /// + public enum AffinityStatus + { + OK, + AffinityKeyNotSet, + AffinityKeyExtractionFailed, + DestinationNotFound + } +} diff --git a/src/ReverseProxy/Service/SessionAffinity/BaseSessionAffinityProvider.cs b/src/ReverseProxy/Service/SessionAffinity/BaseSessionAffinityProvider.cs new file mode 100644 index 000000000..37aac3673 --- /dev/null +++ b/src/ReverseProxy/Service/SessionAffinity/BaseSessionAffinityProvider.cs @@ -0,0 +1,187 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using System.Collections.Generic; +using System.Text; +using Microsoft.AspNetCore.DataProtection; +using Microsoft.AspNetCore.Http; +using Microsoft.Extensions.Logging; +using Microsoft.ReverseProxy.RuntimeModel; + +namespace Microsoft.ReverseProxy.Service.SessionAffinity +{ + internal abstract class BaseSessionAffinityProvider : ISessionAffinityProvider + { + private readonly IDataProtector _dataProtector; + protected static readonly object AffinityKeyId = new object(); + protected readonly ILogger Logger; + + protected BaseSessionAffinityProvider(IDataProtectionProvider dataProtectionProvider, ILogger logger) + { + _dataProtector = dataProtectionProvider?.CreateProtector(GetType().FullName) ?? throw new ArgumentNullException(nameof(dataProtectionProvider)); + Logger = logger ?? throw new ArgumentNullException(nameof(logger)); + } + + public abstract string Mode { get; } + + public virtual void AffinitizeRequest(HttpContext context, in BackendConfig.BackendSessionAffinityOptions options, DestinationInfo destination) + { + if (!options.Enabled) + { + throw new InvalidOperationException($"Session affinity is disabled for backend."); + } + + // Affinity key is set on the response only if it's a new affinity. + if (!context.Items.ContainsKey(AffinityKeyId)) + { + var affinityKey = GetDestinationAffinityKey(destination); + SetAffinityKey(context, options, affinityKey); + } + } + + public virtual AffinityResult FindAffinitizedDestinations(HttpContext context, IReadOnlyList destinations, string backendId, in BackendConfig.BackendSessionAffinityOptions options) + { + if (!options.Enabled) + { + throw new InvalidOperationException($"Session affinity is disabled for backend {backendId}."); + } + + var requestAffinityKey = GetRequestAffinityKey(context, options); + + if (requestAffinityKey.Key == null) + { + return new AffinityResult(null, requestAffinityKey.ExtractedSuccessfully ? AffinityStatus.AffinityKeyNotSet : AffinityStatus.AffinityKeyExtractionFailed); + } + + IReadOnlyList matchingDestinations = null; + if (destinations.Count > 0) + { + for (var i = 0; i < destinations.Count; i++) + { + // TODO: Add fast destination lookup by ID + if (requestAffinityKey.Key.Equals(GetDestinationAffinityKey(destinations[i]))) + { + // It's allowed to affinitize a request to a pool of destinations so as to enable load-balancing among them. + // However, we currently stop after the first match found to avoid performance degradation. + matchingDestinations = destinations[i]; + break; + } + } + Log.DestinationMatchingToAffinityKeyNotFound(Logger, backendId); + } + else + { + Log.AffinityCannotBeEstablishedBecauseNoDestinationsFound(Logger, backendId); + } + + // Empty destination list passed to this method is handled the same way as if no matching destinations are found. + if (matchingDestinations == null) + { + return new AffinityResult(null, AffinityStatus.DestinationNotFound); + } + + context.Items[AffinityKeyId] = requestAffinityKey; + return new AffinityResult(matchingDestinations, AffinityStatus.OK); + } + + protected virtual string GetSettingValue(string key, BackendConfig.BackendSessionAffinityOptions options) + { + if (options.Settings == null || !options.Settings.TryGetValue(key, out var value)) + { + throw new ArgumentException(nameof(options), $"{nameof(CookieSessionAffinityProvider)} couldn't find the required parameter {key} in session affinity settings."); + } + + return value; + } + + protected abstract T GetDestinationAffinityKey(DestinationInfo destination); + + protected abstract (T Key, bool ExtractedSuccessfully) GetRequestAffinityKey(HttpContext context, in BackendConfig.BackendSessionAffinityOptions options); + + protected abstract void SetAffinityKey(HttpContext context, in BackendConfig.BackendSessionAffinityOptions options, T unencryptedKey); + + protected string Protect(string unencryptedKey) + { + if (string.IsNullOrEmpty(unencryptedKey)) + { + return unencryptedKey; + } + + var userData = Encoding.UTF8.GetBytes(unencryptedKey); + + var protectedData = _dataProtector.Protect(userData); + return Convert.ToBase64String(protectedData).TrimEnd('='); + } + + protected (string Key, bool ExtractedSuccessfully) Unprotect(string encryptedRequestKey) + { + if (string.IsNullOrEmpty(encryptedRequestKey)) + { + return (Key: null, ExtractedSuccessfully: true); + } + + try + { + var keyBytes = Convert.FromBase64String(Pad(encryptedRequestKey)); + + var decryptedKeyBytes = _dataProtector.Unprotect(keyBytes); + if (decryptedKeyBytes == null) + { + Log.RequestAffinityKeyDecryptionFailed(Logger, null); + return (Key: null, ExtractedSuccessfully: false); + } + + return (Key: Encoding.UTF8.GetString(decryptedKeyBytes), ExtractedSuccessfully: true); + } + catch (Exception ex) + { + Log.RequestAffinityKeyDecryptionFailed(Logger, ex); + return (Key: null, ExtractedSuccessfully: false); + } + } + + private static string Pad(string text) + { + var padding = 3 - ((text.Length + 3) % 4); + if (padding == 0) + { + return text; + } + return text + new string('=', padding); + } + + private static class Log + { + private static readonly Action _affinityCannotBeEstablishedBecauseNoDestinationsFound = LoggerMessage.Define( + LogLevel.Warning, + EventIds.AffinityCannotBeEstablishedBecauseNoDestinationsFoundOnBackend, + "The request affinity cannot be established because no destinations are found on backend `{backendId}`."); + + private static readonly Action _requestAffinityKeyDecryptionFailed = LoggerMessage.Define( + LogLevel.Error, + EventIds.RequestAffinityKeyDecryptionFailed, + "The request affinity key decryption failed."); + + private static readonly Action _destinationMatchingToAffinityKeyNotFound = LoggerMessage.Define( + LogLevel.Warning, + EventIds.DestinationMatchingToAffinityKeyNotFound, + "Destination matching to the request affinity key is not found on backend `{backnedId}`. Configured failure policy will be applied."); + + public static void AffinityCannotBeEstablishedBecauseNoDestinationsFound(ILogger logger, string backendId) + { + _affinityCannotBeEstablishedBecauseNoDestinationsFound(logger, backendId, null); + } + + public static void RequestAffinityKeyDecryptionFailed(ILogger logger, Exception ex) + { + _requestAffinityKeyDecryptionFailed(logger, ex); + } + + public static void DestinationMatchingToAffinityKeyNotFound(ILogger logger, string backendId) + { + _destinationMatchingToAffinityKeyNotFound(logger, backendId, null); + } + } + } +} diff --git a/src/ReverseProxy/Service/SessionAffinity/CookieSessionAffinityProvider.cs b/src/ReverseProxy/Service/SessionAffinity/CookieSessionAffinityProvider.cs new file mode 100644 index 000000000..a28d225ed --- /dev/null +++ b/src/ReverseProxy/Service/SessionAffinity/CookieSessionAffinityProvider.cs @@ -0,0 +1,46 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using Microsoft.AspNetCore.Http; +using Microsoft.AspNetCore.DataProtection; +using Microsoft.Extensions.Options; +using Microsoft.ReverseProxy.Abstractions.BackendDiscovery.Contract; +using Microsoft.ReverseProxy.RuntimeModel; +using Microsoft.Extensions.Logging; + +namespace Microsoft.ReverseProxy.Service.SessionAffinity +{ + internal class CookieSessionAffinityProvider : BaseSessionAffinityProvider + { + private readonly CookieSessionAffinityProviderOptions _providerOptions; + + public CookieSessionAffinityProvider( + IOptions providerOptions, + IDataProtectionProvider dataProtectionProvider, + ILogger logger) + : base(dataProtectionProvider, logger) + { + _providerOptions = providerOptions?.Value ?? throw new ArgumentNullException(nameof(providerOptions)); + } + + public override string Mode => SessionAffinityConstants.Modes.Cookie; + + protected override string GetDestinationAffinityKey(DestinationInfo destination) + { + return destination.DestinationId; + } + + protected override (string Key, bool ExtractedSuccessfully) GetRequestAffinityKey(HttpContext context, in BackendConfig.BackendSessionAffinityOptions options) + { + var encryptedRequestKey = context.Request.Cookies.TryGetValue(_providerOptions.Cookie.Name, out var keyInCookie) ? keyInCookie : null; + return Unprotect(encryptedRequestKey); + } + + protected override void SetAffinityKey(HttpContext context, in BackendConfig.BackendSessionAffinityOptions options, string unencryptedKey) + { + var affinityCookieOptions = _providerOptions.Cookie.Build(context); + context.Response.Cookies.Append(_providerOptions.Cookie.Name, Protect(unencryptedKey), affinityCookieOptions); + } + } +} diff --git a/src/ReverseProxy/Service/SessionAffinity/CustomHeaderSessionAffinityProvider.cs b/src/ReverseProxy/Service/SessionAffinity/CustomHeaderSessionAffinityProvider.cs new file mode 100644 index 000000000..5e6faf644 --- /dev/null +++ b/src/ReverseProxy/Service/SessionAffinity/CustomHeaderSessionAffinityProvider.cs @@ -0,0 +1,72 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using Microsoft.AspNetCore.DataProtection; +using Microsoft.AspNetCore.Http; +using Microsoft.Extensions.Logging; +using Microsoft.Extensions.Primitives; +using Microsoft.ReverseProxy.Abstractions.BackendDiscovery.Contract; +using Microsoft.ReverseProxy.RuntimeModel; + +namespace Microsoft.ReverseProxy.Service.SessionAffinity +{ + internal class CustomHeaderSessionAffinityProvider : BaseSessionAffinityProvider + { + public static readonly string DefaultCustomHeaderName = "X-Microsoft-Proxy-Affinity"; + private const string CustomHeaderNameKey = "CustomHeaderName"; + + public CustomHeaderSessionAffinityProvider( + IDataProtectionProvider dataProtectionProvider, + ILogger logger) + : base(dataProtectionProvider, logger) + {} + + public override string Mode => SessionAffinityConstants.Modes.CustomHeader; + + protected override string GetDestinationAffinityKey(DestinationInfo destination) + { + return destination.DestinationId; + } + + protected override (string Key, bool ExtractedSuccessfully) GetRequestAffinityKey(HttpContext context, in BackendConfig.BackendSessionAffinityOptions options) + { + var customHeaderName = options.Settings != null && options.Settings.TryGetValue(CustomHeaderNameKey, out var nameInSettings) ? nameInSettings : DefaultCustomHeaderName; + var keyHeaderValues = context.Request.Headers[customHeaderName]; + + if (StringValues.IsNullOrEmpty(keyHeaderValues)) + { + // It means affinity key is not defined that is a successful case + return (Key: null, ExtractedSuccessfully: true); + } + + if (keyHeaderValues.Count > 1) + { + // Multiple values is an ambiguous case which is considered a key extraction failure + Log.RequestAffinityHeaderHasMultipleValues(Logger, customHeaderName, keyHeaderValues.Count); + return (Key: null, ExtractedSuccessfully: false); + } + + return Unprotect(keyHeaderValues[0]); + } + + protected override void SetAffinityKey(HttpContext context, in BackendConfig.BackendSessionAffinityOptions options, string unencryptedKey) + { + var customHeaderName = GetSettingValue(CustomHeaderNameKey, options); + context.Response.Headers.Append(customHeaderName, Protect(unencryptedKey)); + } + + private static class Log + { + private static readonly Action _requestAffinityHeaderHasMultipleValues = LoggerMessage.Define( + LogLevel.Error, + EventIds.RequestAffinityHeaderHasMultipleValues, + "The request affinity header `{headerName}` has `{valueCount}` values."); + + public static void RequestAffinityHeaderHasMultipleValues(ILogger logger, string headerName, int valueCount) + { + _requestAffinityHeaderHasMultipleValues(logger, headerName, valueCount, null); + } + } + } +} diff --git a/src/ReverseProxy/Service/SessionAffinity/IAffinityFailurePolicy.cs b/src/ReverseProxy/Service/SessionAffinity/IAffinityFailurePolicy.cs new file mode 100644 index 000000000..a4ac902e4 --- /dev/null +++ b/src/ReverseProxy/Service/SessionAffinity/IAffinityFailurePolicy.cs @@ -0,0 +1,34 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.Threading.Tasks; +using Microsoft.AspNetCore.Http; +using Microsoft.ReverseProxy.Middleware; +using Microsoft.ReverseProxy.RuntimeModel; + +namespace Microsoft.ReverseProxy.Service.SessionAffinity +{ + /// + /// Affinity failures handling policy. + /// + internal interface IAffinityFailurePolicy + { + /// + /// A unique identifier for this failure policy. This will be referenced from config. + /// + public string Name { get; } + + /// + /// Handles affinity failures. This method assumes the full control on + /// and can change it in any way. + /// + /// Current request's context. + /// Session affinity options set for the backend. + /// Affinity resolution status. + /// + /// if the failure is considered recoverable and the request processing can proceed. + /// Otherwise, indicating that an error response has been generated and the request's processing must be terminated. + /// + public Task Handle(HttpContext context, BackendConfig.BackendSessionAffinityOptions options, AffinityStatus affinityStatus); + } +} diff --git a/src/ReverseProxy/Service/SessionAffinity/ISessionAffinityProvider.cs b/src/ReverseProxy/Service/SessionAffinity/ISessionAffinityProvider.cs new file mode 100644 index 000000000..a5799a20e --- /dev/null +++ b/src/ReverseProxy/Service/SessionAffinity/ISessionAffinityProvider.cs @@ -0,0 +1,39 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.Collections.Generic; +using System.Net.Http; +using Microsoft.AspNetCore.Http; +using Microsoft.ReverseProxy.RuntimeModel; + +namespace Microsoft.ReverseProxy.Service.SessionAffinity +{ + /// + /// Provides session affinity for load-balanced backends. + /// + public interface ISessionAffinityProvider + { + /// + /// A unique identifier for this session affinity implementation. This will be referenced from config. + /// + public string Mode { get; } + + /// + /// Finds to which the current request is affinitized by the affinity key. + /// + /// Current request's context. + /// s available for the request. + /// Target backend ID. + /// Affinity options. + /// carrying the found affinitized destinations if any and the . + public AffinityResult FindAffinitizedDestinations(HttpContext context, IReadOnlyList destinations, string backendId, in BackendConfig.BackendSessionAffinityOptions options); + + /// + /// Affinitize the current request to the given by setting the affinity key extracted from . + /// + /// Current request's context. + /// Affinity options. + /// to which request is to be affinitized. + public void AffinitizeRequest(HttpContext context, in BackendConfig.BackendSessionAffinityOptions options, DestinationInfo destination); + } +} diff --git a/src/ReverseProxy/Service/SessionAffinity/RedistributeAffinityFailurePolicy.cs b/src/ReverseProxy/Service/SessionAffinity/RedistributeAffinityFailurePolicy.cs new file mode 100644 index 000000000..316dd4b9b --- /dev/null +++ b/src/ReverseProxy/Service/SessionAffinity/RedistributeAffinityFailurePolicy.cs @@ -0,0 +1,29 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using System.Threading.Tasks; +using Microsoft.AspNetCore.Http; +using Microsoft.ReverseProxy.Abstractions.BackendDiscovery.Contract; +using Microsoft.ReverseProxy.RuntimeModel; + +namespace Microsoft.ReverseProxy.Service.SessionAffinity +{ + internal class RedistributeAffinityFailurePolicy : IAffinityFailurePolicy + { + public string Name => SessionAffinityConstants.AffinityFailurePolicies.Redistribute; + + public Task Handle(HttpContext context, BackendConfig.BackendSessionAffinityOptions options, AffinityStatus affinityStatus) + { + if (affinityStatus == AffinityStatus.OK + || affinityStatus == AffinityStatus.AffinityKeyNotSet) + { + throw new InvalidOperationException($"{nameof(RedistributeAffinityFailurePolicy)} is called to handle a successful request's affinity status {affinityStatus}."); + } + + // Available destinations list have not been changed in the context, + // so simply allow processing to proceed to load balancing. + return Task.FromResult(true); + } + } +} diff --git a/src/ReverseProxy/Service/SessionAffinity/Return503ErrorAffinityFailurePolicy.cs b/src/ReverseProxy/Service/SessionAffinity/Return503ErrorAffinityFailurePolicy.cs new file mode 100644 index 000000000..40b2a1052 --- /dev/null +++ b/src/ReverseProxy/Service/SessionAffinity/Return503ErrorAffinityFailurePolicy.cs @@ -0,0 +1,28 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using System.Threading.Tasks; +using Microsoft.AspNetCore.Http; +using Microsoft.ReverseProxy.Abstractions.BackendDiscovery.Contract; +using Microsoft.ReverseProxy.RuntimeModel; + +namespace Microsoft.ReverseProxy.Service.SessionAffinity +{ + internal class Return503ErrorAffinityFailurePolicy : IAffinityFailurePolicy + { + public string Name => SessionAffinityConstants.AffinityFailurePolicies.Return503Error; + + public Task Handle(HttpContext context, BackendConfig.BackendSessionAffinityOptions options, AffinityStatus affinityStatus) + { + if (affinityStatus == AffinityStatus.OK + || affinityStatus == AffinityStatus.AffinityKeyNotSet) + { + throw new InvalidOperationException($"{nameof(Return503ErrorAffinityFailurePolicy)} is called to handle a successful request's affinity status {affinityStatus}."); + } + + context.Response.StatusCode = 503; + return Task.FromResult(false); + } + } +} diff --git a/src/ReverseProxy/Service/SessionAffinity/SessionAffinityMiddlewareHelper.cs b/src/ReverseProxy/Service/SessionAffinity/SessionAffinityMiddlewareHelper.cs new file mode 100644 index 000000000..738f88243 --- /dev/null +++ b/src/ReverseProxy/Service/SessionAffinity/SessionAffinityMiddlewareHelper.cs @@ -0,0 +1,50 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using System.Collections.Generic; + +namespace Microsoft.ReverseProxy.Service.SessionAffinity +{ + internal static class SessionAffinityMiddlewareHelper + { + public static IDictionary ToDictionaryByUniqueId(this IEnumerable services, Func idSelector) + { + if (services == null) + { + throw new ArgumentNullException(nameof(services)); + } + + var result = new Dictionary(StringComparer.Ordinal); + + foreach (var service in services) + { + if (!result.TryAdd(idSelector(service), service)) + { + throw new ArgumentException(nameof(services), $"More than one {typeof(T)} found with the same identifier."); + } + } + + return result; + } + + public static IDictionary ToProviderDictionary(this IEnumerable sessionAffinityProviders) + { + return ToDictionaryByUniqueId(sessionAffinityProviders, p => p.Mode); + } + + public static IDictionary ToPolicyDictionary(this IEnumerable affinityFailurePolicies) + { + return ToDictionaryByUniqueId(affinityFailurePolicies, p => p.Name); + } + + public static T GetRequiredServiceById(this IDictionary services, string id) + { + if (!services.TryGetValue(id, out var result)) + { + throw new ArgumentException(nameof(id), $"No {typeof(T)} was found for the id {id}."); + } + return result; + } + } +} diff --git a/test/ReverseProxy.Tests/Middleware/AffinitizeRequestMiddlewareTests.cs b/test/ReverseProxy.Tests/Middleware/AffinitizeRequestMiddlewareTests.cs new file mode 100644 index 000000000..3cec79135 --- /dev/null +++ b/test/ReverseProxy.Tests/Middleware/AffinitizeRequestMiddlewareTests.cs @@ -0,0 +1,138 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Threading.Tasks; +using Microsoft.AspNetCore.Http; +using Microsoft.Extensions.Logging; +using Microsoft.ReverseProxy.Abstractions.Telemetry; +using Microsoft.ReverseProxy.RuntimeModel; +using Microsoft.ReverseProxy.Service.SessionAffinity; +using Moq; +using Xunit; + +namespace Microsoft.ReverseProxy.Middleware +{ + public class AffinitizeRequestMiddlewareTests : AffinityMiddlewareTestBase + { + [Fact] + public async Task Invoke_SingleDestinationChosen_InvokeAffinitizeRequest() + { + var backend = GetBackend(); + var invokedMode = string.Empty; + const string expectedMode = "Mode-B"; + var providers = RegisterAffinityProviders( + false, + Destinations[1], + backend.BackendId, + ("Mode-A", (AffinityStatus?)null, (DestinationInfo[])null, (Action)(p => throw new InvalidOperationException($"Provider {p.Mode} call is not expected."))), + (expectedMode, (AffinityStatus?)null, (DestinationInfo[])null, (Action)(p => invokedMode = p.Mode))); + var nextInvoked = false; + var middleware = new AffinitizeRequestMiddleware(c => { + nextInvoked = true; + return Task.CompletedTask; + }, + providers.Select(p => p.Object), + GetOperationLogger(), + new Mock>().Object); + var context = new DefaultHttpContext(); + context.Features.Set(backend); + var destinationFeature = GetDestinationsFeature(Destinations[1]); + context.Features.Set(destinationFeature); + + await middleware.Invoke(context); + + Assert.Equal(expectedMode, invokedMode); + Assert.True(nextInvoked); + providers[0].VerifyGet(p => p.Mode, Times.Once); + providers[0].VerifyNoOtherCalls(); + providers[1].VerifyAll(); + Assert.Same(destinationFeature.Destinations, Destinations[1]); + } + + [Fact] + public async Task Invoke_MultipleCandidateDestinations_ChooseOneAndInvokeAffinitizeRequest() + { + var backend = GetBackend(); + var invokedMode = string.Empty; + const string expectedMode = "Mode-B"; + var providers = new[] { + GetProviderForRandomDestination("Mode-A", Destinations, p => throw new InvalidOperationException($"Provider {p.Mode} call is not expected.")), + GetProviderForRandomDestination(expectedMode, Destinations, p => invokedMode = p.Mode) + }; + var nextInvoked = false; + var logger = AffinityTestHelper.GetLogger(); + var middleware = new AffinitizeRequestMiddleware(c => { + nextInvoked = true; + return Task.CompletedTask; + }, + providers.Select(p => p.Object), + GetOperationLogger(), + logger.Object); + var context = new DefaultHttpContext(); + context.Features.Set(backend); + var destinationFeature = GetDestinationsFeature(Destinations); + context.Features.Set(destinationFeature); + + await middleware.Invoke(context); + + Assert.Equal(expectedMode, invokedMode); + Assert.True(nextInvoked); + providers[0].VerifyGet(p => p.Mode, Times.Once); + providers[0].VerifyNoOtherCalls(); + providers[1].VerifyAll(); + logger.Verify( + l => l.Log(LogLevel.Warning, EventIds.MultipleDestinationsOnBackendToEstablishRequestAffinity, It.IsAny(), null, (Func)It.IsAny()), + Times.Once); + Assert.Equal(1, destinationFeature.Destinations.Count); + var chosen = destinationFeature.Destinations[0]; + var sameDestinationCount = Destinations.Count(d => chosen == d); + Assert.Equal(1, sameDestinationCount); + } + + [Fact] + public async Task Invoke_NoDestinationChosen_LogWarningAndCallNext() + { + var backend = GetBackend(); + var nextInvoked = false; + var logger = AffinityTestHelper.GetLogger(); + var middleware = new AffinitizeRequestMiddleware(c => { + nextInvoked = true; + return Task.CompletedTask; + }, + new ISessionAffinityProvider[0], + GetOperationLogger(), + logger.Object); + var context = new DefaultHttpContext(); + context.Features.Set(backend); + var destinationFeature = GetDestinationsFeature(new DestinationInfo[0]); + context.Features.Set(destinationFeature); + + await middleware.Invoke(context); + + Assert.True(nextInvoked); + logger.Verify( + l => l.Log(LogLevel.Warning, EventIds.NoDestinationOnBackendToEstablishRequestAffinity, It.IsAny(), null, (Func)It.IsAny()), + Times.Once); + Assert.Equal(0, destinationFeature.Destinations.Count); + } + + private IOperationLogger GetOperationLogger() + { + var result = new Mock>(MockBehavior.Strict); + result.Setup(l => l.Execute(It.IsAny(), It.IsAny())).Callback((string name, Action callback) => callback()); + return result.Object; + } + + private Mock GetProviderForRandomDestination(string mode, IReadOnlyList destinations, Action callback) + { + var provider = new Mock(MockBehavior.Strict); + provider.SetupGet(p => p.Mode).Returns(mode); + provider.Setup(p => p.AffinitizeRequest(It.IsAny(), BackendConfig.SessionAffinityOptions, It.Is(d => destinations.Contains(d)))) + .Callback(() => callback(provider.Object)); + return provider; + } + } +} diff --git a/test/ReverseProxy.Tests/Middleware/AffinitizedDestinationLookupMiddlewareTests.cs b/test/ReverseProxy.Tests/Middleware/AffinitizedDestinationLookupMiddlewareTests.cs new file mode 100644 index 000000000..e425213ab --- /dev/null +++ b/test/ReverseProxy.Tests/Middleware/AffinitizedDestinationLookupMiddlewareTests.cs @@ -0,0 +1,121 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using System.Linq; +using System.Threading.Tasks; +using Microsoft.AspNetCore.Http; +using Microsoft.Extensions.Logging; +using Microsoft.ReverseProxy.Abstractions.Telemetry; +using Microsoft.ReverseProxy.Service.SessionAffinity; +using Microsoft.ReverseProxy.Signals; +using Moq; +using Xunit; + +namespace Microsoft.ReverseProxy.Middleware +{ + public class AffinitizedDestinationLookupMiddlewareTests : AffinityMiddlewareTestBase + { + [Theory] + [InlineData(AffinityStatus.AffinityKeyNotSet, null)] + [InlineData(AffinityStatus.OK, AffinitizedDestinationName)] + public async Task Invoke_SuccessfulFlow_CallNext(AffinityStatus status, string foundDestinationId) + { + var backend = GetBackend(); + var foundDestinations = foundDestinationId != null ? Destinations.Where(d => d.DestinationId == foundDestinationId).ToArray() : null; + var invokedMode = string.Empty; + const string expectedMode = "Mode-B"; + var providers = RegisterAffinityProviders( + true, + Destinations, + backend.BackendId, + ("Mode-A", AffinityStatus.DestinationNotFound, (RuntimeModel.DestinationInfo[])null, (Action)(p => throw new InvalidOperationException($"Provider {p.Mode} call is not expected."))), + (expectedMode, status, foundDestinations, p => invokedMode = p.Mode)); + var nextInvoked = false; + var middleware = new AffinitizedDestinationLookupMiddleware(c => { + nextInvoked = true; + return Task.CompletedTask; + }, + providers.Select(p => p.Object), new IAffinityFailurePolicy[0], + GetOperationLogger(false), + new Mock>().Object); + var context = new DefaultHttpContext(); + context.Features.Set(backend); + var destinationFeature = GetDestinationsFeature(Destinations); + context.Features.Set(destinationFeature); + + await middleware.Invoke(context); + + Assert.Equal(expectedMode, invokedMode); + Assert.True(nextInvoked); + providers[0].VerifyGet(p => p.Mode, Times.Once); + providers[0].VerifyNoOtherCalls(); + providers[1].VerifyAll(); + + if (foundDestinationId != null) + { + Assert.Equal(1, destinationFeature.Destinations.Count); + Assert.Equal(foundDestinationId, destinationFeature.Destinations[0].DestinationId); + } + else + { + Assert.Same(Destinations, destinationFeature.Destinations); + } + } + + [Theory] + [InlineData(AffinityStatus.DestinationNotFound, true)] + [InlineData(AffinityStatus.DestinationNotFound, false)] + [InlineData(AffinityStatus.AffinityKeyExtractionFailed, true)] + [InlineData(AffinityStatus.AffinityKeyExtractionFailed, false)] + public async Task Invoke_ErrorFlow_CallFailurePolicy(AffinityStatus affinityStatus, bool keepProcessing) + { + var backend = GetBackend(); + var providers = RegisterAffinityProviders(true, Destinations, backend.BackendId, ("Mode-B", affinityStatus, null, _ => { })); + var invokedPolicy = string.Empty; + const string expectedPolicy = "Policy-1"; + var failurePolicies = RegisterFailurePolicies( + affinityStatus, + ("Policy-0", false, p => throw new InvalidOperationException($"Policy {p.Name} call is not expected.")), + (expectedPolicy, keepProcessing, p => invokedPolicy = p.Name)); + var nextInvoked = false; + var logger = AffinityTestHelper.GetLogger(); + var middleware = new AffinitizedDestinationLookupMiddleware(c => { + nextInvoked = true; + return Task.CompletedTask; + }, + providers.Select(p => p.Object), failurePolicies.Select(p => p.Object), + GetOperationLogger(true), + logger.Object); + var context = new DefaultHttpContext(); + context.Features.Set(backend); + var destinationFeature = GetDestinationsFeature(Destinations); + context.Features.Set(destinationFeature); + + await middleware.Invoke(context); + + Assert.Equal(expectedPolicy, invokedPolicy); + Assert.Equal(keepProcessing, nextInvoked); + failurePolicies[0].VerifyGet(p => p.Name, Times.Once); + failurePolicies[0].VerifyNoOtherCalls(); + failurePolicies[1].VerifyAll(); + if (!keepProcessing) + { + logger.Verify( + l => l.Log(LogLevel.Warning, EventIds.AffinityResolutionFailedForBackend, It.IsAny(), null, (Func)It.IsAny()), + Times.Once); + } + } + + private IOperationLogger GetOperationLogger(bool callFailurePolicy) + { + var result = new Mock>(MockBehavior.Strict); + result.Setup(l => l.Execute(It.IsAny(), It.IsAny>())).Returns((string name, Func callback) => callback()); + if (callFailurePolicy) + { + result.Setup(l => l.ExecuteAsync(It.IsAny(), It.IsAny>>())).Returns(async (string name, Func> callback) => await callback()); + } + return result.Object; + } + } +} diff --git a/test/ReverseProxy.Tests/Middleware/AffinityMiddlewareTestBase.cs b/test/ReverseProxy.Tests/Middleware/AffinityMiddlewareTestBase.cs new file mode 100644 index 000000000..3ee3c9673 --- /dev/null +++ b/test/ReverseProxy.Tests/Middleware/AffinityMiddlewareTestBase.cs @@ -0,0 +1,87 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using System.Collections.Generic; +using Microsoft.AspNetCore.Http; +using Microsoft.ReverseProxy.RuntimeModel; +using Microsoft.ReverseProxy.Service.Management; +using Microsoft.ReverseProxy.Service.Proxy.Infrastructure; +using Microsoft.ReverseProxy.Service.SessionAffinity; +using Microsoft.ReverseProxy.Signals; +using Moq; + +namespace Microsoft.ReverseProxy.Middleware +{ + public abstract class AffinityMiddlewareTestBase + { + protected const string AffinitizedDestinationName = "dest-B"; + protected readonly IReadOnlyList Destinations = new[] { new DestinationInfo("dest-A"), new DestinationInfo(AffinitizedDestinationName), new DestinationInfo("dest-C") }; + protected readonly BackendConfig BackendConfig = new BackendConfig(default, default, new BackendConfig.BackendSessionAffinityOptions(true, "Mode-B", "Policy-1", null)); + + internal BackendInfo GetBackend() + { + var destinationManager = new Mock(); + destinationManager.SetupGet(m => m.Items).Returns(SignalFactory.Default.CreateSignal(Destinations)); + var backend = new BackendInfo("backend-1", destinationManager.Object, new Mock().Object); + backend.Config.Value = BackendConfig; + return backend; + } + + internal IReadOnlyList> RegisterAffinityProviders( + bool lookupMiddlewareTest, + IReadOnlyList expectedDestinations, + string expectedBackend, + params (string Mode, AffinityStatus? Status, DestinationInfo[] Destinations, Action Callback)[] prototypes) + { + var result = new List>(); + foreach (var (mode, status, destinations, callback) in prototypes) + { + var provider = new Mock(MockBehavior.Strict); + provider.SetupGet(p => p.Mode).Returns(mode); + if (lookupMiddlewareTest) + { + provider.Setup(p => p.FindAffinitizedDestinations( + It.IsAny(), + expectedDestinations, + expectedBackend, + BackendConfig.SessionAffinityOptions)) + .Returns(new AffinityResult(destinations, status.Value)) + .Callback(() => callback(provider.Object)); + } + else + { + provider.Setup(p => p.AffinitizeRequest( + It.IsAny(), + BackendConfig.SessionAffinityOptions, + expectedDestinations[0])) + .Callback(() => callback(provider.Object)); + } + result.Add(provider); + } + return result.AsReadOnly(); + } + + internal IReadOnlyList> RegisterFailurePolicies(AffinityStatus expectedStatus, params (string Name, bool Handled, Action Callback)[] prototypes) + { + var result = new List>(); + foreach (var (name, handled, callback) in prototypes) + { + var policy = new Mock(MockBehavior.Strict); + policy.SetupGet(p => p.Name).Returns(name); + policy.Setup(p => p.Handle(It.IsAny(), It.Is(o => o.AffinityFailurePolicy == name), expectedStatus)) + .ReturnsAsync(handled) + .Callback(() => callback(policy.Object)); + result.Add(policy); + } + return result.AsReadOnly(); + } + + internal IAvailableDestinationsFeature GetDestinationsFeature(IReadOnlyList destinations) + { + var result = new Mock(MockBehavior.Strict); + result.SetupProperty(p => p.Destinations, destinations); + return result.Object; + } + } +} diff --git a/test/ReverseProxy.Tests/Middleware/DestinationInitializerMiddlewareTests.cs b/test/ReverseProxy.Tests/Middleware/DestinationInitializerMiddlewareTests.cs index 40cef94a7..7c46385ad 100644 --- a/test/ReverseProxy.Tests/Middleware/DestinationInitializerMiddlewareTests.cs +++ b/test/ReverseProxy.Tests/Middleware/DestinationInitializerMiddlewareTests.cs @@ -7,8 +7,6 @@ using Microsoft.AspNetCore.Http; using Microsoft.AspNetCore.Routing; using Microsoft.AspNetCore.Routing.Patterns; -using Microsoft.ReverseProxy.Abstractions.Telemetry; -using Microsoft.ReverseProxy.Telemetry; using Microsoft.ReverseProxy.RuntimeModel; using Microsoft.ReverseProxy.Service.Management; using Microsoft.ReverseProxy.Service.Proxy.Infrastructure; @@ -85,7 +83,8 @@ public async Task Invoke_NoHealthyEndpoints_503() proxyHttpClientFactory: proxyHttpClientFactoryMock.Object); backend1.Config.Value = new BackendConfig( new BackendConfig.BackendHealthCheckOptions(enabled: true, Timeout.InfiniteTimeSpan, Timeout.InfiniteTimeSpan, 0, ""), - new BackendConfig.BackendLoadBalancingOptions()); + new BackendConfig.BackendLoadBalancingOptions(), + new BackendConfig.BackendSessionAffinityOptions()); var destination1 = backend1.DestinationManager.GetOrCreateItem( "destination1", destination => diff --git a/test/ReverseProxy.Tests/Middleware/LoadBalancerMiddlewareTests.cs b/test/ReverseProxy.Tests/Middleware/LoadBalancerMiddlewareTests.cs index 1ad19414f..cda5dda5c 100644 --- a/test/ReverseProxy.Tests/Middleware/LoadBalancerMiddlewareTests.cs +++ b/test/ReverseProxy.Tests/Middleware/LoadBalancerMiddlewareTests.cs @@ -41,7 +41,7 @@ public async Task Invoke_Works() backendId: "backend1", destinationManager: new DestinationManager(), proxyHttpClientFactory: proxyHttpClientFactoryMock.Object); - backend1.Config.Value = new BackendConfig(default, new BackendConfig.BackendLoadBalancingOptions(LoadBalancingMode.RoundRobin)); + backend1.Config.Value = new BackendConfig(default, new BackendConfig.BackendLoadBalancingOptions(LoadBalancingMode.RoundRobin), default); var destination1 = backend1.DestinationManager.GetOrCreateItem( "destination1", destination => diff --git a/test/ReverseProxy.Tests/Service/Config/DynamicConfigBuilderTests.cs b/test/ReverseProxy.Tests/Service/Config/DynamicConfigBuilderTests.cs index ccb17d4c8..67c9b9a84 100644 --- a/test/ReverseProxy.Tests/Service/Config/DynamicConfigBuilderTests.cs +++ b/test/ReverseProxy.Tests/Service/Config/DynamicConfigBuilderTests.cs @@ -27,6 +27,8 @@ private IDynamicConfigBuilder CreateConfigBuilder(IBackendsRepo backends, IRoute servicesBuilder.AddSingleton(backends); servicesBuilder.AddSingleton(routes); servicesBuilder.AddSingleton(); + servicesBuilder.AddDataProtection(); + servicesBuilder.AddLogging(); var services = servicesBuilder.BuildServiceProvider(); return services.GetRequiredService(); } diff --git a/test/ReverseProxy.Tests/Service/HealthProbe/BackendProberFactoryTests.cs b/test/ReverseProxy.Tests/Service/HealthProbe/BackendProberFactoryTests.cs index 591d534da..7556a637c 100644 --- a/test/ReverseProxy.Tests/Service/HealthProbe/BackendProberFactoryTests.cs +++ b/test/ReverseProxy.Tests/Service/HealthProbe/BackendProberFactoryTests.cs @@ -57,7 +57,8 @@ public void BackendProberFactory_CreateBackendProber() timeout: TimeSpan.FromSeconds(60), port: 8000, path: "/example"), - loadBalancingOptions: default); + loadBalancingOptions: default, + sessionAffinityOptions: default); var destinationManager = new DestinationManager(); var prober = factory.CreateBackendProber(backendId, backendConfig, destinationManager); diff --git a/test/ReverseProxy.Tests/Service/HealthProbe/BackendProberTests.cs b/test/ReverseProxy.Tests/Service/HealthProbe/BackendProberTests.cs index af4c33ab2..964cc3ed1 100644 --- a/test/ReverseProxy.Tests/Service/HealthProbe/BackendProberTests.cs +++ b/test/ReverseProxy.Tests/Service/HealthProbe/BackendProberTests.cs @@ -52,7 +52,8 @@ public BackendProberTests() timeout: TimeSpan.FromSeconds(60), port: 8000, path: "/example"), - loadBalancingOptions: default); + loadBalancingOptions: default, + sessionAffinityOptions: default); _timer = new VirtualMonotonicTimer(); _semaphore = new AsyncSemaphore(10); _fakeRandom = new Mock(); diff --git a/test/ReverseProxy.Tests/Service/HealthProbe/HealthProbeWorkerTests.cs b/test/ReverseProxy.Tests/Service/HealthProbe/HealthProbeWorkerTests.cs index 5c09c1610..4e9ee1a0b 100644 --- a/test/ReverseProxy.Tests/Service/HealthProbe/HealthProbeWorkerTests.cs +++ b/test/ReverseProxy.Tests/Service/HealthProbe/HealthProbeWorkerTests.cs @@ -41,7 +41,8 @@ public HealthProbeWorkerTests() timeout: TimeSpan.FromSeconds(1), port: 8000, path: "/example"), - loadBalancingOptions: default); + loadBalancingOptions: default, + sessionAffinityOptions: default); // Set up prober. We do not want to let prober really perform any actions. // The behavior of prober should be tested in its own unit test, see "BackendProberTests.cs". @@ -118,7 +119,8 @@ public async Task UpdateTrackedBackends_ProbeNotEnabled_ShouldNotStartProber() timeout: TimeSpan.FromSeconds(20), port: 1234, path: "/"), - loadBalancingOptions: default); + loadBalancingOptions: default, + sessionAffinityOptions: default); }); // Start probing. @@ -196,7 +198,8 @@ public async Task UpdateTrackedBackends_BackendConfigChange_RecreatesProber() timeout: TimeSpan.FromSeconds(1), port: 8000, path: "/newexample"), - loadBalancingOptions: default); + loadBalancingOptions: default, + sessionAffinityOptions: default); await health.UpdateTrackedBackends(); // After the config is updated, the program should discover this change, create a new prober, @@ -228,7 +231,8 @@ public async Task UpdateTrackedBackends_BackendConfigDisabledProbing_StopsProber timeout: TimeSpan.FromSeconds(1), port: 8000, path: "/newexample"), - loadBalancingOptions: default); + loadBalancingOptions: default, + sessionAffinityOptions: default); await health.UpdateTrackedBackends(); // After the config is updated, the program should discover this change, diff --git a/test/ReverseProxy.Tests/Service/RuntimeModel/BackendInfoTests.cs b/test/ReverseProxy.Tests/Service/RuntimeModel/BackendInfoTests.cs index 8c2e747dc..1e80cd9da 100644 --- a/test/ReverseProxy.Tests/Service/RuntimeModel/BackendInfoTests.cs +++ b/test/ReverseProxy.Tests/Service/RuntimeModel/BackendInfoTests.cs @@ -2,6 +2,7 @@ // Licensed under the MIT License. using System; +using Microsoft.AspNetCore.Mvc; using Microsoft.ReverseProxy.Service.Management; using Microsoft.ReverseProxy.Service.Proxy.Infrastructure; using Tests.Common; @@ -76,7 +77,7 @@ public void DynamicState_ReactsToBackendConfigChanges() Assert.NotNull(state1); Assert.Empty(state1.AllDestinations); - backend.Config.Value = new BackendConfig(healthCheckOptions: default, loadBalancingOptions: default); + backend.Config.Value = new BackendConfig(healthCheckOptions: default, loadBalancingOptions: default, sessionAffinityOptions: default); Assert.NotSame(state1, backend.DynamicState.Value); Assert.Empty(backend.DynamicState.Value.AllDestinations); } @@ -145,7 +146,8 @@ private static void EnableHealthChecks(BackendInfo backend) timeout: TimeSpan.FromSeconds(30), port: 30000, path: "/"), - loadBalancingOptions: default); + loadBalancingOptions: default, + sessionAffinityOptions: default); } } } diff --git a/test/ReverseProxy.Tests/Service/SessionAffinity/AffinityTestHelper.cs b/test/ReverseProxy.Tests/Service/SessionAffinity/AffinityTestHelper.cs new file mode 100644 index 000000000..14e4dfdd1 --- /dev/null +++ b/test/ReverseProxy.Tests/Service/SessionAffinity/AffinityTestHelper.cs @@ -0,0 +1,35 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using System.Text; +using Microsoft.AspNetCore.DataProtection; +using Microsoft.Extensions.Logging; +using Moq; + +namespace Microsoft.ReverseProxy.Service.SessionAffinity +{ + public static class AffinityTestHelper + { + public static Mock> GetLogger() + { + var result = new Mock>(); + result.Setup(l => l.IsEnabled(It.IsAny())).Returns(true); + return result; + } + + public static Mock GetDataProtector() + { + var protector = new Mock(); + protector.Setup(p => p.CreateProtector(It.IsAny())).Returns(protector.Object); + protector.Setup(p => p.Protect(It.IsAny())).Returns((byte[] b) => b); + protector.Setup(p => p.Unprotect(It.IsAny())).Returns((byte[] b) => b); + return protector; + } + + public static string ToUTF8BytesInBase64(this string text) + { + return Convert.ToBase64String(Encoding.UTF8.GetBytes(text)); + } + } +} diff --git a/test/ReverseProxy.Tests/Service/SessionAffinity/BaseSesstionAffinityProviderTests.cs b/test/ReverseProxy.Tests/Service/SessionAffinity/BaseSesstionAffinityProviderTests.cs new file mode 100644 index 000000000..9ee81403c --- /dev/null +++ b/test/ReverseProxy.Tests/Service/SessionAffinity/BaseSesstionAffinityProviderTests.cs @@ -0,0 +1,211 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Text; +using Microsoft.AspNetCore.DataProtection; +using Microsoft.AspNetCore.Http; +using Microsoft.Extensions.Logging; +using Microsoft.ReverseProxy.RuntimeModel; +using Moq; +using Xunit; + +namespace Microsoft.ReverseProxy.Service.SessionAffinity +{ + public class BaseSesstionAffinityProviderTests + { + private const string InvalidKeyNull = "!invalid key - null!"; + private const string InvalidKeyThrow = "!invalid key - throw!"; + private const string KeyName = "StubAffinityKey"; + private readonly BackendConfig.BackendSessionAffinityOptions _defaultOptions = new BackendConfig.BackendSessionAffinityOptions(true, "Stub", "Return503", new Dictionary { { "AffinityKeyName", KeyName } }); + + [Theory] + [MemberData(nameof(FindAffinitizedDestinationsCases))] + public void Request_FindAffinitizedDestinations( + HttpContext context, + DestinationInfo[] allDestinations, + AffinityStatus expectedStatus, + DestinationInfo expectedDestination, + byte[] expectedEncryptedKey, + bool unprotectCalled, + LogLevel? expectedLogLevel, + EventId expectedEventId) + { + var dataProtector = GetDataProtector(); + var logger = AffinityTestHelper.GetLogger>(); + var provider = new ProviderStub(dataProtector.Object, logger.Object); + var affinityResult = provider.FindAffinitizedDestinations(context, allDestinations, "backend-1", _defaultOptions); + + if(unprotectCalled) + { + dataProtector.Verify(p => p.Unprotect(It.Is(b => b.SequenceEqual(expectedEncryptedKey))), Times.Once); + } + + Assert.Equal(expectedStatus, affinityResult.Status); + Assert.Same(expectedDestination, affinityResult.Destinations?.FirstOrDefault()); + + if (expectedLogLevel != null) + { + logger.Verify( + l => l.Log(expectedLogLevel.Value, expectedEventId, It.IsAny(), It.IsAny(), (Func)It.IsAny()), + Times.Once); + } + + if (expectedDestination != null) + { + Assert.Equal(1, affinityResult.Destinations.Count); + } + else + { + Assert.Null(affinityResult.Destinations); + } + } + + [Fact] + public void FindAffinitizedDestination_AffinityDisabledOnBackend_ReturnsAffinityDisabled() + { + var provider = new ProviderStub(GetDataProtector().Object, AffinityTestHelper.GetLogger>().Object); + var options = new BackendConfig.BackendSessionAffinityOptions(false, _defaultOptions.Mode, _defaultOptions.AffinityFailurePolicy, _defaultOptions.Settings); + Assert.Throws(() => provider.FindAffinitizedDestinations(new DefaultHttpContext(), new[] { new DestinationInfo("1") }, "backend-1", options)); + } + + [Fact] + public void AffinitizeRequest_AffinitiDisabled_DoNothing() + { + var dataProtector = GetDataProtector(); + var provider = new ProviderStub(dataProtector.Object, AffinityTestHelper.GetLogger>().Object); + Assert.Throws(() => provider.AffinitizeRequest(new DefaultHttpContext(), default, new DestinationInfo("id"))); + } + + [Fact] + public void AffinitizeRequest_RequestIsAffinitized_DoNothing() + { + var dataProtector = GetDataProtector(); + var provider = new ProviderStub(dataProtector.Object, AffinityTestHelper.GetLogger>().Object); + var context = new DefaultHttpContext(); + provider.DirectlySetExtractedKeyOnContext(context, "ExtractedKey"); + provider.AffinitizeRequest(context, _defaultOptions, new DestinationInfo("id")); + Assert.Null(provider.LastSetEncryptedKey); + dataProtector.Verify(p => p.Protect(It.IsAny()), Times.Never); + } + + [Fact] + public void AffinitizeRequest_RequestIsNotAffinitized_SetAffinityKey() + { + var dataProtector = GetDataProtector(); + var provider = new ProviderStub(dataProtector.Object, AffinityTestHelper.GetLogger>().Object); + var destination = new DestinationInfo("dest-A"); + provider.AffinitizeRequest(new DefaultHttpContext(), _defaultOptions, destination); + Assert.Equal("ZGVzdC1B", provider.LastSetEncryptedKey); + var keyBytes = Encoding.UTF8.GetBytes(destination.DestinationId); + dataProtector.Verify(p => p.Protect(It.Is(b => b.SequenceEqual(keyBytes))), Times.Once); + } + + [Fact] + public void FindAffinitizedDestinations_AffinityOptionSettingNotFound_Throw() + { + var provider = new ProviderStub(GetDataProtector().Object, AffinityTestHelper.GetLogger>().Object); + var options = GetOptionsWithUnknownSetting(); + Assert.Throws(() => provider.FindAffinitizedDestinations(new DefaultHttpContext(), new[] { new DestinationInfo("dest-A") }, "backend-1", options)); + } + + [Fact] + public void AffinitizeRequest_AffinityOptionSettingNotFound_Throw() + { + var provider = new ProviderStub(GetDataProtector().Object, AffinityTestHelper.GetLogger>().Object); + var options = GetOptionsWithUnknownSetting(); + Assert.Throws(() => provider.AffinitizeRequest(new DefaultHttpContext(), options, new DestinationInfo("dest-A"))); + } + + [Fact] + public void Ctor_MandatoryArgumentIsNull_Throw() + { + Assert.Throws(() => new ProviderStub(null, new Mock().Object)); + // CreateDataProtector will return null + Assert.Throws(() => new ProviderStub(new Mock().Object, new Mock().Object)); + Assert.Throws(() => new ProviderStub(GetDataProtector().Object, null)); + } + + public static IEnumerable FindAffinitizedDestinationsCases() + { + var destinations = new[] { new DestinationInfo("dest-A"), new DestinationInfo("dest-B"), new DestinationInfo("dest-C") }; + yield return new object[] { GetHttpContext(new[] { ("SomeKey", "SomeValue") }), destinations, AffinityStatus.AffinityKeyNotSet, null, null, false, null, null }; + yield return new object[] { GetHttpContext(new[] { (KeyName, "dest-B") }), destinations, AffinityStatus.OK, destinations[1], Encoding.UTF8.GetBytes("dest-B"), true, null, null }; + yield return new object[] { GetHttpContext(new[] { (KeyName, "dest-Z") }), destinations, AffinityStatus.DestinationNotFound, null, Encoding.UTF8.GetBytes("dest-Z"), true, LogLevel.Warning, EventIds.DestinationMatchingToAffinityKeyNotFound }; + yield return new object[] { GetHttpContext(new[] { (KeyName, "dest-B") }), new DestinationInfo[0], AffinityStatus.DestinationNotFound, null, Encoding.UTF8.GetBytes("dest-B"), true, LogLevel.Warning, EventIds.AffinityCannotBeEstablishedBecauseNoDestinationsFoundOnBackend }; + yield return new object[] { GetHttpContext(new[] { (KeyName, "/////") }, false), destinations, AffinityStatus.AffinityKeyExtractionFailed, null, Encoding.UTF8.GetBytes(InvalidKeyNull), false, LogLevel.Error, EventIds.RequestAffinityKeyDecryptionFailed }; + yield return new object[] { GetHttpContext(new[] { (KeyName, InvalidKeyNull) }), destinations, AffinityStatus.AffinityKeyExtractionFailed, null, Encoding.UTF8.GetBytes(InvalidKeyNull), true, LogLevel.Error, EventIds.RequestAffinityKeyDecryptionFailed }; + yield return new object[] { GetHttpContext(new[] { (KeyName, InvalidKeyThrow) }), destinations, AffinityStatus.AffinityKeyExtractionFailed, null, Encoding.UTF8.GetBytes(InvalidKeyThrow), true, LogLevel.Error, EventIds.RequestAffinityKeyDecryptionFailed }; + } + + private static BackendConfig.BackendSessionAffinityOptions GetOptionsWithUnknownSetting() + { + return new BackendConfig.BackendSessionAffinityOptions(true, "Stub", "Return503", new Dictionary { { "Unknown", "ZZZ" } }); + } + + private static HttpContext GetHttpContext((string Key, string Value)[] items, bool encodeToBase64 = true) + { + var context = new DefaultHttpContext + { + Items = items.ToDictionary(i => (object)i.Key, i => encodeToBase64 ? (object)Convert.ToBase64String(Encoding.UTF8.GetBytes(i.Value)) : i.Value) + }; + return context; + } + + private Mock GetDataProtector() + { + var result = new Mock(); + var nullBytes = Encoding.UTF8.GetBytes(InvalidKeyNull); + var throwBytes = Encoding.UTF8.GetBytes(InvalidKeyThrow); + result.Setup(p => p.Protect(It.IsAny())).Returns((byte[] k) => k); + result.Setup(p => p.Unprotect(It.IsAny())).Returns((byte[] k) => k); + result.Setup(p => p.Unprotect(It.Is(b => b.SequenceEqual(nullBytes)))).Returns((byte[])null); + result.Setup(p => p.Unprotect(It.Is(b => b.SequenceEqual(throwBytes)))).Throws(); + result.Setup(p => p.CreateProtector(It.IsAny())).Returns(result.Object); + return result; + } + + private class ProviderStub : BaseSessionAffinityProvider + { + public static readonly string KeyNameSetting = "AffinityKeyName"; + + public ProviderStub(IDataProtectionProvider dataProtectionProvider, ILogger logger) + : base(dataProtectionProvider, logger) + {} + + public override string Mode => "Stub"; + + public string LastSetEncryptedKey { get; private set; } + + public void DirectlySetExtractedKeyOnContext(HttpContext context, string key) + { + context.Items[AffinityKeyId] = key; + } + + protected override string GetDestinationAffinityKey(DestinationInfo destination) + { + return destination.DestinationId; + } + + protected override (string Key, bool ExtractedSuccessfully) GetRequestAffinityKey(HttpContext context, in BackendConfig.BackendSessionAffinityOptions options) + { + Assert.Equal(Mode, options.Mode); + var keyName = GetSettingValue(KeyNameSetting, options); + // HttpContext.Items is used here to store the request affinity key for simplicity. + // In real world scenario, a provider will extract it from request (e.g. header, cookie, etc.) + var encryptedKey = context.Items.TryGetValue(keyName, out var requestKey) ? requestKey : null; + return Unprotect((string)encryptedKey); + } + + protected override void SetAffinityKey(HttpContext context, in BackendConfig.BackendSessionAffinityOptions options, string unencryptedKey) + { + var keyName = GetSettingValue(KeyNameSetting, options); + var encryptedKey = Protect(unencryptedKey); + context.Items[keyName] = encryptedKey; + LastSetEncryptedKey = encryptedKey; + } + } + } +} diff --git a/test/ReverseProxy.Tests/Service/SessionAffinity/CookieSessionAffinityProviderTests.cs b/test/ReverseProxy.Tests/Service/SessionAffinity/CookieSessionAffinityProviderTests.cs new file mode 100644 index 000000000..d4e24f2e8 --- /dev/null +++ b/test/ReverseProxy.Tests/Service/SessionAffinity/CookieSessionAffinityProviderTests.cs @@ -0,0 +1,121 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using System.Collections.Generic; +using Microsoft.AspNetCore.Http; +using Microsoft.Extensions.Options; +using Microsoft.ReverseProxy.Abstractions.BackendDiscovery.Contract; +using Microsoft.ReverseProxy.RuntimeModel; +using Xunit; + +namespace Microsoft.ReverseProxy.Service.SessionAffinity +{ + public class CookieSessionAffinityProviderTests + { + private readonly CookieSessionAffinityProviderOptions _defaultProviderOptions = new CookieSessionAffinityProviderOptions(); + private readonly BackendConfig.BackendSessionAffinityOptions _defaultOptions = new BackendConfig.BackendSessionAffinityOptions(true, "Cookie", "Return503", null); + private readonly IReadOnlyList _destinations = new[] { new DestinationInfo("dest-A"), new DestinationInfo("dest-B"), new DestinationInfo("dest-C") }; + + [Fact] + public void FindAffinitizedDestination_AffinityKeyIsNotSetOnRequest_ReturnKeyNotSet() + { + var provider = new CookieSessionAffinityProvider( + Options.Create(_defaultProviderOptions), + AffinityTestHelper.GetDataProtector().Object, + AffinityTestHelper.GetLogger().Object); + + Assert.Equal(SessionAffinityConstants.Modes.Cookie, provider.Mode); + + var context = new DefaultHttpContext(); + context.Request.Headers["Cookie"] = new[] { $"Some-Cookie=ZZZ" }; + + var affinityResult = provider.FindAffinitizedDestinations(context, _destinations, "backend-1", _defaultOptions); + + Assert.Equal(AffinityStatus.AffinityKeyNotSet, affinityResult.Status); + Assert.Null(affinityResult.Destinations); + } + + [Fact] + public void FindAffinitizedDestination_AffinityKeyIsSetOnRequest_Success() + { + var provider = new CookieSessionAffinityProvider( + Options.Create(_defaultProviderOptions), + AffinityTestHelper.GetDataProtector().Object, + AffinityTestHelper.GetLogger().Object); + var context = new DefaultHttpContext(); + var affinitizedDestination = _destinations[1]; + context.Request.Headers["Cookie"] = GetCookieWithAffinity(affinitizedDestination); + + var affinityResult = provider.FindAffinitizedDestinations(context, _destinations, "backend-1", _defaultOptions); + + Assert.Equal(AffinityStatus.OK, affinityResult.Status); + Assert.Equal(1, affinityResult.Destinations.Count); + Assert.Same(affinitizedDestination, affinityResult.Destinations[0]); + } + + [Fact] + public void AffinitizedRequest_AffinityKeyIsNotExtracted_SetKeyOnResponse() + { + var provider = new CookieSessionAffinityProvider( + Options.Create(_defaultProviderOptions), + AffinityTestHelper.GetDataProtector().Object, + AffinityTestHelper.GetLogger().Object); + var context = new DefaultHttpContext(); + + provider.AffinitizeRequest(context, _defaultOptions, _destinations[1]); + + var affinityCookieHeader = context.Response.Headers["Set-Cookie"]; + Assert.Equal(".Microsoft.ReverseProxy.Affinity=ZGVzdC1C; path=/; httponly", affinityCookieHeader); + } + + [Fact] + public void AffinitizeRequest_CookieBuilderSettingsChanged_UseNewSettings() + { + var providerOptions = new CookieSessionAffinityProviderOptions(); + providerOptions.Cookie.Domain = "mydomain.my"; + providerOptions.Cookie.HttpOnly = false; + providerOptions.Cookie.IsEssential = true; + providerOptions.Cookie.MaxAge = TimeSpan.FromHours(1); + providerOptions.Cookie.Name = "My.Affinity"; + providerOptions.Cookie.Path = "/some"; + providerOptions.Cookie.SameSite = SameSiteMode.Lax; + providerOptions.Cookie.SecurePolicy = CookieSecurePolicy.Always; + var provider = new CookieSessionAffinityProvider( + Options.Create(providerOptions), + AffinityTestHelper.GetDataProtector().Object, + AffinityTestHelper.GetLogger().Object); + var context = new DefaultHttpContext(); + + provider.AffinitizeRequest(context, _defaultOptions, _destinations[1]); + + var affinityCookieHeader = context.Response.Headers["Set-Cookie"]; + Assert.Equal("My.Affinity=ZGVzdC1C; max-age=3600; domain=mydomain.my; path=/some; secure; samesite=lax", affinityCookieHeader); + } + + [Fact] + public void AffinitizedRequest_AffinityKeyIsExtracted_DoNothing() + { + var provider = new CookieSessionAffinityProvider( + Options.Create(_defaultProviderOptions), + AffinityTestHelper.GetDataProtector().Object, + AffinityTestHelper.GetLogger().Object); + var context = new DefaultHttpContext(); + var affinitizedDestination = _destinations[0]; + context.Request.Headers["Cookie"] = GetCookieWithAffinity(affinitizedDestination); + + var affinityResult = provider.FindAffinitizedDestinations(context, _destinations, "backend-1", _defaultOptions); + + Assert.Equal(AffinityStatus.OK, affinityResult.Status); + + provider.AffinitizeRequest(context, _defaultOptions, affinitizedDestination); + + Assert.False(context.Response.Headers.ContainsKey("Cookie")); + } + + private string[] GetCookieWithAffinity(DestinationInfo affinitizedDestination) + { + return new[] { $"Some-Cookie=ZZZ", $"{_defaultProviderOptions.Cookie.Name}={affinitizedDestination.DestinationId.ToUTF8BytesInBase64()}" }; + } + } +} diff --git a/test/ReverseProxy.Tests/Service/SessionAffinity/CustomHeaderSessionAffinityProviderTests.cs b/test/ReverseProxy.Tests/Service/SessionAffinity/CustomHeaderSessionAffinityProviderTests.cs new file mode 100644 index 000000000..53f2a21af --- /dev/null +++ b/test/ReverseProxy.Tests/Service/SessionAffinity/CustomHeaderSessionAffinityProviderTests.cs @@ -0,0 +1,107 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using System.Collections.Generic; +using Microsoft.AspNetCore.Http; +using Microsoft.ReverseProxy.Abstractions.BackendDiscovery.Contract; +using Microsoft.ReverseProxy.RuntimeModel; +using Xunit; + +namespace Microsoft.ReverseProxy.Service.SessionAffinity +{ + public class CustomHeaderSessionAffinityProviderTests + { + private const string AffinityHeaderName = "X-MyAffinity"; + private readonly BackendConfig.BackendSessionAffinityOptions _defaultOptions = + new BackendConfig.BackendSessionAffinityOptions(true, "Cookie", "Return503", new Dictionary { { "CustomHeaderName", AffinityHeaderName } }); + private readonly IReadOnlyList _destinations = new[] { new DestinationInfo("dest-A"), new DestinationInfo("dest-B"), new DestinationInfo("dest-C") }; + + [Fact] + public void FindAffinitizedDestination_AffinityKeyIsNotSetOnRequest_ReturnKeyNotSet() + { + var provider = new CustomHeaderSessionAffinityProvider(AffinityTestHelper.GetDataProtector().Object, AffinityTestHelper.GetLogger().Object); + + Assert.Equal(SessionAffinityConstants.Modes.CustomHeader, provider.Mode); + + var context = new DefaultHttpContext(); + context.Request.Headers["SomeHeader"] = new[] { "SomeValue" }; + + var affinityResult = provider.FindAffinitizedDestinations(context, _destinations, "backend-1", _defaultOptions); + + Assert.Equal(AffinityStatus.AffinityKeyNotSet, affinityResult.Status); + Assert.Null(affinityResult.Destinations); + } + + [Fact] + public void FindAffinitizedDestination_AffinityKeyIsSetOnRequest_Success() + { + var provider = new CustomHeaderSessionAffinityProvider(AffinityTestHelper.GetDataProtector().Object, AffinityTestHelper.GetLogger().Object); + var context = new DefaultHttpContext(); + context.Request.Headers["SomeHeader"] = new[] { "SomeValue" }; + var affinitizedDestination = _destinations[1]; + context.Request.Headers[AffinityHeaderName] = new[] { affinitizedDestination.DestinationId.ToUTF8BytesInBase64() }; + + var affinityResult = provider.FindAffinitizedDestinations(context, _destinations, "backend-1", _defaultOptions); + + Assert.Equal(AffinityStatus.OK, affinityResult.Status); + Assert.Equal(1, affinityResult.Destinations.Count); + Assert.Same(affinitizedDestination, affinityResult.Destinations[0]); + } + + [Theory] + [MemberData(nameof(FindAffinitizedDestination_CustomHeaderNameIsNotSpecified_Cases))] + public void FindAffinitizedDestination_CustomHeaderNameIsNotSpecified_UseDefaultName(Dictionary settings) + { + var options = new BackendConfig.BackendSessionAffinityOptions(true, "CustomHeader", "Return503", settings); + var provider = new CustomHeaderSessionAffinityProvider(AffinityTestHelper.GetDataProtector().Object, AffinityTestHelper.GetLogger().Object); + var context = new DefaultHttpContext(); + var affinitizedDestination = _destinations[1]; + context.Request.Headers[CustomHeaderSessionAffinityProvider.DefaultCustomHeaderName] = new[] { affinitizedDestination.DestinationId.ToUTF8BytesInBase64() }; + + var affinityResult = provider.FindAffinitizedDestinations(context, _destinations, "backend-1", options); + + Assert.Equal(AffinityStatus.OK, affinityResult.Status); + Assert.Equal(1, affinityResult.Destinations.Count); + Assert.Same(affinitizedDestination, affinityResult.Destinations[0]); + } + + [Fact] + public void AffinitizedRequest_AffinityKeyIsNotExtracted_SetKeyOnResponse() + { + var provider = new CustomHeaderSessionAffinityProvider(AffinityTestHelper.GetDataProtector().Object, AffinityTestHelper.GetLogger().Object); + var context = new DefaultHttpContext(); + var chosenDestination = _destinations[1]; + var expectedAffinityHeaderValue = chosenDestination.DestinationId.ToUTF8BytesInBase64(); + + provider.AffinitizeRequest(context, _defaultOptions, chosenDestination); + + Assert.True(context.Response.Headers.ContainsKey(AffinityHeaderName)); + Assert.Equal(expectedAffinityHeaderValue, context.Response.Headers[AffinityHeaderName]); + } + + [Fact] + public void AffinitizedRequest_AffinityKeyIsExtracted_DoNothing() + { + var provider = new CustomHeaderSessionAffinityProvider(AffinityTestHelper.GetDataProtector().Object, AffinityTestHelper.GetLogger().Object); + var context = new DefaultHttpContext(); + context.Request.Headers["SomeHeader"] = new[] { "SomeValue" }; + var affinitizedDestination = _destinations[1]; + context.Request.Headers[AffinityHeaderName] = new[] { affinitizedDestination.DestinationId.ToUTF8BytesInBase64() }; + + var affinityResult = provider.FindAffinitizedDestinations(context, _destinations, "backend-1", _defaultOptions); + + Assert.Equal(AffinityStatus.OK, affinityResult.Status); + + provider.AffinitizeRequest(context, _defaultOptions, affinitizedDestination); + + Assert.False(context.Response.Headers.ContainsKey(AffinityHeaderName)); + } + + public static IEnumerable FindAffinitizedDestination_CustomHeaderNameIsNotSpecified_Cases() + { + yield return new object[] { null }; + yield return new object[] { new Dictionary { { "SomeSetting", "SomeValue" } } }; + } + } +} diff --git a/test/ReverseProxy.Tests/Service/SessionAffinity/RedistributeAffinityFailurePolicyTests.cs b/test/ReverseProxy.Tests/Service/SessionAffinity/RedistributeAffinityFailurePolicyTests.cs new file mode 100644 index 000000000..5b47c1128 --- /dev/null +++ b/test/ReverseProxy.Tests/Service/SessionAffinity/RedistributeAffinityFailurePolicyTests.cs @@ -0,0 +1,37 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using System.Collections.Generic; +using System.Threading.Tasks; +using Microsoft.AspNetCore.Http; +using Microsoft.ReverseProxy.Abstractions.BackendDiscovery.Contract; +using Xunit; + +namespace Microsoft.ReverseProxy.Service.SessionAffinity +{ + public class RedistributeAffinityFailurePolicyTests + { + [Theory] + [InlineData(AffinityStatus.AffinityKeyExtractionFailed)] + [InlineData(AffinityStatus.DestinationNotFound)] + public async Task Handle_FailedAffinityStatus_ReturnTrue(AffinityStatus status) + { + var policy = new RedistributeAffinityFailurePolicy(); + + Assert.Equal(SessionAffinityConstants.AffinityFailurePolicies.Redistribute, policy.Name); + Assert.True(await policy.Handle(new DefaultHttpContext(), default, status)); + } + + [Theory] + [InlineData(AffinityStatus.OK)] + [InlineData(AffinityStatus.AffinityKeyNotSet)] + public async Task Handle_SuccessfulAffinityStatus_Throw(AffinityStatus status) + { + var policy = new RedistributeAffinityFailurePolicy(); + var context = new DefaultHttpContext(); + + await Assert.ThrowsAsync(() => policy.Handle(context, default, status)); + } + } +} diff --git a/test/ReverseProxy.Tests/Service/SessionAffinity/Return503ErrorAffinityFailurePolicyTests.cs b/test/ReverseProxy.Tests/Service/SessionAffinity/Return503ErrorAffinityFailurePolicyTests.cs new file mode 100644 index 000000000..63d7ff383 --- /dev/null +++ b/test/ReverseProxy.Tests/Service/SessionAffinity/Return503ErrorAffinityFailurePolicyTests.cs @@ -0,0 +1,39 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using System.Threading.Tasks; +using Microsoft.AspNetCore.Http; +using Microsoft.ReverseProxy.Abstractions.BackendDiscovery.Contract; +using Xunit; + +namespace Microsoft.ReverseProxy.Service.SessionAffinity +{ + public class Return503ErrorAffinityFailurePolicyTests + { + [Theory] + [InlineData(AffinityStatus.DestinationNotFound)] + [InlineData(AffinityStatus.AffinityKeyExtractionFailed)] + public async Task Handle_FaultyAffinityStatus_RespondWith503(AffinityStatus status) + { + var policy = new Return503ErrorAffinityFailurePolicy(); + var context = new DefaultHttpContext(); + + Assert.Equal(SessionAffinityConstants.AffinityFailurePolicies.Return503Error, policy.Name); + + Assert.False(await policy.Handle(context, default, status)); + Assert.Equal(503, context.Response.StatusCode); + } + + [Theory] + [InlineData(AffinityStatus.OK)] + [InlineData(AffinityStatus.AffinityKeyNotSet)] + public async Task Handle_SuccessfulAffinityStatus_Throw(AffinityStatus status) + { + var policy = new Return503ErrorAffinityFailurePolicy(); + var context = new DefaultHttpContext(); + + await Assert.ThrowsAsync(() => policy.Handle(context, default, status)); + } + } +}