ProviderHelper.cs

public sealed class ProviderHelper
    {
        private static readonly Lazy<ProviderHelper> Lazy = new Lazy<ProviderHelper>(() => new ProviderHelper());
        private static readonly Dictionary<string, Provider> Providers = new Dictionary<string, Provider>();

        private static readonly LocalDataStoreSlot LocalDataStoreSlot =
            Thread.GetNamedDataSlot(nameof(Provider).ToLower());

        private Input<string>? _accessKey;

        private bool _initialized;
        private Input<string>? _region;
        private Input<string>? _secretKey;

        private ProviderHelper()
        {
        }

        public static ProviderHelper Instance => Lazy.Value;

        public void SetStack(Stack stack)
        {
            var config = new Config("providerHelper");
            var accessKey = config.RequireSecret("awsAccessKey");
            var secretKey = config.RequireSecret("awsSecretKey");
            var region = config.Require("awsRegion");
            var defaultAssumeRoleArn = config.Get("awsDefaultAssumeRoleArn");

            SetStack(stack, accessKey, secretKey, region, defaultAssumeRoleArn);
        }

        private void SetStack(Stack stack, Input<string> accessKey, Input<string> secretKey, string region,
            string? defaultAssumeRoleArn)
        {
            _accessKey = accessKey;
            _secretKey = secretKey;
            _region = region;

            SetCurrentProvider(GetOrCreateProvider(defaultAssumeRoleArn));

            stack.SetPrivate("_transformations",
                ((ImmutableArray<ResourceTransformation>) stack.GetPrivate("_transformations")).Add(args =>
                {
                    args.Options.Provider = GetCurrentProvider();
                    return new ResourceTransformationResult(args.Args, args.Options);
                }));

            _initialized = true;
        }

        public ScopedProvider AssumeRole(Role role)
        {
            return AssumeRole(role.Value);
        }

        public ScopedProvider AssumeRole(string roleArn)
        {
            if (string.IsNullOrEmpty(roleArn)) throw new ArgumentNullException(nameof(roleArn));

            IsInitialized();
            return new ScopedProvider(GetOrCreateProvider(roleArn));
        }

        public Provider GetAssumeRole(Role role)
        {
            return GetAssumeRole(role.Value);
        }

        public Provider GetAssumeRole(string roleArn)
        {
            IsInitialized();
            return GetOrCreateProvider(roleArn);
        }

        public Provider GetCurrentProvider()
        {
            IsInitialized();
            return (Provider) Thread.GetData(LocalDataStoreSlot)!;
        }

        internal void SetCurrentProvider(Provider provider)
        {
            Thread.SetData(LocalDataStoreSlot, provider);
        }

        internal void IsInitialized()
        {
            if (!_initialized) throw new InvalidOperationException("Call SetStack(myStack) to initialize.");
        }

        internal Provider GetOrCreateProvider(string? roleArn = null)
        {
            var providerName = GetProviderName(roleArn);
            if (Providers.TryGetValue(providerName, out var provider)) return provider;

            provider = CreateProvider(providerName, roleArn);
            Providers.Add(providerName, provider);
            return provider;
        }

        internal Provider CreateProvider(
            string providerName,
            string? roleArn = null)
        {
            var providerArgs = new ProviderArgs
            {
                AccessKey = _accessKey,
                SecretKey = _secretKey,
                Region = _region
            };

            if (string.IsNullOrEmpty(roleArn)) return new Provider(providerName, providerArgs);

            providerArgs.AssumeRole = new ProviderAssumeRoleArgs
            {
                RoleArn = roleArn
            };

            return new Provider(providerName, providerArgs);
        }

        internal string GetProviderName(string? roleArn = null)
        {
            if (string.IsNullOrEmpty(roleArn)) return "default";

            var arn = Arn.Parse(roleArn);
            return $"{arn.AccountId}-{arn.Resource.Replace('/', '-')}";
        }
    }

    public class ScopedProvider : IDisposable
    {
        private readonly Provider _previousProvider;

        public ScopedProvider(Provider provider)
        {
            _previousProvider = ProviderHelper.Instance.GetCurrentProvider();
            SetProvider(provider);
        }

        public void Dispose()
        {
            SetProvider(_previousProvider);
        }

        ~ScopedProvider()
        {
            Dispose();
        }

        private void SetProvider(Provider provider)
        {
            ProviderHelper.Instance.SetCurrentProvider(provider);
        }
    }