From c6dbb96f3fcf89a247bf7e7271f3c5b283563144 Mon Sep 17 00:00:00 2001
From: Kaizen Conroy <36202692+kaizencc@users.noreply.github.com>
Date: Tue, 9 Aug 2022 22:58:09 -0400
Subject: [PATCH] fix(kms): imported key ignores environment from arn (#21519)

Fixes #21464. KMS keys imported using `fromKeyArn()` currently take the environment of the stack, not the environment from the arn.

This PR follows the precedent set in https://github.com/aws/aws-cdk/pull/19026 and https://github.com/aws/aws-cdk/pull/18255. It is essentially the same code change and tests. Ideally, we would have a mechanism for testing all `fromXxxArn` APIs to ensure they have the correct behavior. There are still many places where it does not. However, given the significant overhead of creating such a mechanism, I'm creating this one-off PR to unblock users in KMS.

----

### All Submissions:

* [ ] Have you followed the guidelines in our [Contributing guide?](https://github.com/aws/aws-cdk/blob/main/CONTRIBUTING.md)

### Adding new Unconventional Dependencies:

* [ ] This PR adds new unconventional dependencies following the process described [here](https://github.com/aws/aws-cdk/blob/main/CONTRIBUTING.md/#adding-new-unconventional-dependencies)

### New Features

* [ ] Have you added the new feature to an [integration test](https://github.com/aws/aws-cdk/blob/main/INTEGRATION_TESTS.md)?
	* [ ] Did you use `yarn integ` to deploy the infrastructure and generate the snapshot (i.e. `yarn integ` without `--dry-run`)?

*By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*
---
 packages/@aws-cdk/aws-kms/lib/key.ts       | 14 +++++-----
 packages/@aws-cdk/aws-kms/test/key.test.ts | 31 ++++++++++++++++++++++
 2 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/packages/@aws-cdk/aws-kms/lib/key.ts b/packages/@aws-cdk/aws-kms/lib/key.ts
index db334374593ed..0dc523fb73c37 100644
--- a/packages/@aws-cdk/aws-kms/lib/key.ts
+++ b/packages/@aws-cdk/aws-kms/lib/key.ts
@@ -1,6 +1,6 @@
 import * as iam from '@aws-cdk/aws-iam';
 import * as cxschema from '@aws-cdk/cloud-assembly-schema';
-import { FeatureFlags, IResource, Lazy, RemovalPolicy, Resource, Stack, Duration, Token, ContextProvider, Arn, ArnFormat } from '@aws-cdk/core';
+import { FeatureFlags, IResource, Lazy, RemovalPolicy, Resource, ResourceProps, Stack, Duration, Token, ContextProvider, Arn, ArnFormat } from '@aws-cdk/core';
 import * as cxapi from '@aws-cdk/cx-api';
 import { Construct } from 'constructs';
 import { Alias } from './alias';
@@ -94,8 +94,8 @@ abstract class KeyBase extends Resource implements IKey {
    */
   private readonly aliases: Alias[] = [];
 
-  constructor(scope: Construct, id: string) {
-    super(scope, id);
+  constructor(scope: Construct, id: string, props: ResourceProps = {}) {
+    super(scope, id, props);
 
     this.node.addValidation({ validate: () => this.policy?.validateForResourcePolicy() ?? [] });
   }
@@ -464,8 +464,8 @@ export class Key extends KeyBase {
       // policies is really the only option
       protected readonly trustAccountIdentities: boolean = true;
 
-      constructor(keyId: string) {
-        super(scope, id);
+      constructor(keyId: string, props: ResourceProps = {}) {
+        super(scope, id, props);
 
         this.keyId = keyId;
       }
@@ -476,7 +476,9 @@ export class Key extends KeyBase {
       throw new Error(`KMS key ARN must be in the format 'arn:aws:kms:<region>:<account>:key/<keyId>', got: '${keyArn}'`);
     }
 
-    return new Import(keyResourceName);
+    return new Import(keyResourceName, {
+      environmentFromArn: keyArn,
+    });
   }
 
   /**
diff --git a/packages/@aws-cdk/aws-kms/test/key.test.ts b/packages/@aws-cdk/aws-kms/test/key.test.ts
index 43d6fbdf82aaa..5c4c2adee490d 100644
--- a/packages/@aws-cdk/aws-kms/test/key.test.ts
+++ b/packages/@aws-cdk/aws-kms/test/key.test.ts
@@ -1240,3 +1240,34 @@ describe('key specs and key usages', () => {
       .toThrow('key rotation cannot be enabled on asymmetric keys');
   });
 });
+
+describe('Key.fromKeyArn()', () => {
+  let stack: cdk.Stack;
+
+  beforeEach(() => {
+    const app = new cdk.App();
+    stack = new cdk.Stack(app, 'Base', {
+      env: { account: '111111111111', region: 'stack-region' },
+    });
+  });
+
+  describe('for a key in a different account and region', () => {
+    let key: kms.IKey;
+
+    beforeEach(() => {
+      key = kms.Key.fromKeyArn(
+        stack,
+        'iKey',
+        'arn:aws:kms:key-region:222222222222:key:key-name',
+      );
+    });
+
+    test("the key's region is taken from the ARN", () => {
+      expect(key.env.region).toBe('key-region');
+    });
+
+    test("the key's account is taken from the ARN", () => {
+      expect(key.env.account).toBe('222222222222');
+    });
+  });
+});