fix(kms): correctly recognize newly created resources

packages/@aws-cdk/aws-iam/lib/private/immutable-role.ts

@@ -39,6 +39,7 @@ export class ImmutableRole extends Resource implements IRole {
     Dependable.implement(this, {
       dependencyRoots: [role],
     });
+    this.node.defaultChild = role.node.defaultChild;
   }
 
   public attachInlinePolicy(_policy: Policy): void {

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

@@ -2,7 +2,7 @@ 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 * as cxapi from '@aws-cdk/cx-api';
-import { IConstruct, Construct } from 'constructs';
+import { Construct } from 'constructs';
 import { Alias } from './alias';
 import { KeyLookupOptions } from './key-lookup';
 import { CfnKey } from './kms.generated';
@@ -208,28 +208,20 @@ abstract class KeyBase extends Resource implements IKey {
     }
     // this logic should only apply to newly created
     // (= not imported) resources
-    if (!this.principalIsANewlyCreatedResource(grantPrincipal)) {
+    if (!Resource.isOwnedResource(grantPrincipal)) {
       return undefined;
     }
-    // return undefined;
     const keyStack = Stack.of(this);
     const granteeStack = Stack.of(grantPrincipal);
     if (keyStack === granteeStack) {
       return undefined;
     }
+
     return granteeStack.dependencies.includes(keyStack)
       ? granteeStack.account
       : undefined;
   }
 
-  private principalIsANewlyCreatedResource(principal: IConstruct): boolean {
-    // yes, this sucks
-    // this is just a temporary stopgap to stem the bleeding while we work on a proper fix
-    return principal instanceof iam.Role ||
-      principal instanceof iam.User ||
-      principal instanceof iam.Group;
-  }
-
   private isGranteeFromAnotherRegion(grantee: iam.IGrantable): boolean {
     if (!isConstruct(grantee)) {
       return false;

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

@@ -337,6 +337,41 @@ describe('key policies', () => {
     });
   });
 
+  testFutureBehavior('grant for an immutable role', 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');
+    principalStack.addDependency(keyStack);
+    key.grantEncrypt(principal.withoutPolicyUpdates());
+
+    Template.fromStack(keyStack).hasResourceProperties('AWS::KMS::Key', {
+      KeyPolicy: {
+        Statement: Match.arrayWith([{
+          Action: 'kms:*',
+          Effect: 'Allow',
+          Principal: { AWS: { 'Fn::Join': ['', ['arn:', { Ref: 'AWS::Partition' }, ':iam::111111111111:root']] } },
+          Resource: '*',
+        },
+        {
+          Action: [
+            'kms:Encrypt',
+            'kms:ReEncrypt*',
+            'kms:GenerateDataKey*',
+          ],
+          Effect: 'Allow',
+          Principal: { AWS: { 'Fn::Join': ['', ['arn:', { Ref: 'AWS::Partition' }, ':iam::0123456789012:root']] } },
+          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');

packages/@aws-cdk/core/lib/resource.ts

@@ -123,10 +123,17 @@ export abstract class Resource extends Construct implements IResource {
   /**
    * Check whether the given construct is a Resource
    */
-  public static isResource(construct: IConstruct): construct is CfnResource {
+  public static isResource(construct: IConstruct): construct is Resource {
     return construct !== null && typeof(construct) === 'object' && RESOURCE_SYMBOL in construct;
   }
 
+  /**
+   * Returns true if the construct was created by CDK, and false otherwise
+   */
+  public static isOwnedResource(construct: IConstruct): boolean {
+    return construct.node.defaultChild ? CfnResource.isCfnResource(construct.node.defaultChild) : false;
+  }
+
   public readonly stack: Stack;
   public readonly env: ResourceEnvironment;