CSHARP-4273: Cache AWS Credentials Where Possible.

Author: DmitryLukyanov

evergreen/evergreen.yml

@@ -514,6 +514,7 @@ functions:
           cat <<'EOF' > "${PROJECT_DIRECTORY}/prepare_mongodb_aws.sh"
             MONGODB_URI="mongodb://localhost"
           EOF
+          export AWS_EC2_ENABLED=true
           PROJECT_DIRECTORY=${PROJECT_DIRECTORY} evergreen/run-mongodb-aws-test.sh
 
   run-aws-auth-test-with-aws-credentials-as-environment-variables:
@@ -987,7 +988,7 @@ tasks:
         - func: run-aws-auth-test-with-aws-credentials-as-environment-variables
         - func: run-aws-auth-test-with-aws-credentials-and-session-token-as-environment-variables
         - func: run-aws-auth-test-with-aws-EC2-credentials
-        # ECS test is skipped untill testing on Linux becomes possible
+        # ECS test is skipped until testing on Linux becomes possible
 
     - name: stable-api-tests-net472
       commands:

src/MongoDB.Driver.Core/Core/Authentication/External/AwsAuthenticationCredentialsProvider.cs

@@ -25,20 +25,27 @@ namespace MongoDB.Driver.Core.Authentication.External
 {
     internal class AwsCredentials : IExternalCredentials
     {
+        // credentials are considered expired when: Expiration - now > 5 mins
+        private static readonly TimeSpan __overlapWhenExpired = TimeSpan.FromMinutes(5);
+
         private readonly string _accessKeyId;
+        private readonly DateTime? _expiration;
         private readonly SecureString _secretAccessKey;
         private readonly string _sessionToken;
 
-        public AwsCredentials(string accessKeyId, SecureString secretAccessKey, string sessionToken)
+        public AwsCredentials(string accessKeyId, SecureString secretAccessKey, string sessionToken, string expiration)
         {
             _accessKeyId = Ensure.IsNotNull(accessKeyId, nameof(accessKeyId));
+            _expiration = expiration != null ? DateTime.Parse(expiration) : null;
             _secretAccessKey = Ensure.IsNotNull(secretAccessKey, nameof(secretAccessKey));
             _sessionToken = sessionToken; // can be null
         }
 
         public string AccessKeyId => _accessKeyId;
+        public DateTime? Expiration => _expiration;
         public SecureString SecretAccessKey => _secretAccessKey;
         public string SessionToken => _sessionToken;
+        public bool IsExpired => _expiration.HasValue ? (_expiration.Value - DateTime.UtcNow) < __overlapWhenExpired : false;
 
         public BsonDocument GetKmsCredentials()
             => new BsonDocument
@@ -93,7 +100,7 @@ private AwsCredentials CreateAwsCredentialsFromEnvironmentVariables()
                 throw new InvalidOperationException($"When using AWS authentication if a session token is provided via environment variables then an access key ID and a secret access key must be provided also.");
             }
 
-            return new AwsCredentials(accessKeyId, SecureStringHelper.ToSecureString(secretAccessKey), sessionToken);
+            return new AwsCredentials(accessKeyId, SecureStringHelper.ToSecureString(secretAccessKey), sessionToken, expiration: null);
         }
 
         private async Task<AwsCredentials> CreateAwsCredentialsFromEcsResponseAsync(CancellationToken cancellationToken)
