An AWS CDK KMS multi-region key and replica key defined using Kotlin. This aims to solve the issue: Add support for MultiRegion Key with ReplicaKey from https://github.com/aws/aws-cdk/issues/15257.

MultiRegionKey.kt

package com.company.cdk

import software.amazon.awscdk.Stack
import software.amazon.awscdk.customresources.AwsCustomResource
import software.amazon.awscdk.customresources.AwsCustomResourcePolicy
import software.amazon.awscdk.customresources.AwsSdkCall
import software.amazon.awscdk.customresources.PhysicalResourceId
import software.amazon.awscdk.customresources.PhysicalResourceIdReference
import software.amazon.awscdk.services.iam.AccountRootPrincipal
import software.amazon.awscdk.services.iam.PolicyDocument
import software.amazon.awscdk.services.iam.PolicyStatement
import software.amazon.awscdk.services.kms.CfnAlias
import software.amazon.awscdk.services.kms.CfnKey
import software.amazon.awscdk.services.kms.IKey
import software.amazon.awscdk.services.kms.Key
import software.amazon.awscdk.services.kms.KeyLookupOptions
import software.amazon.awscdk.services.kms.KeySpec
import software.amazon.awscdk.services.kms.KeyUsage
import software.constructs.Construct

class MultiRegionKey private constructor(
    key: IKey,
    val replicaRegion: String,
) : IKey by key {
    data class Props(
        val alias: String,
        val description: String,
        val replicaRegion: String, // Currently assumes there is only a single replica region
        // The following are out of the box defaults from CfnKey
        val pendingWindowInDays: Int = 7,
        val enableKeyRotation: Boolean = false,
        val keyPolicy: PolicyDocument = defaultKeyPolicy,
        val keySpec: KeySpec = KeySpec.SYMMETRIC_DEFAULT,
        val keyUsage: KeyUsage = KeyUsage.ENCRYPT_DECRYPT,
    )

    companion object {
        operator fun invoke(
            scope: Construct,
            id: String,
            props: Props,
        ): MultiRegionKey {
            val key =
                if (Stack.of(scope).region != props.replicaRegion) {
                    CfnKey.Builder
                        .create(scope, id)
                        .multiRegion(true)
                        .description(props.description)
                        .keySpec(props.keySpec.name)
                        .keyUsage(props.keyUsage.name)
                        .keyPolicy(props.keyPolicy.toJSON())
                        .enableKeyRotation(props.enableKeyRotation)
                        .pendingWindowInDays(props.pendingWindowInDays)
                        .build()
                        .also {
                            it.overrideLogicalId()
                            it.addAlias("${id}Alias", props.alias)
                            it.replicateKey(
                                replicaRegion = props.replicaRegion,
                                alias = props.alias,
                                pendingWindowInDays = props.pendingWindowInDays,
                            )
                        }
                        .let(Key::fromCfnKey)
                } else {
                    // KMS multi-region replica key is already provisioned by the custom resource in the primary region.
                    // Hence, only return a reference `IKey` of the replica key, looked up using the alias
                    Key.fromLookup(
                        scope,
                        id,
                        KeyLookupOptions
                            .builder()
                            .aliasName(props.alias.asAliasName())
                            .build()
                    )
                }

            return MultiRegionKey(
                key = key,
                replicaRegion = props.replicaRegion,
            )
        }

        private val defaultKeyPolicy: PolicyDocument = allowRootAccess()

        private fun allowRootAccess(): PolicyDocument =
            PolicyDocument.Builder
                .create()
                .statements(listOf(
                    PolicyStatement.Builder
                        .create()
                        .principals(listOf(AccountRootPrincipal()))
                        // This is a resource-based policy attached against a KMS key.
                        // "*" refers to "this KMS key"
                        // Reference: https://docs.aws.amazon.com/kms/latest/developerguide/key-policy-overview.html
                        .resources(listOf("*"))
                        .actions(listOf("kms:*"))
                        .build()
                ))
                .build()
    }
}

