Valid credential threshold (#652)

Author: adam-fowler

Sources/SotoCore/Credential/RotatingCredentialProvider.swift

@@ -24,17 +24,46 @@ import struct Foundation.TimeInterval
 ///
 /// Used for wrapping another credential provider whose `getCredential` method returns an `ExpiringCredential`.
 /// If no credential is available, or the current credentials are going to expire in the near future  the wrapped credential provider
-/// `getCredential` is called again. If current credentials have not expired they are returned otherwise we wait on new
+/// `getCredential` is called again. If current credentials have not expired (within a threshold) they are returned otherwise we wait on new
 /// credentials being provided.
 public final class RotatingCredentialProvider: CredentialProvider {
     let expiringCredential: ExpiringValue<Credential>
-
+    let validCredentialThreshold: TimeInterval
     public let provider: CredentialProvider
 
-    public init(context: CredentialProviderFactory.Context, provider: CredentialProvider, remainingTokenLifetimeForUse: TimeInterval? = nil) {
+    ///  Initialize RotatingCredentialProvider
+    /// - Parameters:
+    ///   - context: Context used to create this credential provider
+    ///   - provider: Credential provider to request credentials from
+    ///   - remainingTokenLifetimeForUse: How near to expiration, before we request new credentials
+    public init(
+        context: CredentialProviderFactory.Context,
+        provider: CredentialProvider,
+        remainingTokenLifetimeForUse: TimeInterval? = nil
+    ) {
+        self.provider = provider
+        self.validCredentialThreshold = 15
+        self.expiringCredential = .init(threshold: remainingTokenLifetimeForUse ?? 165) {
+            try await Self.getCredentialAndExpiration(provider: provider, validCredentialThreshold: 15, logger: context.logger)
+        }
+    }
+
+    ///  Initialize RotatingCredentialProvider
+    /// - Parameters:
+    ///   - context: Context used to create this credential provider
+    ///   - provider: Credential provider to request credentials from
+    ///   - remainingTokenLifetimeForUse: How near to expiration, before we request new credentials
+    ///   - validCredentialThreshold: How near to expiration do we return the current credentials
+    public init(
+        context: CredentialProviderFactory.Context,
+        provider: CredentialProvider,
+        remainingTokenLifetimeForUse: TimeInterval? = nil,
+        validCredentialThreshold: TimeInterval
+    ) {
         self.provider = provider
+        self.validCredentialThreshold = validCredentialThreshold
         self.expiringCredential = .init(threshold: remainingTokenLifetimeForUse ?? 3 * 60) {
-            try await Self.getCredentialAndExpiration(provider: provider, logger: context.logger)
+            try await Self.getCredentialAndExpiration(provider: provider, validCredentialThreshold: validCredentialThreshold, logger: context.logger)
         }
     }
 
@@ -51,17 +80,34 @@ public final class RotatingCredentialProvider: CredentialProvider {
 
     public func getCredential(logger: Logger) async throws -> Credential {
         try await self.expiringCredential.getValue {
-            try await Self.getCredentialAndExpiration(provider: self.provider, logger: logger)
+            try await Self.getCredentialAndExpiration(
+                provider: self.provider,
+                validCredentialThreshold: self.validCredentialThreshold,
+                logger: logger
+            )
         }
     }
 
-    static func getCredentialAndExpiration(provider: CredentialProvider, logger: Logger) async throws -> (Credential, Date) {
-        logger.debug("Refeshing AWS credentials", metadata: ["aws-credential-provider": .string("\(self)(\(provider.description))")])
+    static func getCredentialAndExpiration(
+        provider: CredentialProvider,
+        validCredentialThreshold: TimeInterval,
+        logger: Logger
+    ) async throws -> (Credential, Date) {
+        logger.debug(
+            "Refeshing AWS credentials",
+            metadata: ["aws-credential-provider": .string("\(self)(\(provider.description))")]
+        )
         try Task.checkCancellation()
         let credential = try await provider.getCredential(logger: logger)
-        logger.debug("AWS credentials ready", metadata: ["aws-credential-provider": .string("\(self)(\(provider.description))")])
+        logger.debug(
+            "AWS credentials ready",
+            metadata: ["aws-credential-provider": .string("\(self)(\(provider.description))")]
+        )
         if let expiringCredential = credential as? ExpiringCredential {
-            return (expiringCredential, expiringCredential.expiration)
+            return (
+                expiringCredential,
+                expiringCredential.expiration.addingTimeInterval(-validCredentialThreshold)
+            )
         } else {
             return (credential, Date.distantFuture)
         }

Tests/SotoCoreTests/Credential/RotatingCredentialProviderTests.swift

@@ -104,6 +104,56 @@ class RotatingCredentialProviderTests: XCTestCase {
         XCTAssertEqual(count.load(ordering: .sequentiallyConsistent), 1)
     }
 
+    func testGetCredentialAndGetNewOnesAsAboutToExpire() async throws {
+        let group = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        defer { XCTAssertNoThrow(try group.syncShutdownGracefully()) }
+        let httpClient = HTTPClient(eventLoopGroupProvider: .shared(group))
+        defer { XCTAssertNoThrow(try httpClient.syncShutdown()) }
+
+        let creds = [
+            TestExpiringCredential(
+                accessKeyId: "abc123",
+                secretAccessKey: "abc123",
+                sessionToken: "abc123",
+                expiration: Date(timeIntervalSinceNow: 10)
+            ),
+            TestExpiringCredential(
+                accessKeyId: "def456",
+                secretAccessKey: "def456",
+                sessionToken: "def456",
+                expiration: Date(timeIntervalSinceNow: 10)
+            ),
+        ]
+
+        let count = ManagedAtomic(0)
+        let client = RotatingCredentialTestClient {
+            let cred = creds[count.load(ordering: .sequentiallyConsistent)]
+            count.wrappingIncrement(ordering: .sequentiallyConsistent)
+            return cred
+        }
+        let context = CredentialProviderFactory.Context(httpClient: httpClient, logger: Logger(label: "soto"), options: .init())
+        let provider = RotatingCredentialProvider(context: context, provider: client)
+
+        // get credentials for first time
+        var returned = try await provider.getCredential(logger: Logger(label: "soto"))
+
+        XCTAssertEqual(returned.accessKeyId, creds[0].accessKeyId)
+        XCTAssertEqual(returned.secretAccessKey, creds[0].secretAccessKey)
+        XCTAssertEqual(returned.sessionToken, creds[0].sessionToken)
+        XCTAssertEqual((returned as? TestExpiringCredential)?.expiration, creds[0].expiration)
+
+        // get credentials a second time, callback must not be hit
+        returned = try await provider.getCredential(logger: Logger(label: "soto"))
+
+        XCTAssertEqual(returned.accessKeyId, creds[1].accessKeyId)
+        XCTAssertEqual(returned.secretAccessKey, creds[1].secretAccessKey)
+        XCTAssertEqual(returned.sessionToken, creds[1].sessionToken)
+        XCTAssertEqual((returned as? TestExpiringCredential)?.expiration, creds[1].expiration)
+
+        // ensure callback was hit twice as we
+        XCTAssertEqual(count.load(ordering: .sequentiallyConsistent), 2)
+    }
+
     func testGetCredentialHighlyConcurrent() async throws {
         let group = MultiThreadedEventLoopGroup(numberOfThreads: System.coreCount)
         defer { XCTAssertNoThrow(try group.syncShutdownGracefully()) }