@@ -116,12 +123,23 @@ private async Task<AwsCredentials> CreateAwsCredentialsFromEc2ResponseAsync(Canc
 
         private AwsCredentials CreateAwsCreadentialsFromAwsResponse(string awsResponse)
         {
+            // Response template:
+            //{
+            //    "Code": "Success",
+            //    "LastUpdated": "..",
+            //    "Type": "AWS-HMAC",
+            //    "AccessKeyId": "..",
+            //    "SecretAccessKey": "..",
+            //    "Token": "",
+            //    "Expiration": "YYYY-mm-ddThh:mm:ssZ"
+            //}
             var parsedResponse = BsonDocument.Parse(awsResponse);
             var accessKeyId = parsedResponse.GetValue("AccessKeyId", null)?.AsString;
             var secretAccessKey = parsedResponse.GetValue("SecretAccessKey", null)?.AsString;
             var sessionToken = parsedResponse.GetValue("Token", null)?.AsString;
+            var expiration = parsedResponse.GetValue("Expiration", null)?.AsString;
 
-            return new AwsCredentials(accessKeyId, SecureStringHelper.ToSecureString(secretAccessKey), sessionToken);
+            return new AwsCredentials(accessKeyId, SecureStringHelper.ToSecureString(secretAccessKey), sessionToken, expiration);
         }
 
         // nested types

src/MongoDB.Driver.Core/Core/Authentication/External/CacheableCredentialsProvider.cs

@@ -0,0 +1,93 @@
+/* Copyright 2010-present MongoDB Inc.
+*
+* Licensed under the Apache License, Version 2.0 (the "License");
+* you may not use this file except in compliance with the License.
+* You may obtain a copy of the License at
+*
+* http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+
+using System.Threading;
+using System.Threading.Tasks;
+using MongoDB.Driver.Core.Misc;
+
+namespace MongoDB.Driver.Core.Authentication.External
+{
+    internal interface ICredentialsCache<TCredentials> where TCredentials : IExternalCredentials
+    {
+        TCredentials Credentials { get; }
+        void Clear();
+    }
+
+    internal class CacheableCredentialsProvider<TCredentials> : IExternalAuthenticationCredentialsProvider<TCredentials>, ICredentialsCache<TCredentials>
+        where TCredentials : IExternalCredentials
+    {
+        private TCredentials _cachedCredentials;
+        private readonly IExternalAuthenticationCredentialsProvider<TCredentials> _provider;
+
+        public CacheableCredentialsProvider(IExternalAuthenticationCredentialsProvider<TCredentials> provider)
+        {
+            _provider = Ensure.IsNotNull(provider, nameof(provider));
+        }
+
+        public TCredentials Credentials => _cachedCredentials;
+
+        public TCredentials CreateCredentialsFromExternalSource(CancellationToken cancellationToken = default)
+        {
+            var cachedCredentials = _cachedCredentials;
+            if (IsValidCache(cachedCredentials))
+            {
+                return cachedCredentials;
+            }
+            else
+            {
+                Clear();
+                try
+                {
+                    cachedCredentials = _provider.CreateCredentialsFromExternalSource(cancellationToken);
+                    _cachedCredentials = cachedCredentials;
+                    return cachedCredentials;
+                }
+                catch
+                {
+                    Clear();
+                    throw;
+                }
+            }
+        }
+
+        public async Task<TCredentials> CreateCredentialsFromExternalSourceAsync(CancellationToken cancellationToken = default)
+        {
+            var cachedCredentials = _cachedCredentials;
+            if (IsValidCache(cachedCredentials))
+            {
+                return cachedCredentials;
+            }
+            else
+            {
+                Clear();
+                try
+                {
+                    cachedCredentials = await _provider.CreateCredentialsFromExternalSourceAsync(cancellationToken).ConfigureAwait(false);
+                    _cachedCredentials = cachedCredentials;
+                    return cachedCredentials;
+                }
+                catch
+                {
+                    Clear();
+                    throw;
+                }
+            }
+        }
+
+        // private methods
+        private bool IsValidCache(TCredentials credentials) => credentials != null && !credentials.IsExpired;
+        public void Clear() => _cachedCredentials = default;
+    }
+}

src/MongoDB.Driver.Core/Core/Authentication/External/ExternalCredentialsAuthenticators.cs

@@ -14,7 +14,6 @@
 */
 
 using System;
-using System.Net.Http;
 using MongoDB.Driver.Core.Misc;
 
 namespace MongoDB.Driver.Core.Authentication.External
@@ -38,7 +37,7 @@ internal ExternalCredentialsAuthenticators() : this(new HttpClientHelper())
         internal ExternalCredentialsAuthenticators(HttpClientHelper httpClientHelper)
         {
             _httpClientHelper = Ensure.IsNotNull(httpClientHelper, nameof(httpClientHelper));
-            _awsExternalAuthenticationCredentialsProvider = new Lazy<IExternalAuthenticationCredentialsProvider<AwsCredentials>>(() => new AwsAuthenticationCredentialsProvider(_httpClientHelper), isThreadSafe: true);
+            _awsExternalAuthenticationCredentialsProvider = new Lazy<IExternalAuthenticationCredentialsProvider<AwsCredentials>>(() => new CacheableCredentialsProvider<AwsCredentials>(new AwsAuthenticationCredentialsProvider(_httpClientHelper)), isThreadSafe: true);
             _gcpExternalAuthenticationCredentialsProvider = new Lazy<IExternalAuthenticationCredentialsProvider<GcpCredentials>>(() => new GcpAuthenticationCredentialsProvider(_httpClientHelper), isThreadSafe: true);
         }
 

src/MongoDB.Driver.Core/Core/Authentication/External/GcpAuthenticationCredentialsProvider.cs

@@ -30,6 +30,10 @@ internal class GcpCredentials : IExternalCredentials
 
         public string AccessToken => _accessToken;
 
+        public DateTime? Expiration => null;
+
+        public bool IsExpired => false;
+
         public BsonDocument GetKmsCredentials() => new BsonDocument("accessToken", _accessToken);
     }
 

