Skip to content

Instantly share code, notes, and snippets.

@adam-fowler
Last active October 29, 2020 16:19
Show Gist options
  • Save adam-fowler/31199b6746d6c9727a3b53b9542cb7a4 to your computer and use it in GitHub Desktop.
Save adam-fowler/31199b6746d6c9727a3b53b9542cb7a4 to your computer and use it in GitHub Desktop.
enum RequestProvider<Request> {
    case `static`(Request)
    case dynamic((EventLoop)->EventLoopFuture<Request>)

    func getRequest(on eventLoop: EventLoop) -> EventLoopFuture<Request> {
        switch self {
        case .static(let request):
            return eventLoop.makeSucceededFuture(request)
        case .dynamic(let requestFunction):
            return requestFunction(eventLoop)
        }
    }
}

struct AssumeRoleCredentialProvider: CredentialProviderWithClient {
    let requestProvider: RequestProvider<STS.AssumeRoleRequest>
    let client: AWSClient
    let sts: STS

    init(
        request: RequestProvider<STS.AssumeRoleRequest>,
        credentialProvider: CredentialProviderFactory,
        region: Region,
        httpClient: AWSHTTPClient
    ) {
        self.client = AWSClient(credentialProvider: credentialProvider, httpClientProvider: .shared(httpClient))
        self.sts = STS(client: self.client, region: region)
        self.requestProvider = request
    }

    func getCredential(on eventLoop: EventLoop, logger: Logger) -> EventLoopFuture<Credential> {
        return requestProvider.getRequest(on: eventLoop).flatMap { request in
            sts.assumeRole(request, on: eventLoop)
        }.flatMapThrowing { response in
            guard let credentials = response.credentials else { throw CredentialProviderError.noProvider }
            return RotatingCredential(
                accessKeyId: credentials.accessKeyId,
                secretAccessKey: credentials.secretAccessKey,
                sessionToken: credentials.sessionToken,
                expiration: credentials.expiration
            )
        }
    }
}

extension CredentialProviderFactory {
    /// Use AssumeRole to provide credentials
    /// - Parameters:
    ///   - request: AssumeRole request structure
    ///   - credentialProvider: Credential provider used in client that runs the AssumeRole operation
    ///   - region: Region to run request in
    public static func stsAssumeRole(
        request: STS.AssumeRoleRequest,
        credentialProvider: CredentialProviderFactory = .default,
        region: Region
    ) -> CredentialProviderFactory {
        .custom { context in
            let provider = STS.AssumeRoleCredentialProvider(
                requestProvider: .static(request),
                credentialProvider: credentialProvider,
                region: region,
                httpClient: context.httpClient
            )
            return RotatingCredentialProvider(context: context, provider: provider)
        }
    }

    /// Use AssumeRole to provide credentials
    /// - Parameters:
    ///   - request: Function that returns a EventLoopFuture to be fulfillled with an AssumeRole request structure
    ///   - credentialProvider: Credential provider used in client that runs the AssumeRole operation
    ///   - region: Region to run request in
    public static func stsAssumeRole(
        requestProvider: @escaping (EventLoop)->EventLoopFuture<STS.AssumeRoleRequest>,
        credentialProvider: CredentialProviderFactory = .default,
        region: Region
    ) -> CredentialProviderFactory {
        .custom { context in
            let provider = STS.AssumeRoleCredentialProvider(
                requestProvider: .dynamic(requestProvider),
                credentialProvider: credentialProvider,
                region: region,
                httpClient: context.httpClient
            )
            return RotatingCredentialProvider(context: context, provider: provider)
        }
    }
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment