[PM-5645] Cosmos DB Grant Storage (#3634)

* table storage grants

* simple shard on storage accounts

* use is not

* cosmos grant repo

* remove single storage connection string

* some fixes to dapper grant repo

* pattern matching

* add fallback to base PersistedGrantStore

* service collection extension cleanup

* cleanup

* remove unused Id

* empty string rowkey

* fix sharding method logic

* ttl for cosmos

* make ttl an int

* fixes to cosmos implementation

* fix partition key values

* catch notfound exceptions

* indenting

* update grantitem with custom serialization

* use new transform helpers

* grantloader perf test tool

* ref

* remove grant loader project

* remove table storage implementation

* remove table storage stuff

* all redis fallback to build to null

* revert sln file change

* EOF new line

* remove trailing comma

* lint fixes

* add grant to names

* move cosmos serilaizer to utils

* add some .net 8 keyed service comments

* EnableContentResponseOnWrite

* Fix type in EF grant repository

src/Core/Auth/Entities/Grant.cs

@@ -1,10 +1,28 @@
 #nullable enable
 using System.ComponentModel.DataAnnotations;
+using Bit.Core.Auth.Models.Data;
+using Duende.IdentityServer.Models;
 
 namespace Bit.Core.Auth.Entities;
 
-public class Grant
+public class Grant : IGrant
 {
+    public Grant() { }
+
+    public Grant(PersistedGrant pGrant)
+    {
+        Key = pGrant.Key;
+        Type = pGrant.Type;
+        SubjectId = pGrant.SubjectId;
+        SessionId = pGrant.SessionId;
+        ClientId = pGrant.ClientId;
+        Description = pGrant.Description;
+        CreationDate = pGrant.CreationTime;
+        ExpirationDate = pGrant.Expiration;
+        ConsumedDate = pGrant.ConsumedTime;
+        Data = pGrant.Data;
+    }
+
     public int Id { get; set; }
     [MaxLength(200)]
     public string Key { get; set; } = null!;

src/Core/Auth/Models/Data/GrantItem.cs

@@ -0,0 +1,77 @@
+using System.Text.Json.Serialization;
+using Bit.Core.Auth.Repositories.Cosmos;
+using Duende.IdentityServer.Models;
+
+namespace Bit.Core.Auth.Models.Data;
+
+public class GrantItem : IGrant
+{
+    public GrantItem() { }
+
+    public GrantItem(PersistedGrant pGrant)
+    {
+        Key = pGrant.Key;
+        Type = pGrant.Type;
+        SubjectId = pGrant.SubjectId;
+        SessionId = pGrant.SessionId;
+        ClientId = pGrant.ClientId;
+        Description = pGrant.Description;
+        CreationDate = pGrant.CreationTime;
+        ExpirationDate = pGrant.Expiration;
+        ConsumedDate = pGrant.ConsumedTime;
+        Data = pGrant.Data;
+        SetTtl();
+    }
+
+    public GrantItem(IGrant g)
+    {
+        Key = g.Key;
+        Type = g.Type;
+        SubjectId = g.SubjectId;
+        SessionId = g.SessionId;
+        ClientId = g.ClientId;
+        Description = g.Description;
+        CreationDate = g.CreationDate;
+        ExpirationDate = g.ExpirationDate;
+        ConsumedDate = g.ConsumedDate;
+        Data = g.Data;
+        SetTtl();
+    }
+
+    [JsonPropertyName("id")]
+    [JsonConverter(typeof(Base64IdStringConverter))]
+    public string Key { get; set; }
+    [JsonPropertyName("typ")]
+    public string Type { get; set; }
+    [JsonPropertyName("sub")]
+    public string SubjectId { get; set; }
+    [JsonPropertyName("sid")]
+    public string SessionId { get; set; }
+    [JsonPropertyName("cid")]
+    public string ClientId { get; set; }
+    [JsonPropertyName("des")]
+    public string Description { get; set; }
+    [JsonPropertyName("cre")]
+    public DateTime CreationDate { get; set; } = DateTime.UtcNow;
+    [JsonPropertyName("exp")]
+    public DateTime? ExpirationDate { get; set; }
+    [JsonPropertyName("con")]
+    public DateTime? ConsumedDate { get; set; }
+    [JsonPropertyName("data")]
+    public string Data { get; set; }
+    // https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/how-to-time-to-live?tabs=dotnet-sdk-v3#set-time-to-live-on-an-item-using-an-sdk
+    [JsonPropertyName("ttl")]
+    public int Ttl { get; set; } = -1;
+
+    public void SetTtl()
+    {
+        if (ExpirationDate != null)
+        {
+            var sec = (ExpirationDate.Value - DateTime.UtcNow).TotalSeconds;
+            if (sec > 0)
+            {
+                Ttl = (int)sec;
+            }
+        }
+    }
+}

src/Core/Auth/Models/Data/IGrant.cs

@@ -0,0 +1,15 @@
+namespace Bit.Core.Auth.Models.Data;
+
+public interface IGrant
+{
+    string Key { get; set; }
+    string Type { get; set; }
+    string SubjectId { get; set; }
+    string SessionId { get; set; }
+    string ClientId { get; set; }
+    string Description { get; set; }
+    DateTime CreationDate { get; set; }
+    DateTime? ExpirationDate { get; set; }
+    DateTime? ConsumedDate { get; set; }
+    string Data { get; set; }
+}

src/Core/Auth/Repositories/Cosmos/Base64IdStringConverter.cs

@@ -0,0 +1,32 @@
+using System.Text.Json;
+using System.Text.Json.Serialization;
+using Bit.Core.Utilities;
+
+namespace Bit.Core.Auth.Repositories.Cosmos;
+
+public class Base64IdStringConverter : JsonConverter<string>
+{
+    public override string Read(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options) =>
+        ToKey(reader.GetString());
+
+    public override void Write(Utf8JsonWriter writer, string value, JsonSerializerOptions options) =>
+        writer.WriteStringValue(ToId(value));
+
+    public static string ToId(string key)
+    {
+        if (key == null)
+        {
+            return null;
+        }
+        return CoreHelpers.TransformToBase64Url(key);
+    }
+
+    public static string ToKey(string id)
+    {
+        if (id == null)
+        {
+            return null;
+        }
+        return CoreHelpers.TransformFromBase64Url(id);
+    }
+}

src/Core/Auth/Repositories/Cosmos/GrantRepository.cs

@@ -0,0 +1,81 @@
+using System.Net;
+using System.Text.Json;
+using System.Text.Json.Serialization;
+using Bit.Core.Auth.Models.Data;
+using Bit.Core.Settings;
+using Bit.Core.Utilities;
+using Microsoft.Azure.Cosmos;
+
+namespace Bit.Core.Auth.Repositories.Cosmos;
+
+public class GrantRepository : IGrantRepository
+{
+    private readonly CosmosClient _client;
+    private readonly Database _database;
+    private readonly Container _container;
+
+    public GrantRepository(GlobalSettings globalSettings)
+        : this(globalSettings.IdentityServer.CosmosConnectionString)
+    { }
+
+    public GrantRepository(string cosmosConnectionString)
+    {
+        var options = new CosmosClientOptions
+        {
+            Serializer = new SystemTextJsonCosmosSerializer(new JsonSerializerOptions
+            {
+                DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull,
+                WriteIndented = false
+            })
+        };
+        // TODO: Perhaps we want to evaluate moving this to DI as a keyed service singleton in .NET 8
+        _client = new CosmosClient(cosmosConnectionString, options);
+        _database = _client.GetDatabase("identity");
+        _container = _database.GetContainer("grant");
+    }
+
+    public async Task<IGrant> GetByKeyAsync(string key)
+    {
+        var id = Base64IdStringConverter.ToId(key);
+        try
+        {
+            var response = await _container.ReadItemAsync<GrantItem>(id, new PartitionKey(id));
+            return response.Resource;
+        }
+        catch (CosmosException e)
+        {
+            if (e.StatusCode == HttpStatusCode.NotFound)
+            {
+                return null;
+            }
+            throw;
+        }
+    }
+
+    public Task<ICollection<IGrant>> GetManyAsync(string subjectId, string sessionId, string clientId, string type)
+        => throw new NotImplementedException();
+
+    public async Task SaveAsync(IGrant obj)
+    {
+        if (obj is not GrantItem item)
+        {
+            item = new GrantItem(obj);
+        }
+        item.SetTtl();
+        var id = Base64IdStringConverter.ToId(item.Key);
+        await _container.UpsertItemAsync(item, new PartitionKey(id), new ItemRequestOptions
+        {
+            // ref: https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/best-practice-dotnet#best-practices-for-write-heavy-workloads
+            EnableContentResponseOnWrite = false
+        });
+    }
+
+    public async Task DeleteByKeyAsync(string key)
+    {
+        var id = Base64IdStringConverter.ToId(key);
+        await _container.DeleteItemAsync<IGrant>(id, new PartitionKey(id));
+    }
+
+    public Task DeleteManyAsync(string subjectId, string sessionId, string clientId, string type)
+        => throw new NotImplementedException();
+}

src/Core/Auth/Repositories/IGrantRepository.cs

@@ -1,12 +1,12 @@
-using Bit.Core.Auth.Entities;
+using Bit.Core.Auth.Models.Data;
 
 namespace Bit.Core.Auth.Repositories;
 
 public interface IGrantRepository
 {
-    Task<Grant> GetByKeyAsync(string key);
-    Task<ICollection<Grant>> GetManyAsync(string subjectId, string sessionId, string clientId, string type);
-    Task SaveAsync(Grant obj);
+    Task<IGrant> GetByKeyAsync(string key);
+    Task<ICollection<IGrant>> GetManyAsync(string subjectId, string sessionId, string clientId, string type);
+    Task SaveAsync(IGrant obj);
     Task DeleteByKeyAsync(string key);
     Task DeleteManyAsync(string subjectId, string sessionId, string clientId, string type);
 }

src/Core/Core.csproj

@@ -37,7 +37,7 @@
     <PackageReference Include="Microsoft.Azure.NotificationHubs" Version="4.1.0" />
     <PackageReference Include="Microsoft.Data.SqlClient" Version="5.0.1" />
     <!-- Azure.Identity is a explicit dependency to Microsoft.Data.SqlClient -->
-    <PackageReference Include="Azure.Identity" Version="1.10.2"/>
+    <PackageReference Include="Azure.Identity" Version="1.10.2" />
     <PackageReference Include="Microsoft.Extensions.Configuration.EnvironmentVariables" Version="6.0.1" />
     <PackageReference Include="Microsoft.Extensions.Configuration.UserSecrets" Version="6.0.1" />
     <PackageReference Include="Microsoft.Extensions.Identity.Stores" Version="6.0.25" />

src/Core/Settings/GlobalSettings.cs

@@ -327,6 +327,7 @@ public class GlobalSettings : IGlobalSettings
         public string CertificateThumbprint { get; set; }
         public string CertificatePassword { get; set; }
         public string RedisConnectionString { get; set; }
+        public string CosmosConnectionString { get; set; }
         public string LicenseKey { get; set; } = "eyJhbGciOiJQUzI1NiIsImtpZCI6IklkZW50aXR5U2VydmVyTGljZW5zZWtleS83Y2VhZGJiNzgxMzA0NjllODgwNjg5MTAyNTQxNGYxNiIsInR5cCI6ImxpY2Vuc2Urand0In0.eyJpc3MiOiJodHRwczovL2R1ZW5kZXNvZnR3YXJlLmNvbSIsImF1ZCI6IklkZW50aXR5U2VydmVyIiwiaWF0IjoxNzAxODIwODAwLCJleHAiOjE3MzM0NDMyMDAsImNvbXBhbnlfbmFtZSI6IkJpdHdhcmRlbiBJbmMuIiwiY29udGFjdF9pbmZvIjoiY29udGFjdEBkdWVuZGVzb2Z0d2FyZS5jb20iLCJlZGl0aW9uIjoiU3RhcnRlciIsImlkIjoiNDMxOSIsImZlYXR1cmUiOlsiaXN2IiwidW5saW1pdGVkX2NsaWVudHMiXSwicHJvZHVjdCI6IkJpdHdhcmRlbiJ9.iLA771PffgIh0ClRS8OWHbg2cAgjhgOkUjRRkLNr9dpQXhYZkVKdpUn-Gw9T7grsGcAx0f4p-TQmtcCpbN9EJCF5jlF0-NfsRTp_gmCgQ5eXyiE4DzJp2OCrz_3STf07N1dILwhD3nk9rzcA6SRQ4_kja8wAMHKnD5LisW98r5DfRDBecRs16KS5HUhg99DRMR5fd9ntfydVMTC_E23eEOHVLsR4YhiSXaEINPjFDG1czyOBClJItDW8g9X8qlClZegr630UjnKKg06A4usoL25VFHHn8Ew3v-_-XdlWoWsIpMMVvacwZT8rwkxjIesFNsXG6yzuROIhaxAvB1297A";
     }
 

src/Core/Utilities/CoreHelpers.cs

@@ -338,16 +338,50 @@ public static class CoreHelpers
         return Encoding.UTF8.GetString(Base64UrlDecode(input));
     }
 