/**
 * This extension function handles the creation, update, and deletion on the replica key in the replica region.
 * Given a multi-region CfnKey instance, the replica key is created using custom resources that uses AWS SDK for
 * JavaScript under the hood.
 *
 * Creating the replica key can be broken down into two steps:
 * 1. Using `replicateKey` SDK call [1], create the replica key in the replica region
 * 2. Using `createAlias` SDK call [2], set the key alias of the replica key to be the same alias as the primary key
 *
 * References:
 * - [1]: https://docs.aws.amazon.com/AWSJavaScriptSDK/latest/AWS/KMS.html#replicateKey-property
 * - [2]: https://docs.aws.amazon.com/AWSJavaScriptSDK/latest/AWS/KMS.html#createAlias-property
 */
private fun CfnKey.replicateKey(
    replicaRegion: String,
    alias: String,
    pendingWindowInDays: Int,
): ReplicaKey =
    ReplicaKey(
        alias = alias,
        pendingWindowInDays = pendingWindowInDays,
        primaryKey = this,
        replicaRegion = replicaRegion,
    ).also { replicaKey ->
        val primaryKey = this

        val replicaKeyProvider = AwsCustomResource.Builder
            .create(primaryKey, "ReplicaKeyProvider")
            .onCreate(
                AwsSdkCall
                    .builder()
                    .service("KMS")
                    .action("replicateKey")
                    .parameters(mapOf(
                        "KeyId" to primaryKey.keyId,
                        "ReplicaRegion" to replicaKey.replicaRegion,
                    ))
                    .physicalResourceId(PhysicalResourceId.of(replicaKey.replicaRegion))
                    .build()
            )
            .onDelete(
                AwsSdkCall
                    .builder()
                    .service("KMS")
                    .action("scheduleKeyDeletion")
                    .region(replicaKey.replicaRegion)
                    .parameters(mapOf(
                        "KeyId" to replicaKey.keyId,
                        "PendingWindowInDays" to replicaKey.pendingWindowInDays,
                    ))
                    .physicalResourceId(PhysicalResourceId.of(replicaKey.replicaRegion))
                    .build()
            )
            .policy(
                AwsCustomResourcePolicy.fromStatements(
                    listOf(
                        // Allow replicateKey API call for primary key
                        PolicyStatement.Builder
                            .create()
                            .actions(listOf("kms:ReplicateKey"))
                            .resources(listOf(primaryKey.keyArn))
                            .build(),
                        // Allow createKey API (must be used in tandem with replicateKey)
                        PolicyStatement.Builder
                            .create()
                            .actions(listOf("kms:CreateKey"))
                            // "*" is expected to be used since `CreateKey` does not use any particular AWS KMS resource
                            // https://docs.aws.amazon.com/kms/latest/developerguide/customer-managed-policies.html#iam-policy-example-create-key
                            .resources(listOf("*"))
                            .build(),
                        // Allow deletion of the replica key
                        PolicyStatement.Builder
                            .create()
                            .actions(listOf(
                                // Adding tags to the replica key is not yet supported in this construct
                                // "kms:TagResource",
                                "kms:ScheduleKeyDeletion",
                            ))
                            .resources(listOf(replicaKey.keyArn))
                            .build(),
                    )
                )
            )
            .build()

        /**
         * This construct assumes that the key can only be replicated in a single replica region.
         * This assumption does not reflect the actual behavior of a multi-region KMS key but is done for simplifying
         * the management of a single replica key across a region. Conveniently enough, this also covers our current
         * use-case.
         *
         * With this assumption, if the replica region changed, then we want to mark the update behavior to be a `replacement`:
         *  1. Create a new replica key in the new replica region
         *  2. Mark for deletion the old replica key in the old region
         *
         * By default, when `onCreate` is not specified and only `onUpdate` is specified, CDK will use the `onUpdate` for both
         * create and update events.
         */
        AwsCustomResource.Builder
            .create(primaryKey, "ReplicaKeyAliasProvider")
            .onUpdate(
                AwsSdkCall
                    .builder()
                    .service("KMS")
                    .action("createAlias")
                    .region(replicaKey.replicaRegion)
                    .parameters(mapOf(
                        "TargetKeyId" to replicaKey.keyId,
                        "AliasName" to replicaKey.aliasName,
                    ))
                    .physicalResourceId(PhysicalResourceId.of(replicaKey.aliasName))
                    .build()
            )
            .onDelete(
                AwsSdkCall
                    .builder()
                    .service("KMS")
                    .action("deleteAlias")
                    .region(replicaKey.replicaRegion)
                    .parameters(
                        mapOf(
                            "AliasName" to PhysicalResourceIdReference(),
                        )
                    )
                    .build()
            )
            .policy(
                AwsCustomResourcePolicy.fromStatements(
                    listOf(
                        PolicyStatement.Builder
                            .create()
                            .actions(listOf("kms:CreateAlias"))
                            .resources(listOf(
                                replicaKey.keyArn,
                                replicaKey.aliasArn,
                            ))
                            .build(),
                        PolicyStatement.Builder
                            .create()
                            .actions(listOf("kms:DeleteAlias"))
                            .resources(listOf(
                                replicaKey.keyArn,
                                // Using "*" since the DeleteAlias operation should be able to delete the old alias during an
                                // update with replacement event. CDK does not have the reference to the old alias, just the
                                // new one
                                KmsArn.aliasArn(
                                    region = replicaKey.replicaRegion,
                                    account = replicaKey.account,
                                    alias = "*",
                                ),
                            ))
                            .build()
                    )
                )
            )
            .build()
            .also {
                it.node.addDependency(replicaKeyProvider)
            }
    }

