Cosmos: Support AAD RBAC via the ClientSecretCredential

[JeremyLikness]

The Azure Cosmos DB SDK now supports role-based access control (RBAC) via Azure Active Directory (AAD) token credentials.

Instructions: Configure role-based access control for your Azure Cosmos DB account with Azure Active Directory

This more direct support should be recommended over resource tokens.

[jeanpaulsmit]

Great, I hope this will end up high on the priority list!
Not having to store the Cosmos db access keys somewhere, is the best security measure.

[surenderssm]

@ajcvickers - Any high level ETA on this?
We are on a mission of having Zero Secrets in our components by using Azure managed identity.
This is one of our blocker :( !!

[ajcvickers]

@surenderssm As the milestone indicates, this is currently planned for the next release; that is EF Core 7.0, which will be released next November.

[surenderssm]

@ajcvickers - Thank you for the update.
Are there any interim reliable solutions which can help us on this front?

[ajcvickers]

/cc @JeremyLikness

[dmitriyi-affinity]

EF Core and CosmosDb provider are pretty extensible.
Yes now passthrough TokenCredential to the underlying CosmosClient is not supported, but it possible to override ISingletonCosmosClientWrapper and CosmosOptionsExtension to extended/patched versions and add support yourself.

[wkoeter]

It is a lot more trouble than you let seem. It really is not that extendable.

You need to extend/implement/duplicate (at least)

• CosmosOptionsExtension
• CosmosServiceCollectionExtensions
• CosmosSingletonOptions
• ISingletonCosmosClientWrapper
• ICosmosClientWrapper
• OptionsExtension

The already swapped CosmosOptionsExtension is referenced hardcoded in CosmosClientWrapper constructor using a servicelocator leading to a NullReferenceException because you implemented a custom one. Duplicating will cascade to more internals that need duplication.

• Duplication of CosmosClientWrapper will cascade
• Implementing CosmosClientWrapper yourself is overkill
• Registering CustomCosmosOptionsExtension under parent type CosmosOptionsExtension using the generic arguments is not an option. did not work because somewhere internal in DbContextOptions<TContext> it does not use the T generic argument supplied, but again does a GetType() when registering internally. So it did not register as CosmosOptionsExtension but as CustomCosmosOptionsExtension

Working solution:
Got it to work by adding both the framework's CosmosOptionsExtension and the custom CustomCosmosOptionsExtension to the optionsBuilder. This because there are hardcoded references to CosmosOptionsExtension. You cannot skip it..

Make sure to add the CustomCosmosOptionsExtension first, because that makes sure the custom ISingletonCosmosClientWrapper is registered before the standard one. And there is where you add the ManagedIdentityClientId for the DefaultAzureCredential.

[ChrisKlug]

Is there a sample of how to do this anywhere? I've tried reading the above message, but I am not familiar enough with the internals of EF Core to figure it out unfortunately...

[dmitriyi-affinity]

I forced to get it work in our production project. Will try to find a time to share the code.

[wkoeter]

Hi @ChrisKlug
Here is a gist of the changes I made to get it to work. It involves some duplication and customization.

https://gist.github.com/wkoeter/4ed90c7c8f61e3b3a52d2667d5a7c856

[ChrisKlug]

Thank you so much! That is VERY helpful!

[dmitriyi-affinity]

The new DI API that support registration Cosmos with the TokenCredential:

public static class CosmosEntityFrameworkDIExtensions
{
    [SuppressMessage(
        "Usage",
        "EF1001:Internal EF Core API usage.",
        Justification = "Yes we extending internals, hope with EF 7.0 it will be out of the box.")]
    public static DbContextOptionsBuilder UseCosmos(
        this DbContextOptionsBuilder optionsBuilder,
        string accountEndpoint,
        TokenCredential tokenCredential,
        string databaseName,
        Action<CosmosDbContextOptionsBuilder>? cosmosOptionsAction = null)
    {
        if (optionsBuilder == null)
        {
            throw new ArgumentNullException(nameof(optionsBuilder));
        }

        if (tokenCredential == null)
        {
            throw new ArgumentNullException(nameof(tokenCredential));
        }

        var extension = optionsBuilder.Options.FindExtension<CosmosOptionsExtension>() ?? new CosmosOptionsExtension();
        extension = extension
            .WithAccountEndpoint(accountEndpoint)
            .WithAccountKey("__TokenCredential__")
            .WithDatabaseName(databaseName);

        ((IDbContextOptionsBuilderInfrastructure)optionsBuilder).AddOrUpdateExtension(extension);

        var tokenCredentialExtension = new CosmosTokenCredentialOptionsExtension(tokenCredential);
        ((IDbContextOptionsBuilderInfrastructure)optionsBuilder).AddOrUpdateExtension(tokenCredentialExtension);

        cosmosOptionsAction?.Invoke(new CosmosDbContextOptionsBuilder(optionsBuilder));

        // ------------------- Substitution is happening HERE --------------------------
        optionsBuilder.ReplaceService<ISingletonCosmosClientWrapper, CosmosClientWrapperWithTokensSupport>();
        return optionsBuilder;
    }
}

Supporting classes that needs by the new DI API:

[SuppressMessage("Design", "CA1063:Implement IDisposable Correctly", Justification = "Simplified pattern")]
[SuppressMessage("Usage", "CA1816:Dispose methods should call SuppressFinalize", Justification = "Simplified pattern")]
[SuppressMessage("Usage", "EF1001:Internal EF Core API usage.", Justification = "EF Shim")]
[UsedImplicitly(ImplicitUseKindFlags.InstantiatedWithFixedConstructorSignature)]
public class CosmosClientWrapperWithTokensSupport : ISingletonCosmosClientWrapper
{
    private static readonly string UserAgent = " Microsoft.EntityFrameworkCore.Cosmos/" + ProductInfo.GetVersion();
    private readonly CosmosClientOptions _options;
    private readonly string? _endpoint;
    private readonly string? _key;
    private readonly string? _connectionString;
    private readonly TokenCredential? _tokenCredential;
    private CosmosClient? _client;

    [SuppressMessage("Design", "CA1062:Validate arguments of public methods", Justification = "EF Codebase style.")]
    public CosmosClientWrapperWithTokensSupport(ICosmosSingletonOptions options, CosmosTokenCredentialSingletonOptions singletonOptions)
    {
        _endpoint = options.AccountEndpoint;
        _key = options.AccountKey;
        _connectionString = options.ConnectionString;
        var configuration = new CosmosClientOptions
            { ApplicationName = UserAgent, Serializer = new JsonCosmosSerializer() };

        if (options.Region != null)
        {
            configuration.ApplicationRegion = options.Region;
        }

        if (options.LimitToEndpoint != null)
        {
            configuration.LimitToEndpoint = options.LimitToEndpoint.Value;
        }

        if (options.ConnectionMode != null)
        {
            configuration.ConnectionMode = options.ConnectionMode.Value;
        }

        if (options.WebProxy != null)
        {
            configuration.WebProxy = options.WebProxy;
        }

        if (options.RequestTimeout != null)
        {
            configuration.RequestTimeout = options.RequestTimeout.Value;
        }

        if (options.OpenTcpConnectionTimeout != null)
        {
            configuration.OpenTcpConnectionTimeout = options.OpenTcpConnectionTimeout.Value;
        }

        if (options.IdleTcpConnectionTimeout != null)
        {
            configuration.IdleTcpConnectionTimeout = options.IdleTcpConnectionTimeout.Value;
        }

        if (options.GatewayModeMaxConnectionLimit != null)
        {
            configuration.GatewayModeMaxConnectionLimit = options.GatewayModeMaxConnectionLimit.Value;
        }

        if (options.MaxTcpConnectionsPerEndpoint != null)
        {
            configuration.MaxTcpConnectionsPerEndpoint = options.MaxTcpConnectionsPerEndpoint.Value;
        }

        if (options.MaxRequestsPerTcpConnection != null)
        {
            configuration.MaxRequestsPerTcpConnection = options.MaxRequestsPerTcpConnection.Value;
        }

        if (options.HttpClientFactory != null)
        {
            configuration.HttpClientFactory = options.HttpClientFactory;
        }

        _tokenCredential = singletonOptions.TokenCredential;

        _options = configuration;
    }

    public virtual CosmosClient Client
        => _client ??= CreateCosmosClient();

    /// <inheritdoc />
    public void Dispose()
    {
        _client?.Dispose();
        _client = null;
    }

    private CosmosClient CreateCosmosClient()
    {
        if (_tokenCredential != null)
        {
            return new CosmosClient(_endpoint, _tokenCredential, _options);
        }

        return string.IsNullOrEmpty(_connectionString)
            ? new CosmosClient(_endpoint, _key, _options)
            : new CosmosClient(_connectionString, _options);
    }
}

[SuppressMessage("Usage", "EF1001:Internal EF Core API usage.", Justification = "OK, we will track all changes.")]
[SuppressMessage("ReSharper", "ClassWithVirtualMembersNeverInherited.Global", Justification = "Reviewed")]
public class CosmosTokenCredentialOptionsExtension : IDbContextOptionsExtension
{
    private DbContextOptionsExtensionInfo? _info;

    public CosmosTokenCredentialOptionsExtension(TokenCredential? tokenCredential)
    {
        TokenCredential = tokenCredential;
    }

    public TokenCredential? TokenCredential { get; }

    public DbContextOptionsExtensionInfo Info => _info ??= new ExtensionInfo(this);

    public void ApplyServices(IServiceCollection services)
    {
        new EntityFrameworkServicesBuilder(services)
            .TryAdd<ISingletonOptions, CosmosTokenCredentialSingletonOptions>(
                sp => sp.GetRequiredService<CosmosTokenCredentialSingletonOptions>())
            .TryAddProviderSpecificServices(
                sm => sm.TryAddSingleton(
                    sp => new CosmosTokenCredentialSingletonOptions()));
    }

    public void Validate(IDbContextOptions options)
    {
        // Do nothing.
    }

    private sealed class ExtensionInfo : DbContextOptionsExtensionInfo
    {
        private int? _serviceProviderHash;

        public ExtensionInfo(IDbContextOptionsExtension extension)
            : base(extension)
        {
        }

        public override bool IsDatabaseProvider
            => true;

        public override string LogFragment => string.Empty;

        private new CosmosTokenCredentialOptionsExtension Extension
            => (CosmosTokenCredentialOptionsExtension)base.Extension;

        public override int GetServiceProviderHashCode()
        {
            if (_serviceProviderHash == null)
            {
#pragma warning disable SA1129 // Do not use default value type constructor
                var hashCode = new HashCode();
#pragma warning restore SA1129 // Do not use default value type constructor
                hashCode.Add(Extension.TokenCredential);

                _serviceProviderHash = hashCode.ToHashCode();
            }

            return _serviceProviderHash.Value;
        }

        public override bool ShouldUseSameServiceProvider(DbContextOptionsExtensionInfo other)
        {
            return other is ExtensionInfo otherInfo
                   && ReferenceEquals(Extension.TokenCredential, otherInfo.Extension.TokenCredential);
        }

        public override void PopulateDebugInfo(IDictionary<string, string> debugInfo)
        {
            if (debugInfo == null)
            {
                throw new ArgumentNullException(nameof(debugInfo));
            }

            debugInfo["Cosmos:" + nameof(Extension.TokenCredential)] =
                (Extension.TokenCredential?.GetHashCode() ?? 0L).ToString(CultureInfo.InvariantCulture);
        }
    }
}

[SuppressMessage("Usage", "EF1001:Internal EF Core API usage.", Justification = "EF Shim")]
[UsedImplicitly(ImplicitUseKindFlags.InstantiatedNoFixedConstructorSignature)]
public class CosmosTokenCredentialSingletonOptions : ISingletonOptions
{
    public virtual TokenCredential? TokenCredential { get; private set; }

    public void Initialize(IDbContextOptions options)
    {
        // ReSharper disable once UsePatternMatching
        var tokenCredentialOptions = options.FindExtension<CosmosTokenCredentialOptionsExtension>();
        if (tokenCredentialOptions != null)
        {
            TokenCredential = tokenCredentialOptions.TokenCredential;
        }
    }

    public void Validate(IDbContextOptions options)
    {
        // ReSharper disable once UsePatternMatching
#pragma warning disable CA1062 // Validate arguments of public methods
        var tokenCredentialOptions = options.FindExtension<CosmosTokenCredentialOptionsExtension>();
#pragma warning restore CA1062 // Validate arguments of public methods
        if (tokenCredentialOptions != null && TokenCredential != tokenCredentialOptions.TokenCredential)
        {
            throw new InvalidOperationException("Singleton options changed.");
        }
    }
}