+    /// <summary>
+    /// Encodes a Base64 URL formatted string.
+    /// </summary>
+    /// <param name="input">Byte data</param>
+    /// <returns>Base64 URL formatted string</returns>
     public static string Base64UrlEncode(byte[] input)
     {
-        var output = Convert.ToBase64String(input)
+        // Standard base64 encoder
+        var standardB64 = Convert.ToBase64String(input);
+        return TransformToBase64Url(standardB64);
+    }
+
+    /// <summary>
+    /// Transforms a Base64 standard formatted string to a Base64 URL formatted string.
+    /// </summary>
+    /// <param name="input">Base64 standard formatted string</param>
+    /// <returns>Base64 URL formatted string</returns>
+    public static string TransformToBase64Url(string input)
+    {
+        var output = input
             .Replace('+', '-')
             .Replace('/', '_')
             .Replace("=", string.Empty);
         return output;
     }
 
+    /// <summary>
+    /// Decodes a Base64 URL formatted string.
+    /// </summary>
+    /// <param name="input">Base64 URL formatted string</param>
+    /// <returns>Data as bytes</returns>
     public static byte[] Base64UrlDecode(string input)
+    {
+        var standardB64 = TransformFromBase64Url(input);
+        // Standard base64 decoder
+        return Convert.FromBase64String(standardB64);
+    }
+
+    /// <summary>
+    /// Transforms a Base64 URL formatted string to a Base64 standard formatted string.
+    /// </summary>
+    /// <param name="input">Base64 URL formatted string</param>
+    /// <returns>Base64 standard formatted string</returns>
+    public static string TransformFromBase64Url(string input)
     {
         var output = input;
         // 62nd char of encoding
@@ -370,8 +404,8 @@ public static class CoreHelpers
                 throw new InvalidOperationException("Illegal base64url string!");
         }
 
