fix(kms): cross-environment usage fails when trustAccountIdentities is set (#12925)

When the `trustAccountIdentities` flag is set (either directly, or set by
default via the `@aws-cdk/aws-kms:defaultKeyPolicies` feature flag), any grant
statements are always added to the principal OR resource, defaulting to the
principal if possible.

In cross-environment usage, this means that the principal IAM policy is updated,
but the key policy is never updated. However, cross-environment usage always
requires the key policy to be updated. Attempting to use the key in a
cross-environment usage with `trustAccountIdentities` enabled initially presents
itself as a cross-environment reference issue, but even if the stack was
successfully deployed, the delegated account would not have access to use/admin
the key.

This change updates the logic to make use of the existing cross-environment
logic whenever cross-environment usage is detected, regardless of the value of
`trustAccountIdentities`. This has the impact of fixing the cross-env references
and ensuring the key policy is properly updated.

fixes #12921, fixes #12741


----

*By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*

Author: njlynch

packages/@aws-cdk/aws-kms/lib/key.ts

@@ -156,7 +156,7 @@ abstract class KeyBase extends Resource implements IKey {
       resourceArns: [this.keyArn],
       resourceSelfArns: crossEnvironment ? undefined : ['*'],
     };
-    if (this.trustAccountIdentities) {
+    if (this.trustAccountIdentities && !crossEnvironment) {
       return iam.Grant.addToPrincipalOrResource(grantOptions);
     } else {
       return iam.Grant.addToPrincipalAndResource({

packages/@aws-cdk/aws-kms/test/key.test.ts

@@ -245,6 +245,103 @@ describe('key policies', () => {
     });
   });
 
+  testFutureBehavior('grant for a principal in a different region', flags, cdk.App, (app) => {
+    const principalStack = new cdk.Stack(app, 'PrincipalStack', { env: { region: 'testregion1' } });
+    const principal = new iam.Role(principalStack, 'Role', {
+      assumedBy: new iam.AnyPrincipal(),
+      roleName: 'MyRolePhysicalName',
+    });
+
+    const keyStack = new cdk.Stack(app, 'KeyStack', { env: { region: 'testregion2' } });
+    const key = new kms.Key(keyStack, 'Key');
+
+    key.grantEncrypt(principal);
+
+    expect(keyStack).toHaveResourceLike('AWS::KMS::Key', {
+      KeyPolicy: {
+        Statement: arrayWith(
+          {
+            Action: [
+              'kms:Encrypt',
+              'kms:ReEncrypt*',
+              'kms:GenerateDataKey*',
+            ],
+            Effect: 'Allow',
+            Principal: { AWS: { 'Fn::Join': ['', ['arn:', { Ref: 'AWS::Partition' }, ':iam::', { Ref: 'AWS::AccountId' }, ':role/MyRolePhysicalName']] } },
+            Resource: '*',
+          },
+        ),
+        Version: '2012-10-17',
+      },
+    });
+    expect(principalStack).toHaveResourceLike('AWS::IAM::Policy', {
+      PolicyDocument: {
+        Statement: [
+          {
+            Action: [
+              'kms:Encrypt',
+              'kms:ReEncrypt*',
+              'kms:GenerateDataKey*',
+            ],
+            Effect: 'Allow',
+            Resource: '*',
+          },
+        ],
+        Version: '2012-10-17',
+      },
+    });
+  });
+
+  testFutureBehavior('grant for a principal in a different account', flags, cdk.App, (app) => {
+    const principalStack = new cdk.Stack(app, 'PrincipalStack', { env: { account: '0123456789012' } });
+    const principal = new iam.Role(principalStack, 'Role', {
+      assumedBy: new iam.AnyPrincipal(),
+      roleName: 'MyRolePhysicalName',
+    });
+
+    const keyStack = new cdk.Stack(app, 'KeyStack', { env: { account: '111111111111' } });
+    const key = new kms.Key(keyStack, 'Key');
+
+    key.grantEncrypt(principal);
+
+    expect(keyStack).toHaveResourceLike('AWS::KMS::Key', {
+      KeyPolicy: {
+        Statement: [
+          {
+            // Default policy, unmodified
+          },
+          {
+            Action: [
+              'kms:Encrypt',
+              'kms:ReEncrypt*',
+              'kms:GenerateDataKey*',
+            ],
+            Effect: 'Allow',
+            Principal: { AWS: { 'Fn::Join': ['', ['arn:', { Ref: 'AWS::Partition' }, ':iam::0123456789012:role/MyRolePhysicalName']] } },
+            Resource: '*',
+          },
+        ],
+        Version: '2012-10-17',
+      },
+    });
+    expect(principalStack).toHaveResourceLike('AWS::IAM::Policy', {
+      PolicyDocument: {
+        Statement: [
+          {
+            Action: [
+              'kms:Encrypt',
+              'kms:ReEncrypt*',
+              'kms:GenerateDataKey*',
+            ],
+            Effect: 'Allow',
+            Resource: '*',
+          },
+        ],
+        Version: '2012-10-17',
+      },
+    });
+  });
+
   testFutureBehavior('additional key admins can be specified (with imported/immutable principal)', flags, cdk.App, (app) => {
     const stack = new cdk.Stack(app);
     const adminRole = iam.Role.fromRoleArn(stack, 'Admin', 'arn:aws:iam::123456789012:role/TrustedAdmin');