Skip to content

Commit c56e350

Browse files
dayaffedayaffe
committed
some fixing up
1 parent e960c5f commit c56e350

File tree

4 files changed

+194
-33
lines changed

4 files changed

+194
-33
lines changed

‎Sources/Core/AWSSDKIdentity/Sources/AWSSDKIdentity/AWSCredentialIdentityResolvers/CognitoAWSCredentialIdentityResolver.swift

Lines changed: 72 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -10,38 +10,64 @@ import struct Smithy.Attributes
1010
import ClientRuntime
1111
import class Foundation.ProcessInfo
1212
import enum Smithy.ClientError
13-
import class Foundation.NSLock
13+
1414
import struct Foundation.Date
15+
@_spi(FileBasedConfig) import AWSSDKCommon
1516

1617
/// A credential identity resolver that resolves credentials using AWS Cognito Identity.
17-
public struct CognitoAWSCredentialIdentityResolver: AWSCredentialIdentityResolver {
18+
public actor CognitoAWSCredentialIdentityResolver: AWSCredentialIdentityResolver {
1819
private let identityPoolId: String?
1920
private let identityId: String?
2021
private let accountId: String?
2122
private var logins: [String: String]?
2223
private let customRoleArn: String?
2324
private let cognitoPoolRegion: String
25+
private let configFilePath: String?
26+
private let credentialsFilePath: String?
27+
private let profileName: String?
2428

2529
private var cachedIdentityId: String?
2630
private var cachedCredentials: AWSCredentialIdentity?
2731
private var cachedLogins: [String: String]?
28-
private let refreshLock = NSLock()
32+
// Actor provides thread safety, no need for manual locking
2933

3034
public init(
3135
identityPoolId: String? = nil,
3236
identityId: String? = nil,
3337
accountId: String? = nil,
3438
logins: [String: String]? = nil,
3539
customRoleArn: String? = nil,
36-
cognitoPoolRegion: String? = nil
40+
cognitoPoolRegion: String? = nil,
41+
configFilePath: String? = nil,
42+
credentialsFilePath: String? = nil,
43+
profileName: String? = nil
3744
) throws {
3845
let env = ProcessInfo.processInfo.environment
3946

40-
// Resolve from environment variables if not provided
41-
let resolvedIdentityPoolId = identityPoolId ?? env["AWS_COGNITO_IDENTITY_POOL_ID"]
42-
let resolvedIdentityId = identityId ?? env["AWS_COGNITO_IDENTITY_ID"]
43-
let resolvedAccountId = accountId ?? env["AWS_ACCOUNT_ID"]
44-
let resolvedCustomRoleArn = customRoleArn ?? env["AWS_COGNITO_CUSTOM_ROLE_ARN"]
47+
self.configFilePath = configFilePath
48+
self.credentialsFilePath = credentialsFilePath
49+
self.profileName = profileName
50+
51+
// Load config files
52+
let fileBasedConfig = try CRTFileBasedConfiguration(
53+
configFilePath: configFilePath,
54+
credentialsFilePath: credentialsFilePath
55+
)
56+
let resolvedProfileName = profileName ?? env["AWS_PROFILE"] ?? "default"
57+
58+
// Resolve configuration using helper functions
59+
let resolvedIdentityPoolId = Self.resolveOptionalField(
60+
identityPoolId, "AWS_COGNITO_IDENTITY_POOL_ID", "cognito_identity_pool_id", fileBasedConfig, resolvedProfileName
61+
)
62+
let resolvedIdentityId = Self.resolveOptionalField(
63+
identityId, "AWS_COGNITO_IDENTITY_ID", "cognito_identity_id", fileBasedConfig, resolvedProfileName
64+
)
65+
let resolvedAccountId = Self.resolveOptionalField(
66+
accountId, "AWS_ACCOUNT_ID", "account_id", fileBasedConfig, resolvedProfileName
67+
)
68+
let resolvedCustomRoleArn = Self.resolveOptionalField(
69+
customRoleArn, "AWS_COGNITO_CUSTOM_ROLE_ARN", "cognito_custom_role_arn", fileBasedConfig, resolvedProfileName
70+
)
4571

4672
guard resolvedIdentityPoolId != nil || resolvedIdentityId != nil else {
4773
throw ClientError.invalidValue("Either identityPoolId or identityId must be provided")
@@ -56,18 +82,19 @@ public struct CognitoAWSCredentialIdentityResolver: AWSCredentialIdentityResolve
5682
self.logins = logins
5783
self.customRoleArn = resolvedCustomRoleArn
5884

59-
// Priority: cognitoPoolRegion parameter -> AWS_COGNITO_POOL_REGION -> AWS_REGION
60-
self.cognitoPoolRegion = try cognitoPoolRegion
61-
?? env["AWS_COGNITO_POOL_REGION"]
62-
?? env["AWS_REGION"]
63-
?? { throw ClientError.dataNotFound("AWS region not configured") }()
85+
// Resolve region with fallback from cognito_pool_region to region
86+
self.cognitoPoolRegion = try Self.resolveOptionalField(
87+
cognitoPoolRegion, "AWS_COGNITO_POOL_REGION", "cognito_pool_region", fileBasedConfig, resolvedProfileName
88+
) ?? Self.resolveOptionalField(
89+
nil, "AWS_REGION", "region", fileBasedConfig, resolvedProfileName
90+
) ?? { throw ClientError.dataNotFound("AWS region not configured") }()
6491

6592
self.cachedIdentityId = nil
6693
self.cachedCredentials = nil
67-
self.cachedLogins = nil
94+
self.cachedLogins = logins
6895
}
6996

70-
public mutating func getIdentity(identityProperties: Attributes?) async throws -> AWSCredentialIdentity {
97+
public func getIdentity(identityProperties: Attributes?) async throws -> AWSCredentialIdentity {
7198
guard let identityProperties, let internalCognitoClient = identityProperties.get(
7299
key: InternalClientKeys.internalCognitoIdentityClientKey
73100
) else {
@@ -76,9 +103,6 @@ public struct CognitoAWSCredentialIdentityResolver: AWSCredentialIdentityResolve
76103
+ "Missing IdentityProvidingCognitoIdentityClient in identity properties."
77104
)
78105
}
79-
80-
refreshLock.lock()
81-
defer { refreshLock.unlock() }
82106

83107
// Check if cached credentials are still valid
84108
if let cached = cachedCredentials,
@@ -101,16 +125,12 @@ public struct CognitoAWSCredentialIdentityResolver: AWSCredentialIdentityResolve
101125
return credentials
102126
}
103127

104-
private mutating func resolveIdentityId(using client: IdentityProvidingCognitoIdentityClient) async throws -> String {
128+
private func resolveIdentityId(using client: IdentityProvidingCognitoIdentityClient) async throws -> String {
105129
if let existingId = identityId {
106130
return existingId
107131
}
108132

109-
// If logins changed, clear cached identity ID to force refresh
110-
if loginsChanged() {
111-
cachedIdentityId = nil
112-
}
113-
133+
// Use cached identity ID if logins haven't changed
114134
if let cached = cachedIdentityId, !loginsChanged() {
115135
return cached
116136
}
@@ -126,25 +146,47 @@ public struct CognitoAWSCredentialIdentityResolver: AWSCredentialIdentityResolve
126146
)
127147

128148
cachedIdentityId = resolvedId
149+
cachedLogins = logins
129150
return resolvedId
130151
}
131152

132153
private func loginsChanged() -> Bool {
133-
guard let cached = cachedLogins else {
134-
return logins != nil
154+
// Both nil - no change
155+
if cachedLogins == nil && logins == nil {
156+
return false
135157
}
136158

137-
guard let current = logins else {
159+
// One is nil, other isn't - changed
160+
guard let cached = cachedLogins, let current = logins else {
138161
return true
139162
}
140163

164+
// Compare the dictionaries
141165
return cached != current
142166
}
143167

144-
public mutating func updateLogins(_ newLogins: [String: String]?) {
168+
public func updateLogins(_ newLogins: [String: String]?) {
145169
logins = newLogins
146170
// Clear cached data when logins change
147171
cachedIdentityId = nil
148172
cachedCredentials = nil
173+
cachedLogins = nil
174+
}
175+
176+
private static func resolveOptionalField(
177+
_ configValue: String?,
178+
_ envVarName: String,
179+
_ configFieldName: String,
180+
_ config: CRTFileBasedConfiguration,
181+
_ profileName: String
182+
) -> String? {
183+
FieldResolver(
184+
configValue: configValue,
185+
envVarName: envVarName,
186+
configFieldName: configFieldName,
187+
fileBasedConfig: config,
188+
profileName: profileName,
189+
converter: { String($0) }
190+
).value
149191
}
150-
}
192+
}

‎Sources/Core/AWSSDKIdentity/Sources/AWSSDKIdentity/IdentityProvidingClients/IdentityProvidingCognitoIdentityClient.swift

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
//
77

88
import protocol SmithyIdentity.AWSCredentialIdentityResolver
9+
import struct SmithyIdentity.AWSCredentialIdentity
910

1011
/// Protocol for providing Cognito Identity credentials
1112
public protocol IdentityProvidingCognitoIdentityClient {

‎Sources/Core/AWSSDKIdentity/Tests/AWSSDKIdentityTests/AWSCredentialIdentityResolverTests/CognitoCredentialResolverTests.swift

Lines changed: 115 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,10 @@
77

88
import XCTest
99
import Foundation
10-
@testable import YourModule // Replace with actual module name
10+
@testable import AWSSDKIdentity
11+
import struct Smithy.Attributes
12+
import enum Smithy.ClientError
13+
import struct SmithyIdentity.AWSCredentialIdentity
1114

1215
class CognitoCredentialResolverTests: XCTestCase {
1316

@@ -43,7 +46,7 @@ class CognitoCredentialResolverTests: XCTestCase {
4346
)
4447

4548
let mockClient = MockCognitoClient()
46-
let attributes = Attributes()
49+
var attributes = Attributes()
4750
attributes.set(key: InternalClientKeys.internalCognitoIdentityClientKey, value: mockClient)
4851

4952
let credentials = try await resolver.getIdentity(identityProperties: attributes)
@@ -66,7 +69,7 @@ class CognitoCredentialResolverTests: XCTestCase {
6669
)
6770

6871
let mockClient = MockCognitoClient()
69-
let attributes = Attributes()
72+
var attributes = Attributes()
7073
attributes.set(key: InternalClientKeys.internalCognitoIdentityClientKey, value: mockClient)
7174

7275
let credentials = try await resolver.getIdentity(identityProperties: attributes)
@@ -89,6 +92,115 @@ class CognitoCredentialResolverTests: XCTestCase {
8992
}
9093
}
9194

95+
func testTC4_ExpiredCredentialsRefreshed() async throws {
96+
// TC4: Expired credentials are refreshed using existing IdentityId
97+
var resolver = try CognitoAWSCredentialIdentityResolver(
98+
identityId: "us-west-2:12345678-1234-1234-1234-123456789012",
99+
cognitoPoolRegion: "us-west-2"
100+
)
101+
102+
let mockClient = MockCognitoClient()
103+
// Set expired credentials initially
104+
mockClient.mockCredentials = AWSCredentialIdentity(
105+
accessKey: "ASIAEXAMPLEKEY",
106+
secret: "exampleSecretKey123456789",
107+
expiration: Date(timeIntervalSince1970: 100), // Expired
108+
sessionToken: "exampleSessionToken..."
109+
)
110+
111+
var attributes = Attributes()
112+
attributes.set(key: InternalClientKeys.internalCognitoIdentityClientKey, value: mockClient)
113+
114+
// First call with expired credentials
115+
let credentials1 = try await resolver.getIdentity(identityProperties: attributes)
116+
117+
// Update mock to return fresh credentials
118+
mockClient.mockCredentials = AWSCredentialIdentity(
119+
accessKey: "ASIAEXAMPLEKEY",
120+
secret: "exampleSecretKey123456789",
121+
expiration: Date(timeIntervalSince1970: 2735689600), // Future
122+
sessionToken: "exampleSessionToken..."
123+
)
124+
125+
// Second call should refresh credentials
126+
let credentials2 = try await resolver.getIdentity(identityProperties: attributes)
127+
128+
// Should call GetCredentialsForIdentity twice (no GetId since identityId provided)
129+
XCTAssertEqual(mockClient.getIdCalls.count, 0)
130+
XCTAssertEqual(mockClient.getCredentialsCalls.count, 2)
131+
XCTAssertEqual(mockClient.getCredentialsCalls[0].identityId, "us-west-2:12345678-1234-1234-1234-123456789012")
132+
XCTAssertEqual(mockClient.getCredentialsCalls[1].identityId, "us-west-2:12345678-1234-1234-1234-123456789012")
133+
}
134+
135+
func testTC5_IdentityRefreshedWhenLoginsUpdated() async throws {
136+
// TC5: Identity is refreshed when logins are updated
137+
var resolver = try CognitoAWSCredentialIdentityResolver(
138+
identityPoolId: "us-west-2:test-pool",
139+
cognitoPoolRegion: "us-west-2"
140+
)
141+
142+
let mockClient = MockCognitoClient()
143+
var attributes = Attributes()
144+
attributes.set(key: InternalClientKeys.internalCognitoIdentityClientKey, value: mockClient)
145+
146+
// First call without logins
147+
let credentials1 = try await resolver.getIdentity(identityProperties: attributes)
148+
149+
// Update logins
150+
resolver.updateLogins(["accounts.google.com": "new-google-token"])
151+
152+
// Update mock to return different identity ID for new logins
153+
mockClient.mockIdentityId = "us-west-2:12345678-1234-1234-1234-333333333333"
154+
155+
// Second call with updated logins
156+
let credentials2 = try await resolver.getIdentity(identityProperties: attributes)
157+
158+
// Should call GetId twice (once without logins, once with logins)
159+
XCTAssertEqual(mockClient.getIdCalls.count, 2)
160+
XCTAssertNil(mockClient.getIdCalls[0].logins)
161+
XCTAssertEqual(mockClient.getIdCalls[1].logins?["accounts.google.com"], "new-google-token")
162+
163+
// Should call GetCredentialsForIdentity twice
164+
XCTAssertEqual(mockClient.getCredentialsCalls.count, 2)
165+
XCTAssertNil(mockClient.getCredentialsCalls[0].logins)
166+
XCTAssertEqual(mockClient.getCredentialsCalls[1].logins?["accounts.google.com"], "new-google-token")
167+
}
168+
169+
func testTC6_IdentityNotRefreshedWhenLoginsUnchanged() async throws {
170+
// TC6: Identity is NOT refreshed when logins remain unchanged
171+
var resolver = try CognitoAWSCredentialIdentityResolver(
172+
identityPoolId: "us-west-2:test-pool",
173+
logins: ["accounts.google.com": "google-token"],
174+
cognitoPoolRegion: "us-west-2"
175+
)
176+
177+
let mockClient = MockCognitoClient()
178+
// Set credentials with future expiration
179+
mockClient.mockCredentials = AWSCredentialIdentity(
180+
accessKey: "ASIAEXAMPLEKEY",
181+
secret: "exampleSecretKey123456789",
182+
expiration: Date(timeIntervalSinceNow: 3600), // 1 hour from now
183+
sessionToken: "exampleSessionToken..."
184+
)
185+
186+
var attributes = Attributes()
187+
attributes.set(key: InternalClientKeys.internalCognitoIdentityClientKey, value: mockClient)
188+
189+
// First call
190+
let credentials1 = try await resolver.getIdentity(identityProperties: attributes)
191+
192+
// Second call with same logins (should use cached credentials)
193+
let credentials2 = try await resolver.getIdentity(identityProperties: attributes)
194+
195+
// Should call GetId only once
196+
XCTAssertEqual(mockClient.getIdCalls.count, 1)
197+
XCTAssertEqual(mockClient.getIdCalls[0].logins?["accounts.google.com"], "google-token")
198+
199+
// Should call GetCredentialsForIdentity only once (second call uses cache)
200+
XCTAssertEqual(mockClient.getCredentialsCalls.count, 1)
201+
XCTAssertEqual(mockClient.getCredentialsCalls[0].logins?["accounts.google.com"], "google-token")
202+
}
203+
92204
func testTC7_FallbackToAWSRegion() throws {
93205
// TC7: Fall back to AWS_REGION when a Cognito pool region is not explicitly provided
94206
setenv("AWS_REGION", "us-west-2", 1)

‎codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/customization/credentialresolverservices/InternalModelIntegration.kt

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,11 @@ class InternalModelIntegration : SwiftIntegration {
2121
listOf(
2222
"com.amazonaws.ssooidc#CreateToken",
2323
)
24+
private val cognitoOps =
25+
listOf(
26+
"com.amazonaws.cognito.identity#GetId",
27+
"com.amazonaws.cognito.identity#GetCredentialsForIdentity",
28+
)
2429

2530
override fun enabledForService(
2631
model: Model,
@@ -48,6 +53,7 @@ class InternalModelIntegration : SwiftIntegration {
4853
"STS" -> stsOps
4954
"SSO" -> ssoOps
5055
"SSO OIDC" -> ssoOIDCOps
56+
"Cognito Identity" -> cognitoOps
5157
else -> emptyList()
5258
}
5359
}

0 commit comments

Comments
 (0)