-        // Standard base64 decoder
-        return Convert.FromBase64String(output);
+        // Standard base64 string output
+        return output;
     }
 
     public static string PunyEncode(string text)

src/Core/Utilities/SystemTextJsonCosmosSerializer.cs

@@ -0,0 +1,40 @@
+using System.Text.Json;
+using Azure.Core.Serialization;
+using Microsoft.Azure.Cosmos;
+
+namespace Bit.Core.Utilities;
+
+// ref: https://github.com/Azure/azure-cosmos-dotnet-v3/blob/master/Microsoft.Azure.Cosmos.Samples/Usage/SystemTextJson/CosmosSystemTextJsonSerializer.cs
+public class SystemTextJsonCosmosSerializer : CosmosSerializer
+{
+    private readonly JsonObjectSerializer _systemTextJsonSerializer;
+
+    public SystemTextJsonCosmosSerializer(JsonSerializerOptions jsonSerializerOptions)
+    {
+        _systemTextJsonSerializer = new JsonObjectSerializer(jsonSerializerOptions);
+    }
+
+    public override T FromStream<T>(Stream stream)
+    {
+        using (stream)
+        {
+            if (stream.CanSeek && stream.Length == 0)
+            {
+                return default;
+            }
+            if (typeof(Stream).IsAssignableFrom(typeof(T)))
+            {
+                return (T)(object)stream;
+            }
+            return (T)_systemTextJsonSerializer.Deserialize(stream, typeof(T), default);
+        }
+    }
+
+    public override Stream ToStream<T>(T input)
+    {
+        var streamPayload = new MemoryStream();
+        _systemTextJsonSerializer.Serialize(streamPayload, input, input.GetType(), default);
+        streamPayload.Position = 0;
+        return streamPayload;
+    }
+}