[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);
 }