[PM-5293] Redis for Grants (#3577)

* Add Initial Redis Implementation

* Format

* Add Key to PersistedGrant

* Reference Identity In Microbenchmark Project

* Allow Filterable Benchmarks

* Use Shorter Key And Cast to RedisKey Once

* Add RedisPersistedGrantStore Benchmarks

* Run restore

* Format

* Update ID4 References

* Make RedisGrantStore Singleton

* Use MessagePack

* Use Cached Options

* Turn off Compression

* Minor Feedback

* Add Docs to StorablePersistedGrant

* Use existing Identity Redis

---------

Co-authored-by: Matt Bishop <mbishop@bitwarden.com>

src/Identity/Identity.csproj

@@ -13,6 +13,7 @@
   </ItemGroup>
 
   <ItemGroup>
+    <PackageReference Include="MessagePack" Version="2.5.140" />
     <PackageReference Include="Swashbuckle.AspNetCore.SwaggerGen" Version="6.5.0" />
   </ItemGroup>
 

src/Identity/IdentityServer/RedisPersistedGrantStore.cs

@@ -0,0 +1,181 @@
+using System.Diagnostics;
+using Duende.IdentityServer.Models;
+using Duende.IdentityServer.Stores;
+using MessagePack;
+using StackExchange.Redis;
+
+namespace Bit.Identity.IdentityServer;
+
+/// <summary>
+/// A <see cref="IPersistedGrantStore"/> that persists its grants on a Redis DB
+/// </summary>
+/// <remarks>
+/// This store also allows a fallback to another store in the case that a key was not found
+/// in the Redis DB or the Redis DB happens to be down.
+/// </remarks>
+public class RedisPersistedGrantStore : IPersistedGrantStore
+{
+    private static readonly MessagePackSerializerOptions _options = MessagePackSerializerOptions.Standard;
+    private readonly IConnectionMultiplexer _connectionMultiplexer;
+    private readonly ILogger<RedisPersistedGrantStore> _logger;
+    private readonly IPersistedGrantStore _fallbackGrantStore;
+
+    public RedisPersistedGrantStore(
+        IConnectionMultiplexer connectionMultiplexer,
+        ILogger<RedisPersistedGrantStore> logger,
+        IPersistedGrantStore fallbackGrantStore)
+    {
+        _connectionMultiplexer = connectionMultiplexer;
+        _logger = logger;
+        _fallbackGrantStore = fallbackGrantStore;
+    }
+
+    public Task<IEnumerable<PersistedGrant>> GetAllAsync(PersistedGrantFilter filter)
+    {
+        _logger.LogWarning("Redis does not implement 'GetAllAsync', Skipping.");
+        return Task.FromResult(Enumerable.Empty<PersistedGrant>());
+    }
+    public async Task<PersistedGrant> GetAsync(string key)
+    {
+        try
+        {
+            if (!_connectionMultiplexer.IsConnected)
+            {
+                // Redis is down, fallback to using SQL table
+                _logger.LogWarning("This is not connected, using fallback store to execute 'GetAsync' with {Key}.", key);
+                return await _fallbackGrantStore.GetAsync(key);
+            }
+
+            var redisKey = CreateRedisKey(key);
+
+            var redisDb = _connectionMultiplexer.GetDatabase();
+            var redisValueAndExpiry = await redisDb.StringGetWithExpiryAsync(redisKey);
+
+            if (!redisValueAndExpiry.Value.HasValue)
+            {
+                // It wasn't found, there is a chance is was instead stored in the fallback store
+                _logger.LogWarning("Could not find grant in primary store, using fallback one.");
+                return await _fallbackGrantStore.GetAsync(key);
+            }
+
+            Debug.Assert(redisValueAndExpiry.Expiry.HasValue, "Redis entry is expected to have an expiry.");
+
+            var storablePersistedGrant = MessagePackSerializer.Deserialize<StorablePersistedGrant>(redisValueAndExpiry.Value, _options);
+
+            return new PersistedGrant
+            {
+                Key = key,
+                Type = storablePersistedGrant.Type,
+                SubjectId = storablePersistedGrant.SubjectId,
+                SessionId = storablePersistedGrant.SessionId,
+                ClientId = storablePersistedGrant.ClientId,
+                Description = storablePersistedGrant.Description,
+                CreationTime = storablePersistedGrant.CreationTime,
+                ConsumedTime = storablePersistedGrant.ConsumedTime,
+                Data = storablePersistedGrant.Data,
+                Expiration = storablePersistedGrant.CreationTime.Add(redisValueAndExpiry.Expiry.Value),
+            };
+        }
+        catch (Exception ex)
+        {
+            _logger.LogWarning(ex, "Failure in 'GetAsync' using primary grant store, falling back.");
+            return await _fallbackGrantStore.GetAsync(key);
+        }
+    }
+
+    public Task RemoveAllAsync(PersistedGrantFilter filter)
+    {
+        _logger.LogWarning("This does not implement 'RemoveAllAsync', Skipping.");
+        return Task.CompletedTask;
+    }
+
+    // This method is not actually expected to get called and instead redis will just get rid of the expired items
+    public async Task RemoveAsync(string key)
+    {
+        if (!_connectionMultiplexer.IsConnecting)
+        {
+            _logger.LogWarning("Redis is not connected, using fallback store to execute 'RemoveAsync', with {Key}", key);
+            await _fallbackGrantStore.RemoveAsync(key);
+        }
+
+        var redisDb = _connectionMultiplexer.GetDatabase();
+        await redisDb.KeyDeleteAsync(CreateRedisKey(key));
+    }
+
+    public async Task StoreAsync(PersistedGrant grant)
+    {
+        try
+        {
+            if (!_connectionMultiplexer.IsConnected)
+            {
+                _logger.LogWarning("Redis is not connected, using fallback store to execute 'StoreAsync', with {Key}", grant.Key);
+                await _fallbackGrantStore.StoreAsync(grant);
+            }
+
+            if (!grant.Expiration.HasValue)
+            {
+                throw new ArgumentException("A PersistedGrant is always expected to include an expiration time.");
+            }
+
+            var redisDb = _connectionMultiplexer.GetDatabase();
+
+            var redisKey = CreateRedisKey(grant.Key);
+
+            var serializedGrant = MessagePackSerializer.Serialize(new StorablePersistedGrant
+            {
+                Type = grant.Type,
+                SubjectId = grant.SubjectId,
+                SessionId = grant.SessionId,
+                ClientId = grant.ClientId,
+                Description = grant.Description,
+                CreationTime = grant.CreationTime,
+                ConsumedTime = grant.ConsumedTime,
+                Data = grant.Data,
+            }, _options);
+
+            await redisDb.StringSetAsync(redisKey, serializedGrant, grant.Expiration.Value - grant.CreationTime);
+        }
+        catch (Exception ex)
+        {
+            _logger.LogWarning(ex, "Failure in 'StoreAsync' using primary grant store, falling back.");
+            await _fallbackGrantStore.StoreAsync(grant);
+        }
+    }
+
+    private static RedisKey CreateRedisKey(string key)
+    {
+        return $"grant-{key}";
+    }
+
+    // This is a slimmer version of PersistedGrant that removes the Key since that will be used as the key in Redis
+    // it also strips out the ExpirationDate since we use that to store the TTL and read that out when we retrieve the
+    // object and can use that information to fill in the Expiration property on PersistedGrant
+    // TODO: .NET 8 Make all properties required
+    [MessagePackObject]
+    public class StorablePersistedGrant
+    {
+        [Key(0)]
+        public string Type { get; set; }
+
+        [Key(1)]
+        public string SubjectId { get; set; }
+
+        [Key(2)]
+        public string SessionId { get; set; }
+
+        [Key(3)]
+        public string ClientId { get; set; }
+
+        [Key(4)]
+        public string Description { get; set; }
+
+        [Key(5)]
+        public DateTime CreationTime { get; set; }
+
+        [Key(6)]
+        public DateTime? ConsumedTime { get; set; }
+
+        [Key(7)]
+        public string Data { get; set; }
+    }
+}

src/Identity/Utilities/ServiceCollectionExtensions.cs

@@ -1,10 +1,12 @@
 using Bit.Core.IdentityServer;
 using Bit.Core.Settings;
+using Bit.Core.Utilities;
 using Bit.Identity.IdentityServer;
 using Bit.SharedWeb.Utilities;
 using Duende.IdentityServer.ResponseHandling;
 using Duende.IdentityServer.Services;
 using Duende.IdentityServer.Stores;
+using StackExchange.Redis;
 
 namespace Bit.Identity.Utilities;
 
@@ -45,11 +47,34 @@ public static class ServiceCollectionExtensions
             .AddCustomTokenRequestValidator<CustomTokenRequestValidator>()
             .AddProfileService<ProfileService>()
             .AddResourceOwnerValidator<ResourceOwnerPasswordValidator>()
-            .AddPersistedGrantStore<PersistedGrantStore>()
             .AddClientStore<ClientStore>()
             .AddIdentityServerCertificate(env, globalSettings)
             .AddExtensionGrantValidator<WebAuthnGrantValidator>();
 
+        if (CoreHelpers.SettingHasValue(globalSettings.IdentityServer.RedisConnectionString))
+        {
+            // If we have redis, prefer it
+
+            // Add the original persisted grant store via it's implementation type
+            // so we can inject it right after.
+            services.AddSingleton<PersistedGrantStore>();
+
+            services.AddSingleton<IPersistedGrantStore>(sp =>
+            {
+                return new RedisPersistedGrantStore(
+                    // TODO: .NET 8 create a keyed service for this connection multiplexer and even PersistedGrantStore
+                    ConnectionMultiplexer.Connect(globalSettings.IdentityServer.RedisConnectionString),
+                    sp.GetRequiredService<ILogger<RedisPersistedGrantStore>>(),
+                    sp.GetRequiredService<PersistedGrantStore>() // Fallback grant store
+                );
+            });
+        }
+        else
+        {
+            // Use the original grant store
+            identityServerBuilder.AddPersistedGrantStore<PersistedGrantStore>();
+        }
+
         services.AddTransient<ICorsPolicyService, CustomCorsPolicyService>();
         return identityServerBuilder;
     }