Recently we started a big refactor of an application that was running on another team's AWS Account.
Our plan is to move to a serverless architecture using our own account to gradually dismiss the old one.

In order to achieve that, we deployed a Lambda on one account which was in charge of forwarding every request to other AWS Services in another account.
Some of the services we invoke cross-account are SSM ( to retrieve credentials and other configuration stuff), SQS ( to push POST and PATCH requests to a queue to handle batch processing and retries, DynamoDB, to read data in case of GET requests.

Create a Cross Account Role

To access resources CrossAccount we need to a Role that one account can Assume to "impersonate" the other account, and also create a Policy with the right permissions about the service and actions you want to trigger.

In the new account / CDK Stack,

   const crossAccountRole = new Role(this, `cross-account-role`, {
            assumedBy: new AccountPrincipal(this.node.tryGetContext('otherAccountId')),
            roleName: `ross-account-invoker-role`,
            description: `The role grants rights to otherAccount to access resources on this account`,
        })

Here the only interesting part is that we need to add to our cdk.context.json the reference to otherAccountID.

Add Policies

Then assign Policies for every service we will be using:

        crossAccountRole.attachInlinePolicy(
            new Policy(this, `lambda-ssm-policy`, {
                policyName: `lambda-ssm-policy`,
                statements: [ 
                    new PolicyStatement({
                        actions: ['ssm:GetParameter', 'ssm:DescribeParameters', 'ssm:GetParameterHistory', 'ssm:GetParameters'],
                        effect: Effect.ALLOW,
                        resources: [
                            SSM_PARAM_ARN,
                        ],
                    }),
                ],
            }),
        )

The first policy allows the role to getParameters from SSM.
The resource ARN must be specified.
If SSM parameter was created via CDK it is possible to reference it like this: mySSM.parameterArn but hardcoding the reference again in the cdk.context.json and using it like this arn:aws:ssm:${this.region}:${this.account}:parameter${ssmKey} is also a possibility.

        crossAccountRole.attachInlinePolicy(
            new Policy(this, `lambda-dynamodb-policy`, {
                policyName: `lambda-dynamodb-policy`,
                statements: [
                    new PolicyStatement({
                        actions: ['dynamodb:GetItem'],
                        resources: [
                           TABLE_ARN,
                        ],
                    }),
                ],
            }),
        )

Similarly, here we are assigning Permissions to retrieve items from DynamoDB.

While in the following we grant rights to push items to a queue.

        crossAccountREDSTPRole.attachInlinePolicy(
            new Policy(this, `lambda-queue-policy`, {
                policyName: `lambda-queue-policy`,
                statements: [
                    new PolicyStatement({
                        actions: ['sqs:SendMessage'],
                        resources: [
                            QUEUE_ARN,
                        ],
                    }),
                ],
            }),
        )

Having this granular policies in your cross account role are very important to expose access only to the services you plan to use and not have security issues.
It might be a bit harder at the beginning when you are not sure if you are doing everything right, or confusing when you in few weeks or months add another service ( or like in our example, try to Put an Item in Dynamo) and it does not work until you update the policies but it is better to stick to the The Principle of Least Privilege.

a subject should be given only those privileges needed for it to complete its task. If a subject does not need an access right, the subject should not have that right.

Now that we have our role and policies set up in our main account, let's go to our Lambda in the other account and assume the CrossAccount Role so that we can init our AWS Services with the right Credentials.

First implementation (kinda working, until..)

let credentials: Credentials
const getCrossAccountCredentials = async () => {
    const client = new STSClient({region: process.env.REGION})
    const params = {
        RoleArn: process.env.CROSS_ACCOUNT_ROLE_ARN!, //role created on main role which allows this role to invoke resources on the main role.
        RoleSessionName: `assume-role-to-call-main-account`
    }
    const command = new AssumeRoleCommand(params)
    return client.send(command)
        .then(data => {
            if (!data.Credentials) {
                throw new Error(`Got successful AssumeRoleResponse but Credentials are undefined  ${JSON.stringify(data)}`)
            }
            const  cred = new Credentials({
                accessKeyId: data.Credentials.AccessKeyId!,
                secretAccessKey: data.Credentials.SecretAccessKey!,
                sessionToken: data.Credentials.SessionToken,
            })
            return cred
        }).catch(error => {
            throw new Error(`Cannot get credentials to call the lambda \n ${error}`)
        })
}

In our handler we were first checking if we had stored the credentials and if not we invoke STS AssumeRole.
Then we passed the loaded credentials to our AWS Client.

    if (!credentials) {
        credentials = await getCrossAccountCredentials()
    }
 const client = new SSMClient({credentials, region: process.env.REGION})

After deploying and testing everything worked, until we realised that when the Lambda Container stays up for a long time ( due to frequent requests), the Credentials expire.

Just check for Expire?

First quick fix was checking in the handler if the Credentials were expired and eventually reloading it.

    if (!credentials || credentials.expired) {
        credentials = await getCrossAccountCredentials()
    }

But it didn't feel right and especially with multiple services and AWS clients we wanted to have some better code in place in order to refresh the credentials automatically.

Extend Credentials class to handle refresh

By checking the docs we found some interesting methods:

• needsRefresh
• refresh

which could be used by subclasses overriding them.

Therefore we changed our method into a Class:

class CrossAccountCredentials extends Credentials {
    public static async build(callback: () => Promise<{accessKeyId: string, secretAccessKey: string, sessionToken: string, expireTime: Date }>) : Promise<CrossAccountCredentials> {
        const {
            accessKeyId,
            secretAccessKey,
            sessionToken,
            expireTime,
        } = await callback()
        const cred = new CrossAccountCredentials({
            accessKeyId,
            secretAccessKey,
            sessionToken,
        })
        cred.expireTime = expireTime
        return cred
    }

    async refresh(callback: (err?: AWSError) => void = console.error): Promise<void> {
        const {
            accessKeyId,
            secretAccessKey,
            sessionToken,
            expireTime,
        } = await assumeRole()

        this.accessKeyId = accessKeyId
        this.secretAccessKey = secretAccessKey
        this.sessionToken = sessionToken
        this.expireTime = expireTime
        return super.refresh(callback)
    }
}

export default CrossAccountCredentials.build(assumeRole)

The key part here is that in the assumeRole we don't return directly the Parent Class, but an object which contains also ExpiredTime information

// inside AssumeRole method - previously getCrossAccountCredentials

 return {
                accessKeyId: data.Credentials.AccessKeyId!,
                secretAccessKey: data.Credentials.SecretAccessKey!,
                sessionToken: data.Credentials.SessionToken!,
                expireTime: data.Credentials.Expiration!,
            }

Also, instead of manually checking if Credentials exist and are expired and eventually reload them, we could safely every time invoke await credentials.getPromise() which does that for us.

Not yet optimal

This solution was working fine, but there was something that was really bothering me.

We were storing the Credentials and refresh them, but we were, at every run of the lambda handler, instantiate a new AWS Client in order for it to receive valid credentials.

const  ssm = new SSMClient(
{credentials, region: process.env.REGION}
))

