Add uncached password provider support (#5399)

NinoFloris · web-flow · commit 69ebb3b8eb87 · 2023-11-17T18:50:42.000+01:00
Closes #5186
diff --git a/src/Npgsql/NpgsqlDataSource.cs b/src/Npgsql/NpgsqlDataSource.cs
@@ -3,7 +3,6 @@
 using System.Data.Common;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
-using System.Linq;
 using System.Net.Security;
 using System.Security.Cryptography.X509Certificates;
 using System.Threading;
@@ -44,6 +43,8 @@ public abstract class NpgsqlDataSource : DbDataSource
     internal RemoteCertificateValidationCallback? UserCertificateValidationCallback { get; }
     internal Action<X509CertificateCollection>? ClientCertificatesCallback { get; }
 
+    readonly Func<NpgsqlConnectionStringBuilder, string>? _passwordProvider;
+    readonly Func<NpgsqlConnectionStringBuilder, CancellationToken, ValueTask<string>>? _passwordProviderAsync;
     readonly Func<NpgsqlConnectionStringBuilder, CancellationToken, ValueTask<string>>? _periodicPasswordProvider;
     readonly TimeSpan _periodicPasswordSuccessRefreshInterval, _periodicPasswordFailureRefreshInterval;
 
@@ -52,7 +53,7 @@ public abstract class NpgsqlDataSource : DbDataSource
     internal Action<NpgsqlConnection>? ConnectionInitializer { get; }
     internal Func<NpgsqlConnection, Task>? ConnectionInitializerAsync { get; }
 
-    readonly Timer? _passwordProviderTimer;
+    readonly Timer? _periodicPasswordProviderTimer;
     readonly CancellationTokenSource? _timerPasswordProviderCancellationTokenSource;
     readonly Task _passwordRefreshTask = null!;
     string? _password;
@@ -101,6 +102,8 @@ internal NpgsqlDataSource(
                 IntegratedSecurityHandler,
                 UserCertificateValidationCallback,
                 ClientCertificatesCallback,
+                _passwordProvider,
+                _passwordProviderAsync,
                 _periodicPasswordProvider,
                 _periodicPasswordSuccessRefreshInterval,
                 _periodicPasswordFailureRefreshInterval,
@@ -112,6 +115,8 @@ internal NpgsqlDataSource(
             = dataSourceConfig;
         _connectionLogger = LoggingConfiguration.ConnectionLogger;
 
+        Debug.Assert(_passwordProvider is null || _passwordProviderAsync is not null);
+
         // TODO probably want this on the options so it can devirt unconditionally.
         _resolver = new TypeInfoResolverChain(resolverChain);
         _password = settings.Password;
@@ -123,7 +128,7 @@ internal NpgsqlDataSource(
             _timerPasswordProviderCancellationTokenSource = new();
 
             // Create the timer, but don't start it; the manual run below will will schedule the first refresh.
-            _passwordProviderTimer = new Timer(state => _ = RefreshPassword(), null, Timeout.InfiniteTimeSpan, Timeout.InfiniteTimeSpan);
+            _periodicPasswordProviderTimer = new Timer(state => _ = RefreshPassword(), null, Timeout.InfiniteTimeSpan, Timeout.InfiniteTimeSpan);
             // Trigger the first refresh attempt right now, outside the timer; this allows us to capture the Task so it can be observed
             // in GetPasswordAsync.
             _passwordRefreshTask = Task.Run(RefreshPassword);
@@ -293,28 +298,48 @@ public string Password
     {
         set
         {
-            if (_periodicPasswordProvider is not null)
+            if (_passwordProvider is not null || _periodicPasswordProvider is not null)
                 throw new NotSupportedException(NpgsqlStrings.CannotSetBothPasswordProviderAndPassword);
 
             _password = value;
         }
     }
 
-    internal async ValueTask<string?> GetPassword(bool async, CancellationToken cancellationToken = default)
+    internal ValueTask<string?> GetPassword(bool async, CancellationToken cancellationToken = default)
     {
+        if (_passwordProvider is not null)
+            return GetPassword(async, cancellationToken);
+
         // A periodic password provider is configured, but the first refresh hasn't completed yet (race condition).
-        // Wait until it completes.
         if (_password is null && _periodicPasswordProvider is not null)
+            return GetInitialPeriodicPassword(async);
+
+        return new(_password);
+
+        async ValueTask<string?> GetInitialPeriodicPassword(bool async)
         {
             if (async)
                 await _passwordRefreshTask.ConfigureAwait(false);
             else
                 _passwordRefreshTask.GetAwaiter().GetResult();
-
             Debug.Assert(_password is not null);
+
+            return _password;
         }
 
-        return _password;
+        async ValueTask<string?> GetPassword(bool async, CancellationToken cancellationToken)
+        {
+            try
+            {
+                return async ? await _passwordProviderAsync!(Settings, cancellationToken).ConfigureAwait(false) : _passwordProvider(Settings);
+            }
+            catch (Exception e)
+            {
+                _connectionLogger.LogError(e, "Password provider threw an exception");
+
+                throw new NpgsqlException("An exception was thrown from the password provider", e);
+            }
+        }
     }
 
     async Task RefreshPassword()
@@ -323,13 +348,13 @@ async Task RefreshPassword()
         {
             _password = await _periodicPasswordProvider!(Settings, _timerPasswordProviderCancellationTokenSource!.Token).ConfigureAwait(false);
 
-            _passwordProviderTimer!.Change(_periodicPasswordSuccessRefreshInterval, Timeout.InfiniteTimeSpan);
+            _periodicPasswordProviderTimer!.Change(_periodicPasswordSuccessRefreshInterval, Timeout.InfiniteTimeSpan);
         }
         catch (Exception e)
         {
             _connectionLogger.LogError(e, "Periodic password provider threw an exception");
 
-            _passwordProviderTimer!.Change(_periodicPasswordFailureRefreshInterval, Timeout.InfiniteTimeSpan);
+            _periodicPasswordProviderTimer!.Change(_periodicPasswordFailureRefreshInterval, Timeout.InfiniteTimeSpan);
 
             throw new NpgsqlException("An exception was thrown from the periodic password provider", e);
         }
@@ -448,7 +473,7 @@ protected virtual void DisposeBase()
             cancellationTokenSource.Dispose();
         }
 
-        _passwordProviderTimer?.Dispose();
+        _periodicPasswordProviderTimer?.Dispose();
         _setupMappingsSemaphore.Dispose();
         MetricsReporter.Dispose(); // TODO: This is probably too early, dispose only when all connections have been closed?
 
@@ -475,12 +500,12 @@ protected virtual async ValueTask DisposeAsyncBase()
             cancellationTokenSource.Dispose();
         }
 
-        if (_passwordProviderTimer is not null)
+        if (_periodicPasswordProviderTimer is not null)
         {
 #if NET5_0_OR_GREATER
-            await _passwordProviderTimer.DisposeAsync().ConfigureAwait(false);
+            await _periodicPasswordProviderTimer.DisposeAsync().ConfigureAwait(false);
 #else
-            _passwordProviderTimer.Dispose();
+            _periodicPasswordProviderTimer.Dispose();
 #endif
         }
 
diff --git a/src/Npgsql/NpgsqlDataSourceBuilder.cs b/src/Npgsql/NpgsqlDataSourceBuilder.cs
@@ -232,6 +232,30 @@ public NpgsqlDataSourceBuilder UsePeriodicPasswordProvider(
         return this;
     }
 
+    /// <summary>
+    /// Configures a password provider, which is called by the data source when opening connections.
+    /// </summary>
+    /// <param name="passwordProvider">
+    /// A callback that may be invoked during <see cref="NpgsqlConnection.Open()" /> which returns the password to be sent to PostgreSQL.
+    /// </param>
+    /// <param name="passwordProviderAsync">
+    /// A callback that may be invoked during <see cref="NpgsqlConnection.OpenAsync(CancellationToken)" /> which returns the password to be sent to PostgreSQL.
+    /// </param>
+    /// <returns>The same builder instance so that multiple calls can be chained.</returns>
+    /// <remarks>
+    /// <para>
+    /// The provided callback is invoked when opening connections. Therefore its important the callback internally depends on cached
+    /// data or returns quickly otherwise. Any unnecessary delay will affect connection opening time.
+    /// </para>
+    /// </remarks>
+    public NpgsqlDataSourceBuilder UsePasswordProvider(
+        Func<NpgsqlConnectionStringBuilder, string>? passwordProvider,
+        Func<NpgsqlConnectionStringBuilder, CancellationToken, ValueTask<string>>? passwordProviderAsync)
+    {
+        _internalBuilder.UsePasswordProvider(passwordProvider, passwordProviderAsync);
+        return this;
+    }
+
     #endregion Authentication
 
     #region Type mapping
diff --git a/src/Npgsql/NpgsqlDataSourceConfiguration.cs b/src/Npgsql/NpgsqlDataSourceConfiguration.cs
@@ -14,6 +14,8 @@ sealed record NpgsqlDataSourceConfiguration(string? Name,
     IntegratedSecurityHandler userCertificateValidationCallback,
     RemoteCertificateValidationCallback? UserCertificateValidationCallback,
     Action<X509CertificateCollection>? ClientCertificatesCallback,
+    Func<NpgsqlConnectionStringBuilder, string>? PasswordProvider,
+    Func<NpgsqlConnectionStringBuilder, CancellationToken, ValueTask<string>>? PasswordProviderAsync,
     Func<NpgsqlConnectionStringBuilder, CancellationToken, ValueTask<string>>? PeriodicPasswordProvider,
     TimeSpan PeriodicPasswordSuccessRefreshInterval,
     TimeSpan PeriodicPasswordFailureRefreshInterval,
diff --git a/src/Npgsql/NpgsqlSlimDataSourceBuilder.cs b/src/Npgsql/NpgsqlSlimDataSourceBuilder.cs
@@ -33,15 +33,18 @@ public sealed class NpgsqlSlimDataSourceBuilder : INpgsqlTypeMapper
 
     IntegratedSecurityHandler _integratedSecurityHandler = new();
 
+    Func<NpgsqlConnectionStringBuilder, string>? _passwordProvider;
+    Func<NpgsqlConnectionStringBuilder, CancellationToken, ValueTask<string>>? _passwordProviderAsync;
+
     Func<NpgsqlConnectionStringBuilder, CancellationToken, ValueTask<string>>? _periodicPasswordProvider;
     TimeSpan _periodicPasswordSuccessRefreshInterval, _periodicPasswordFailureRefreshInterval;
 
     PgTypeInfoResolverChainBuilder _resolverChainBuilder = new(); // mutable struct, don't make readonly.
 
     readonly UserTypeMapper _userTypeMapper;
 
-    Action<NpgsqlConnection>? _syncConnectionInitializer;
-    Func<NpgsqlConnection, Task>? _asyncConnectionInitializer;
+    Action<NpgsqlConnection>? _connectionInitializer;
+    Func<NpgsqlConnection, Task>? _connectionInitializerAsync;
 
     /// <summary>
     /// A connection string builder that can be used to configured the connection string on the builder.
@@ -239,6 +242,34 @@ public NpgsqlSlimDataSourceBuilder UsePeriodicPasswordProvider(
         return this;
     }
 
+    /// <summary>
+    /// Configures a password provider, which is called by the data source when opening connections.
+    /// </summary>
+    /// <param name="passwordProvider">
+    /// A callback that may be invoked during <see cref="NpgsqlConnection.Open()" /> which returns the password to be sent to PostgreSQL.
+    /// </param>
+    /// <param name="passwordProviderAsync">
+    /// A callback that may be invoked during <see cref="NpgsqlConnection.OpenAsync(CancellationToken)" /> which returns the password to be sent to PostgreSQL.
+    /// </param>
+    /// <returns>The same builder instance so that multiple calls can be chained.</returns>
+    /// <remarks>
+    /// <para>
+    /// The provided callback is invoked when opening connections. Therefore its important the callback internally depends on cached
+    /// data or returns quickly otherwise. Any unnecessary delay will affect connection opening time.
+    /// </para>
+    /// </remarks>
+    public NpgsqlSlimDataSourceBuilder UsePasswordProvider(
+        Func<NpgsqlConnectionStringBuilder, string>? passwordProvider,
+        Func<NpgsqlConnectionStringBuilder, CancellationToken, ValueTask<string>>? passwordProviderAsync)
+    {
+        if (passwordProvider is null != passwordProviderAsync is null)
+            throw new ArgumentException(NpgsqlStrings.SyncAndAsyncPasswordProvidersRequired);
+
+        _passwordProvider = passwordProvider;
+        _passwordProviderAsync = passwordProviderAsync;
+        return this;
+    }
+
     #endregion Authentication
 
     #region Type mapping
@@ -455,8 +486,8 @@ public NpgsqlSlimDataSourceBuilder UsePhysicalConnectionInitializer(
         if (connectionInitializer is null != connectionInitializerAsync is null)
             throw new ArgumentException(NpgsqlStrings.SyncAndAsyncConnectionInitializersRequired);
 
-        _syncConnectionInitializer = connectionInitializer;
-        _asyncConnectionInitializer = connectionInitializerAsync;
+        _connectionInitializer = connectionInitializer;
+        _connectionInitializerAsync = connectionInitializerAsync;
 
         return this;
     }
@@ -504,7 +535,12 @@ NpgsqlDataSourceConfiguration PrepareConfiguration()
             throw new InvalidOperationException(NpgsqlStrings.TransportSecurityDisabled);
         }
 
-        if (_periodicPasswordProvider is not null &&
+        if (_passwordProvider is not null && _periodicPasswordProvider is not null)
+        {
+            throw new NotSupportedException(NpgsqlStrings.CannotSetMultiplePasswordProviderKinds);
+        }
+
+        if ((_passwordProvider is not null || _periodicPasswordProvider is not null) &&
             (ConnectionStringBuilder.Password is not null || ConnectionStringBuilder.Passfile is not null))
         {
             throw new NotSupportedException(NpgsqlStrings.CannotSetBothPasswordProviderAndPassword);
@@ -519,14 +555,16 @@ _loggerFactory is null
             _integratedSecurityHandler,
             _userCertificateValidationCallback,
             _clientCertificatesCallback,
+            _passwordProvider,
+            _passwordProviderAsync,
             _periodicPasswordProvider,
             _periodicPasswordSuccessRefreshInterval,
             _periodicPasswordFailureRefreshInterval,
             _resolverChainBuilder.Build(ConfigureResolverChain),
             HackyEnumMappings(),
             DefaultNameTranslator,
-            _syncConnectionInitializer,
-            _asyncConnectionInitializer);
+            _connectionInitializer,
+            _connectionInitializerAsync);
 
         List<HackyEnumTypeMapping> HackyEnumMappings()
         {
diff --git a/src/Npgsql/Properties/NpgsqlStrings.Designer.cs b/src/Npgsql/Properties/NpgsqlStrings.Designer.cs
diff --git a/src/Npgsql/Properties/NpgsqlStrings.resx b/src/Npgsql/Properties/NpgsqlStrings.resx
@@ -42,6 +42,12 @@
     <data name="CannotSetBothPasswordProviderAndPassword" xml:space="preserve">
         <value>When registering a password provider, a password or password file may not be set.</value>
     </data>
+    <data name="CannotSetMultiplePasswordProviderKinds" xml:space="preserve">
+        <value>Multiple kinds of password providers were found, only one kind may be configured per DbDataSource.</value>
+    </data>
+    <data name="SyncAndAsyncPasswordProvidersRequired" xml:space="preserve">
+        <value>Both sync and async password providers must be provided.</value>
+    </data>
     <data name="PasswordProviderMissing" xml:space="preserve">
         <value>The right type of password provider (sync or async) was not found.</value>
     </data>
diff --git a/src/Npgsql/PublicAPI.Unshipped.txt b/src/Npgsql/PublicAPI.Unshipped.txt
@@ -15,6 +15,7 @@ Npgsql.NpgsqlDataSourceBuilder.MapEnum(System.Type! clrType, string? pgName = nu
 Npgsql.NpgsqlDataSourceBuilder.Name.get -> string?
 Npgsql.NpgsqlDataSourceBuilder.Name.set -> void
 Npgsql.NpgsqlDataSourceBuilder.UnmapEnum(System.Type! clrType, string? pgName = null, Npgsql.INpgsqlNameTranslator? nameTranslator = null) -> bool
+Npgsql.NpgsqlDataSourceBuilder.UsePasswordProvider(System.Func<Npgsql.NpgsqlConnectionStringBuilder!, string!>? passwordProvider, System.Func<Npgsql.NpgsqlConnectionStringBuilder!, System.Threading.CancellationToken, System.Threading.Tasks.ValueTask<string!>>? passwordProviderAsync) -> Npgsql.NpgsqlDataSourceBuilder!
 Npgsql.NpgsqlDataSourceBuilder.UseRootCertificate(System.Security.Cryptography.X509Certificates.X509Certificate2? rootCertificate) -> Npgsql.NpgsqlDataSourceBuilder!
 Npgsql.NpgsqlDataSourceBuilder.UseRootCertificateCallback(System.Func<System.Security.Cryptography.X509Certificates.X509Certificate2!>? rootCertificateCallback) -> Npgsql.NpgsqlDataSourceBuilder!
 Npgsql.NpgsqlSlimDataSourceBuilder
@@ -50,6 +51,7 @@ Npgsql.NpgsqlSlimDataSourceBuilder.UseClientCertificate(System.Security.Cryptogr
 Npgsql.NpgsqlSlimDataSourceBuilder.UseClientCertificates(System.Security.Cryptography.X509Certificates.X509CertificateCollection? clientCertificates) -> Npgsql.NpgsqlSlimDataSourceBuilder!
 Npgsql.NpgsqlSlimDataSourceBuilder.UseClientCertificatesCallback(System.Action<System.Security.Cryptography.X509Certificates.X509CertificateCollection!>? clientCertificatesCallback) -> Npgsql.NpgsqlSlimDataSourceBuilder!
 Npgsql.NpgsqlSlimDataSourceBuilder.UseLoggerFactory(Microsoft.Extensions.Logging.ILoggerFactory? loggerFactory) -> Npgsql.NpgsqlSlimDataSourceBuilder!
+Npgsql.NpgsqlSlimDataSourceBuilder.UsePasswordProvider(System.Func<Npgsql.NpgsqlConnectionStringBuilder!, string!>? passwordProvider, System.Func<Npgsql.NpgsqlConnectionStringBuilder!, System.Threading.CancellationToken, System.Threading.Tasks.ValueTask<string!>>? passwordProviderAsync) -> Npgsql.NpgsqlSlimDataSourceBuilder!
 Npgsql.NpgsqlSlimDataSourceBuilder.UsePeriodicPasswordProvider(System.Func<Npgsql.NpgsqlConnectionStringBuilder!, System.Threading.CancellationToken, System.Threading.Tasks.ValueTask<string!>>? passwordProvider, System.TimeSpan successRefreshInterval, System.TimeSpan failureRefreshInterval) -> Npgsql.NpgsqlSlimDataSourceBuilder!
 Npgsql.NpgsqlSlimDataSourceBuilder.UsePhysicalConnectionInitializer(System.Action<Npgsql.NpgsqlConnection!>? connectionInitializer, System.Func<Npgsql.NpgsqlConnection!, System.Threading.Tasks.Task!>? connectionInitializerAsync) -> Npgsql.NpgsqlSlimDataSourceBuilder!
 Npgsql.NpgsqlSlimDataSourceBuilder.UseRootCertificate(System.Security.Cryptography.X509Certificates.X509Certificate2? rootCertificate) -> Npgsql.NpgsqlSlimDataSourceBuilder!
diff --git a/test/Npgsql.Tests/AuthenticationTests.cs b/test/Npgsql.Tests/AuthenticationTests.cs
