Cross-account access to resources in different AWS accounts using Step Functions

To Nha Notes | Oct. 18, 2023, 4 p.m.

Step Functions provides cross-account access to resources configured in different AWS accounts in your workflows. Using Step Functions service integrations, you can invoke any cross-account AWS resource even if that AWS service doesn’t support resource-based policies or cross-account calls.

For example, we create a step function in staging env to call a lambda function in dev env to trigger MWAA DAGs as following.

Source: Step function state machine:

SAM template

AWSTemplateFormatVersion: "2010-09-09"

Transform: AWS::Serverless-2016-10-31

Description: >

  dpt

 

  Data Platform Toolkit

 

Resources:

  DataOnboardingNextStateMachine:

    Type: AWS::Serverless::StateMachine # More info about State Machine Resource: https://docs.aws.amazon.com/serverless-application-model/latest/developerguide/sam-resource-statemachine.html

    Properties:

      DefinitionUri: statemachine/data-onboarding-next.asl.json

      Role: !GetAtt DataOnboardingNextStateMachineExecutionRole.Arn

 

  DataOnboardingNextStateMachineExecutionRole:

    Type: AWS::IAM::Role

    Properties:

      AssumeRolePolicyDocument:

        Version: '2012-10-17'

        Statement:

          - Effect: Allow

            Principal:

              Service:

                - states.amazonaws.com

            Action: sts:AssumeRole

      Policies:

        - PolicyName: LambdaInvokePolicy

          PolicyDocument:

            Version: '2012-10-17'

            Statement:

              - Effect: Allow

                Action:

                  - lambda:InvokeFunction

                Resource: '*'

        - PolicyName: AssumeTargetRolePolicy

          PolicyDocument:

            Version: '2012-10-17'

            Statement:

              - Effect: Allow

                Action:

                  - sts:AssumeRole

                Resource: '*'

 

Outputs:

  # Find out more about other implicit resources you can reference within SAM

  # https://docs.aws.amazon.com/serverless-application-model/latest/developerguide/sam-specification-generated-resources.html

  DataOnboardingNextStateMachineArn:

    Description: "Data Onboarding Next State machine ARN"

    Value: !Ref DataOnboardingNextStateMachine

  DataOnboardingNextStateMachineRoleArn:

    Description: "IAM Role created for Data Onboarding State machine based on the specified SAM Policy Templates"

    Value: !GetAtt DataOnboardingNextStateMachineExecutionRole.Arn

State machine

{

    "Comment": "A state machine that does data onboarding for new clients.",

    "StartAt": "RDS Hack Update Task",

    "States": {

        "RDS Hack Update Task": {

            "Type": "Task",

            "Resource": "${RDSHackUpdateFunctionArn}",

            "ResultPath": "$.status",

            "Next": "Trigger DSP MWAA DAGs"

        },

        "Trigger DSP MWAA DAGs": {

            "Type": "Task",

            "Resource": "arn:aws:states:::lambda:invoke",

            "Credentials": {

                "RoleArn": "<TARGET_LAMBDA_EXECUTION_ROLE_ARN>"

            },

            "Parameters": {

                "FunctionName": "<TARGET_LAMBDA_ARN>",

                "Payload":{  

                    "mwaa_env_name.$":"$.mwaa_env_name",

                    "mwaa_dag_name.$":"$.mwaa_dag_name",

                    "mwaa_dag_conf.$":"$.mwaa_dag_conf"

                }

            },

            "End": true

        }

    }

}

Execution role policies

AssumeTargetRolePolicy

{
    "Statement": [
        {
            "Action": "sts:AssumeRole",
            "Resource": "<TARGET_LAMBDA_EXECUTION_ROLE_ARN>",
            "Effect": "Allow",
            "Sid": "LambdaInvokePolicy"
        }
    ]
}

LambdaInvokePolicy

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Action": [
                "lambda:InvokeFunction"
            ],
            "Resource": "*",
            "Effect": "Allow"
        }
    ]
}

Trusted entities

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Principal": {
                "Service": "states.amazonaws.com"
            },
            "Action": "sts:AssumeRole"
        }
    ]
}

Target: Lambda function:

SAM template

AWSTemplateFormatVersion: '2010-09-09'

Transform: AWS::Serverless-2016-10-31

Description: >

  dpt

 

  Data Platform Toolkit - MWAA DAGs trigger lambda

 

# More info about Globals: https://github.com/awslabs/serverless-application-model/blob/master/docs/globals.rst

Globals:

  Function:

    Timeout: 3

 

Resources:

  MWAADAGsTriggerFunction:

    Type: AWS::Serverless::Function # More info about Function Resource: https://github.com/awslabs/serverless-application-model/blob/master/versions/2016-10-31.md#awsserverlessfunction

    Properties:

      CodeUri: functions/mwaa_dags_trigger/

      Handler: app.lambda_handler

      Runtime: python3.11

      Timeout: 300

      Architectures:

        - x86_64

      Policies:

        - Statement:

          - Sid: LambdaInvokePolicy

            Effect: Allow

            Action:

              - lambda:InvokeFunction

            Resource: '*'

        - Statement:

          - Sid: MWAADAGsTriggerPolicy

            Effect: Allow

            Action:

              - airflow:ListTagsForResource

              - airflow:CreateCliToken

              - airflow:GetEnvironment

              - airflow:ListEnvironments

            Resource: '*'

     

 

Outputs:

  # ServerlessRestApi is an implicit API created out of Events key under Serverless::Function

  # Find out more about other implicit resources you can reference within SAM

  # https://github.com/awslabs/serverless-application-model/blob/master/docs/internals/generated_resources.rst#api

  MWAADAGsTriggerFunction:

    Description: "MWAA DAGs Trigger Lambda Function ARN"

    Value: !GetAtt MWAADAGsTriggerFunction.Arn

  MWAADAGsTriggerFunctionIamRole:

    Description: "Implicit IAM Role created for Hello World function"

    Value: !GetAtt MWAADAGsTriggerFunctionRole.Arn

Execution role policies

MWAADAGsTriggerFunctionRolePolicy0

{
    "Statement": [
        {
            "Action": [
                "lambda:InvokeFunction"
            ],
            "Resource": "<TARGET_LAMBDA_FUNCTION_ARN>",
            "Effect": "Allow",
            "Sid": "LambdaInvokePolicy"
        }
    ]
}

MWAADAGsTriggerFunctionRolePolicy1

{
    "Statement": [
        {
            "Action": [
                "airflow:ListTagsForResource",
                "airflow:CreateCliToken",
                "airflow:GetEnvironment",
                "airflow:ListEnvironments"
            ],
            "Resource": "*",
            "Effect": "Allow",
            "Sid": "MWAADAGsTriggerPolicy"
        }
    ]
}

Trusted entities

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Principal": {
                "Service": "lambda.amazonaws.com"
            },
            "Action": "sts:AssumeRole"
        },
        {
            "Effect": "Allow",
            "Principal": {
                "AWS": "<SOURCE_STEP_FUNCTION_EXECUTION_ROLE_ARN>"
            },
            "Action": "sts:AssumeRole",
            "Condition": {
                "StringEquals": {
                    "sts:ExternalId": "<SOURCE_STEP_FUNCTION_ARN>"
                }
            }
        }
    ]
}

References

https://docs.aws.amazon.com/step-functions/latest/dg/concepts-access-cross-acct-resources.html

https://docs.aws.amazon.com/step-functions/latest/dg/tutorial-access-cross-acct-resources.html

https://docs.aws.amazon.com/step-functions/latest/dg/connect-lambda.html