I can't tell how much this is a problem, but it is a general recommendation to initialize SDK clients outside of the function handler. Subsequent invocations processed by the same instance of your function can reuse these resources. This saves cost by reducing function run time.

Therefore I dug in the docs even more, I checked the source code of AWS SDK Client in my IDE and and also thanks to Typescript, I noticed something interesting.

Every Client accepts a configuration object which contains the Credentials to be used. Either the Credentials, or a Provider of Credentials!

What does that Provider do?

A function that, when invoked, returns a promise that will be fulfilled with a value of type T.
Example:A function that reads credentials from shared SDK configuration files, assuming roles and collecting MFA tokens as necessary.

How to use it?
By looking up the CredentialsProvider in the new AWS SDK V3 docs I immediately found what I was looking for .fromTemporaryCredentials
that returns a CredentialProvider that retrieves temporary credentials from STS AssumeRole API.

I throw away most of the code that was written previously to extend the Credential Class and handle the refresh method and I simply passed the ARN of my CrossAccountRole to the fromTemporaryCredentials to get a Provider<Credentials> which will take care of the refresh automatically.

import {AssumeRoleCommandInput} from '@aws-sdk/client-sts'
import {fromTemporaryCredentials} from '@aws-sdk/credential-providers'

const assumeRoleParams: AssumeRoleCommandInput = {
     RoleArn: process.env.CROSS_ACCOUNT_ROLE_ARN!, //role created on main role which allows this role to invoke resources on the main role.
     RoleSessionName: `assume-role-to-call-main-account`
    }

By simply passing that CredentialsProvider in the ServiceConfig, I can initiate my services outside the Lambda handler and reuse them, but be sure that when the Credentials are expired, they will be refreshed, without me even knowing it!

const sqs = new SQSClient({
    credentials: fromTemporaryCredentials({params: assumeRoleParams, clientConfig: {region: process.env.REGION}}),
    region: process.env.REGION,
})

It took us a bit to get to this final version (unless someone has some hints or different approaches - feel free to leave a comment! ) but now it shines in its simplicity.

Of course when developing a feature it is important to get shit done and deliver, and as this story shows, we did go live with the first working implementation but then we iterate on it until we found the proper way of doing it.

Never stop at the first quick solution that works, try to improve your code, and most importantly read the docs to really understand what you are using.

Hope it helps!

Photo by Marc-Olivier Jodoin on Unsplash