private data class ReplicaKey(
    private val alias: String,
    val pendingWindowInDays: Int,
    private val primaryKey: CfnKey,
    val replicaRegion: String,
) {
    val keyId = primaryKey.keyId // In KMS, the primary key and replica key share the same key ID (`mrk-xxxxxx`)
    val aliasName = alias.asAliasName()
    val account = Stack.of(primaryKey).account
    val keyArn = KmsArn.keyArn(region = replicaRegion, account = account, keyId = keyId)
    val aliasArn = KmsArn.aliasArn(region = replicaRegion, account = account, alias = alias)
}

private fun CfnKey.addAlias(id: String, alias: String) =
    CfnAlias.Builder
        .create(this, id)
        .aliasName(alias.asAliasName())
        .targetKeyId(keyId)
        .build()
        .also {
            it.overrideLogicalId()
        }

private val CfnKey.keyArn: String
    get() =
        Stack.of(this)
            .let { stack ->
                KmsArn.keyArn(
                    region = stack.region,
                    account = stack.account,
                    keyId = keyId,
                )
            }

private val CfnKey.keyId: String
    get() = attrKeyId

private fun String.asAliasName() = "alias/$this"

private object KmsArn {
    fun keyArn(
        region: String,
        account: String,
        keyId: String,
    ): String = arn(region = region, account = account, resourceId = "key/$keyId")

    fun aliasArn(
        region: String,
        account: String,
        alias: String,
    ): String = arn(region = region, account = account, resourceId = alias.asAliasName())

    private fun arn(
        region: String,
        account: String,
        resourceId: String,
    ): String = "arn:aws:kms:$region:$account:$resourceId"
}

fun IConstruct.overrideLogicalId(logicalId: String = node.id) {
    when (this) {
        is CfnResource -> this.overrideLogicalId(logicalId)
        else -> (this.node.defaultChild as CfnResource).overrideLogicalId(logicalId)
    }
}

SampleUsage.kt

fun main() {
    val key = MultiRegionKey(
        scope = stack,
        id = "MultiRegionKey",
        props = MultiRegionKey.Props(
            alias = "my-multi-region-key",
            description = "a multi-region kms key",
            replicaRegion = "eu-west-1",
            pendingWindowInDays = 7,
            enableKeyRotation = true,
            keySpec = KeySpec.SYMMETRIC_DEFAULT,
            keyUsage = KeyUsage.ENCRYPT_DECRYPT,
        )
    )
}