src/MongoDB.Driver.Core/Core/Authentication/External/IExternalCredentials.cs

@@ -13,12 +13,15 @@
 * limitations under the License.
 */
 
+using System;
 using MongoDB.Bson;
 
 namespace MongoDB.Driver.Core.Authentication.External
 {
     internal interface IExternalCredentials
     {
+        DateTime? Expiration { get; }
+        bool IsExpired { get; }
         BsonDocument GetKmsCredentials();
     }
 }

src/MongoDB.Driver.Core/Core/Authentication/MongoAWSAuthenticator.cs

@@ -17,6 +17,8 @@
 using System.Collections.Generic;
 using System.Linq;
 using System.Security;
+using System.Threading;
+using System.Threading.Tasks;
 using MongoDB.Bson;
 using MongoDB.Bson.Serialization;
 using MongoDB.Driver.Core.Authentication.External;
@@ -51,26 +53,28 @@ private static MongoAWSMechanism CreateMechanism(
             UsernamePasswordCredential credential,
             IEnumerable<KeyValuePair<string, string>> properties,
             IRandomByteGenerator randomByteGenerator,
+            IExternalAuthenticationCredentialsProvider<AwsCredentials> externalAuthenticationCredentialsProvider,
             IClock clock)
         {
             if (credential.Source != "$external")
             {
                 throw new ArgumentException("MONGODB-AWS authentication may only use the $external source.", nameof(credential));
             }
 
-            return CreateMechanism(credential.Username, credential.Password, properties, randomByteGenerator, clock);
+            return CreateMechanism(credential.Username, credential.Password, properties, randomByteGenerator, externalAuthenticationCredentialsProvider, clock);
         }
 
         private static MongoAWSMechanism CreateMechanism(
             string username,
             SecureString password,
             IEnumerable<KeyValuePair<string, string>> properties,
             IRandomByteGenerator randomByteGenerator,
+            IExternalAuthenticationCredentialsProvider<AwsCredentials> externalAuthenticationCredentialsProvider,
             IClock clock)
         {
             var awsCredentials =
                 CreateAwsCredentialsFromMongoCredentials(username, password, properties) ??
-                ExternalCredentialsAuthenticators.Instance.Aws.CreateCredentialsFromExternalSource();
+                externalAuthenticationCredentialsProvider.CreateCredentialsFromExternalSource();
 
             return new MongoAWSMechanism(awsCredentials, randomByteGenerator, clock);
         }
@@ -97,7 +101,7 @@ private static AwsCredentials CreateAwsCredentialsFromMongoCredentials(string us
                 throw new InvalidOperationException("When using MONGODB-AWS authentication if a session token is provided via settings then a username and password must be provided also.");
             }
 
-            return new AwsCredentials(accessKeyId: username, secretAccessKey: password, sessionToken);
+            return new AwsCredentials(accessKeyId: username, secretAccessKey: password, sessionToken, expiration: null);
         }
 
         private static string ExtractSessionTokenFromMechanismProperties(IEnumerable<KeyValuePair<string, string>> properties)
@@ -131,18 +135,9 @@ private static void ValidateMechanismProperties(IEnumerable<KeyValuePair<string,
         }
         #endregion
 
-        // constructors
-        /// <summary>
-        /// Initializes a new instance of the <see cref="MongoAWSAuthenticator"/> class.
-        /// </summary>
-        /// <param name="credential">The credentials.</param>
-        /// <param name="properties">The properties.</param>
-        [Obsolete("Use the newest overload instead.")]
-        public MongoAWSAuthenticator(UsernamePasswordCredential credential, IEnumerable<KeyValuePair<string, string>> properties)
-            : this(credential, properties, serverApi: null)
-        {
-        }
+        private readonly ICredentialsCache<AwsCredentials> _credentialsCache;
 
