AWS Step Functionsのタスクから別アカウントのStep Functionsを実行する

こんにちは。株式会社AlphaDriveでWebアプリケーションエンジニアをしている @tatsuhiro_hori と申します。
この記事では、AWS Step Functionsのタスクから別アカウントのStep Functionsを実行する方法について紹介したいと思います。
この記事は「[AlphaDrive Advent Calendar 2023](https://qiita.com/advent-calendar/2023/alphadrive-tech)」の12/22のエントリーです。

# この記事を書くに至った背景
私の勤めているユーザベースグループでは、様々なSaaSを各グループ企業が開発、サービス提供しています。
その中にはAWSのマネージドサービスを利用して構築されたものもありますが、AWSアカウントはグループ共通ではなく、各会社がそれぞれ別で管理しています。
なので各SaaS同士の連携機能はAWSアカウントを跨いだものになります。
そして私が対応することになったのが、「アカウントAのStepFunctionからアカウントBのStepFunctionを実行する」というケースです。
「アカウントを跨いだStepFunctionsの実行ってどうやってやるんだろう？（よくわからんけど面倒くさそう）」と思いましたが、調べてみると別アカウントのRoleにスイッチする方法（AssumeRole）があることが分かり、やってみたら意外とすぐできたので、今回紹介しようと思いました。

「AssumeRoleとはなんぞや？」という話はこちらの記事などで紹介されていますのでご参考ください。
[AssumeRoleとはなんぞや！](https://qiita.com/ms2geki/items/26eba61c2c405acae365)

# 全体像
今回説明する例では、AWSアカウントとStepFunctionsのStateMachineがそれぞれ2つずつ登場するのですが、便宜上、別アカウントのStateMachineを実行する側を「Operator」、実行される側を「Subject」と名付けます（A、Bで説明するとどっちがどっちか分からなくなるので）。
すなわち、「OperatorアカウントのOperatorStateMachineからSubjectアカウントのSubjectStateMachineをAssumeRoleを利用して実行する」というケースになります。
そしてそれを図にするとこんな感じになります。

![assume_role.drawio.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/627873/47877238-0381-a56f-381e-ba1475ca6a59.png)

各StateMachineにはIAM Roleが紐づいています（図の「OperatorStateMachine IAMRole」と「SubjectStateMachine IAMRole）。
これらはそのStateMachineがリソース（S3やLambda関数など）にアクセスしたり実行するための権限を設定するためのものです。
これとは別にSubjectアカウントには「SubjectStateMachineを実行するためのIAM Role」を用意します（図の「SubjectStateMachine ExecuteRole」）。
このExecuteRoleをAssumeRoleすることでOperatorStateMachineはSubjectStateMachineを実行できるようになります。

# 設定手順
[こちらの記事](https://dev.classmethod.jp/articles/aws-step-functions-now-allows-you-to-assumerole-for-each-task/)では同様にAssumeRoleを使ってStepFunctionsからLambdaを実行する手順について紹介されていますが、基本的な流れは同じです。
全体の流れとしては以下の通りです。
1. SubjectStateMachineのExecuteRoleの設定
2. OperatorStateMachineのIAM Role設定
3. OperatorStateMachineのTaskStateの設定

## 1. SubjectStateMachineのExecte Roleの設定
SubjectStateMachineをOperatorアカウントのStateMachineから実行できるようにするために、以下の設定がされたExecuteRoleを用意します。
- SubjectStateMachineを実行する権限
- OperatorアカウントのOperatorStateMachineを信頼する設定

Permission PoliciesにはシンプルにSubjectStateMachineを実行するための権限を設定します。
```json
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Sid": "ExecuteSubjectStateMachine",
            "Effect": "Allow",
            "Action": "states:StartExecution",
            "Resource": [
                "arn:aws:states:ap-northeast-1:${SubjectアカウントのID}:stateMachine:subject-statemachine"
            ]
        }
    ]
}
```

Trusted PoliciesにはOperatorアカウントのOperatorStateMachineからAssumeRoleを許可する設定を追加します。
```json
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Sid": "OperatorStateMachine",
            "Effect": "Allow",
            "Principal": {
                "AWS": "arn:aws:iam::${OperatorアカウントのID}:role/operator-statemachine-iam-role"
            },
            "Action": "sts:AssumeRole",
            "Condition": {
                "StringEquals": {
                    "sts:ExternalId": "arn:aws:states:ap-northeast-1:${OperatorアカウントのID}:stateMachine:operator-statemachine"
                }
            }
        }
    ]
}
```
## 2. OperatorStateMachineのIAM Roleの設定
OperatorStateMachineのIAM Roleに、SubjectアカウントのSubjectStateMachineを実行するための権限を設定します。
Permission Policiesに、SubjectStateMachineを実行するためのStartExecutionと、SubjectStateMachine　ExectionRole（1で設定したRole）のAssumeRoleを追加しています。
```json
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Action": "states:StartExecution",
            "Resource": "arn:aws:states:ap-northeast-1:${SubjectアカウントのID}:stateMachine:subject-statemachine",
            "Effect": "Allow"
        },
        {
            "Action": "sts:AssumeRole",
            "Resource": "arn:aws:iam::${SubjectアカウントのID}:role/execute-subject-statemachine-role",
            "Effect": "Allow"
        },
}
```

## 3. OperatorStateMachineのTaskStateの設定
StateMachineのASLは以下のようになります。
```json
{
  "Comment": "A description of my state machine",
  "StartAt": "Step Functions StartExecution",
  "States": {
    "Step Functions StartExecution": {
      "Type": "Task",
      "Resource": "arn:aws:states:::states:startExecution.sync:2",
      "Parameters": {
        "StateMachineArn": "arn:aws:states:ap-northeast-1:${SubjectアカウントのID}:stateMachine:subject-statemachine",
        "Input": {
          "StatePayload": "Hello from Step Functions!",
          "AWS_STEP_FUNCTIONS_STARTED_BY_EXECUTION_ID.$": "$$.Execution.Id"
        }
      },
      "End": true,
      "Credentials": {
        "RoleArn": "arn:aws:iam::${SubjectアカウントのID}:role/subject-statemachine-execute-role"
      }
    }
  }
}
```
StateMachineArnにSubjectアカウントのSubjectStateMachineを指定し、CredentialsにSubjectStateMachine ExecuteRole（1で設定したRole）を指定します。
ビジュアルエディタだと「IAM role for cross-account access」というオプションで設定できます。
![ビジュアルエディタ.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/627873/3e5ee61e-e500-e297-2982-3ba4a1b136a6.png)

私が扱っているOperatorStateMachineは、CDK（TypeScript）で実装されており、そこにSubjectStateMachineを実行するためのTaskStateを追加する実装をしました。
その際のコードもサンプルコードとして紹介します。
```ts
import {
    ArnFormat,
    aws_iam as iam,
    aws_stepfunctions as sfn,
    aws_stepfunctions_tasks as tasks,
    Stack,
} from "aws-cdk-lib";
import { Construct } from "constructs";

export class ResourceStateAwsSubjectAccount extends Construct {

    constructor(scope: Construct, id: string) {
        super(scope, id);
    }

    taskStateExecuteSubjectStateMachine(accountId: string, region: string, stateMachineName: string, executeRoleName: string): sfn.IChainable & sfn.INextable {
        return new tasks.StepFunctionsStartExecution(
            this,
            "Execute SubjectStateMachine in AWS SubjectAccount",
            {
                // SubjectStateMachine
                stateMachine: sfn.StateMachine.fromStateMachineArn(
                    this,
                    "subject-statemachine",
                    this.getStateMachineArn(
                        accountId,
                        region,
                        statemachineName
                    )
                ),
                // ASLのCredentialsの部分
                credentials: {
                    role: this.getCrossAccountRole(
                        accountId,
                        executeRoleName
                    ),
                },
            }
        );
    }

    // SubjectStateMachineのARNの組み立て
    getStateMachineArn(
        accountId: string,
        region: string,
        statemachineName: string
    ): string {
        return Stack.of(this).formatArn({
            partition: "aws",
            service: "states",
            region: region,
            account: accountId,
            resource: "stateMachine",
            resourceName: statemachineName,
            arnFormat: ArnFormat.COLON_RESOURCE_NAME,
        });
    }

    // SubjectStateMachine ExecuteRoleのARNの組み立て
    getCrossAccountRole(accountId: string, executeRoleName: string): sfn.TaskRole {
        const roleArn = Stack.of(this).formatArn({
            partition: "aws",
            service: "iam",
            region: "",
            account: accountId,
            resource: "role",
            resourceName: executeRoleName,
            arnFormat: ArnFormat.SLASH_RESOURCE_NAME,
        });
        return sfn.TaskRole.fromRole(
            iam.Role.fromRoleArn(this, executeRoleName, roleArn, {
                mutable: false,
            })
        );
    }
}

```

## おわりに
今回紹介したAssumeRoleはStepFunctions以外のサービスへ（Lambdaなど）にも利用できます。（参考： https://aws.amazon.com/jp/about-aws/whats-new/2022/11/simplify-cross-account-access-aws-services-step-functions/）
StepFunctionsは、AWSの様々なサービスを組み合わせてワークフローを作成することができる便利なサービスなので、これからも何かとお世話になることが増えてくると思います。
クロスアカウントのStepFunctions->StepFunctionsの実行についてお困りの方にとってこの記事が参考になれば幸いです。

## 参考にした記事
[AWS Step Functionsでタスク毎に別アカウントのロールにAssumeRole（スイッチロール）出来るようになりました！](https://dev.classmethod.jp/articles/aws-step-functions-now-allows-you-to-assumerole-for-each-task/)
[AssumeRoleとはなんぞや！](https://zenn.dev/frusciante/articles/28cd351fe9de60)