+        // constructors
         /// <summary>
         /// Initializes a new instance of the <see cref="MongoAWSAuthenticator"/> class.
         /// </summary>
@@ -153,18 +148,13 @@ public MongoAWSAuthenticator(
             UsernamePasswordCredential credential,
             IEnumerable<KeyValuePair<string, string>> properties,
             ServerApi serverApi)
-            : this(credential, properties, new DefaultRandomByteGenerator(), SystemClock.Instance, serverApi)
-        {
-        }
-
-        /// <summary>
-        /// Initializes a new instance of the <see cref="MongoAWSAuthenticator"/> class.
-        /// </summary>
-        /// <param name="username">The username.</param>
-        /// <param name="properties">The properties.</param>
-        [Obsolete("Use the newest overload instead.")]
-        public MongoAWSAuthenticator(string username, IEnumerable<KeyValuePair<string, string>> properties)
-            : this(username, properties, serverApi: null)
+            : this(
+                  credential,
+                  properties,
+                  new DefaultRandomByteGenerator(),
+                  ExternalCredentialsAuthenticators.Instance.Aws,
+                  SystemClock.Instance,
+                  serverApi)
         {
         }
 
@@ -178,28 +168,38 @@ public MongoAWSAuthenticator(
             string username,
             IEnumerable<KeyValuePair<string, string>> properties,
             ServerApi serverApi)
-            : this(username, properties, new DefaultRandomByteGenerator(), SystemClock.Instance, serverApi)
+            : this(
+                  username,
+                  properties,
+                  new DefaultRandomByteGenerator(),
+                  ExternalCredentialsAuthenticators.Instance.Aws,
+                  SystemClock.Instance,
+                  serverApi)
         {
         }
 
         internal MongoAWSAuthenticator(
             UsernamePasswordCredential credential,
             IEnumerable<KeyValuePair<string, string>> properties,
             IRandomByteGenerator randomByteGenerator,
+            IExternalAuthenticationCredentialsProvider<AwsCredentials> externalAuthenticationCredentialsProvider,
             IClock clock,
             ServerApi serverApi)
-            : base(CreateMechanism(credential, properties, randomByteGenerator, clock), serverApi)
+            : base(CreateMechanism(credential, properties, randomByteGenerator, externalAuthenticationCredentialsProvider, clock), serverApi)
         {
+            _credentialsCache = externalAuthenticationCredentialsProvider as ICredentialsCache<AwsCredentials>; // can be null
         }
 
         internal MongoAWSAuthenticator(
             string username,
             IEnumerable<KeyValuePair<string, string>> properties,
             IRandomByteGenerator randomByteGenerator,
+            IExternalAuthenticationCredentialsProvider<AwsCredentials> externalAuthenticationCredentialsProvider,
             IClock clock,
             ServerApi serverApi)
-            : base(CreateMechanism(username, null, properties, randomByteGenerator, clock), serverApi)
+            : base(CreateMechanism(username, null, properties, randomByteGenerator, externalAuthenticationCredentialsProvider, clock), serverApi)
         {
+            _credentialsCache = externalAuthenticationCredentialsProvider as ICredentialsCache<AwsCredentials>; // can be null
         }
 
         /// <inheritdoc/>
@@ -208,6 +208,34 @@ public override string DatabaseName
             get { return "$external"; }
         }
 
+        /// <inheritdoc/>
+        public override void Authenticate(IConnection connection, ConnectionDescription description, CancellationToken cancellationToken)
+        {
+            try
+            {
+                base.Authenticate(connection, description, cancellationToken);
+            }
+            catch
+            {
+                _credentialsCache?.Clear();
+                throw;
+            }
+        }
+
+        /// <inheritdoc/>
+        public override async Task AuthenticateAsync(IConnection connection, ConnectionDescription description, CancellationToken cancellationToken)
+        {
+            try
+            {
+                await base.AuthenticateAsync(connection, description, cancellationToken).ConfigureAwait(false);
+            }
+            catch
+            {
+                _credentialsCache?.Clear();
+                throw;
+            }
+        }
+
         // nested classes
         private class MongoAWSMechanism : ISaslMechanism
         {