From b107b2880284abfda17aee3f572432e1565a9356 Mon Sep 17 00:00:00 2001
From: Piradeep Kandasamy <piradeep@amazon.com>
Date: Tue, 11 Jun 2019 02:18:19 -0700
Subject: [PATCH] fix(aws-ecs): downscope instance drain hook permissions for
 ASG and ECS Cluster actions

---
 .../aws-autoscaling/lib/auto-scaling-group.ts | 24 ++++++++++++++++-
 packages/@aws-cdk/aws-ecs/lib/cluster.ts      |  8 +-----
 .../lib/drain-hook/instance-drain-hook.ts     | 26 ++++++++++++++-----
 3 files changed, 44 insertions(+), 14 deletions(-)

diff --git a/packages/@aws-cdk/aws-autoscaling/lib/auto-scaling-group.ts b/packages/@aws-cdk/aws-autoscaling/lib/auto-scaling-group.ts
index 671ed8d8a6208..4982214e9370e 100644
--- a/packages/@aws-cdk/aws-autoscaling/lib/auto-scaling-group.ts
+++ b/packages/@aws-cdk/aws-autoscaling/lib/auto-scaling-group.ts
@@ -5,7 +5,7 @@ import elbv2 = require('@aws-cdk/aws-elasticloadbalancingv2');
 import iam = require('@aws-cdk/aws-iam');
 import sns = require('@aws-cdk/aws-sns');
 
-import { AutoScalingRollingUpdate, Construct, Fn, IResource, Resource, Tag, Token } from '@aws-cdk/cdk';
+import { AutoScalingRollingUpdate, Construct, Fn, IResource, Resource, Stack, Tag, Token } from '@aws-cdk/cdk';
 import { CfnAutoScalingGroup, CfnAutoScalingGroupProps, CfnLaunchConfiguration } from './autoscaling.generated';
 import { BasicLifecycleHookProps, LifecycleHook } from './lifecycle-hook';
 import { BasicScheduledActionProps, ScheduledAction } from './scheduled-action';
@@ -196,6 +196,7 @@ export interface AutoScalingGroupProps extends CommonAutoScalingGroupProps {
 abstract class AutoScalingGroupBase extends Resource implements IAutoScalingGroup {
 
   public abstract autoScalingGroupName: string;
+  public abstract autoScalingGroupArn: string;
   protected albTargetGroup?: elbv2.ApplicationTargetGroup;
 
   /**
@@ -318,6 +319,11 @@ export class AutoScalingGroup extends AutoScalingGroupBase implements
   public static fromAutoScalingGroupName(scope: Construct, id: string, autoScalingGroupName: string): IAutoScalingGroup {
     class Import extends AutoScalingGroupBase {
       public autoScalingGroupName = autoScalingGroupName;
+      public autoScalingGroupArn = Stack.of(this).formatArn({
+        service: 'autoscaling',
+        resource: 'autoScalingGroup:*:autoScalingGroupName',
+        resourceName: this.autoScalingGroupName
+      });
     }
 
     return new Import(scope, id);
@@ -343,6 +349,11 @@ export class AutoScalingGroup extends AutoScalingGroupBase implements
    */
   public readonly autoScalingGroupName: string;
 
+  /**
+   * Arn of the AutoScalingGroup
+   */
+  public readonly autoScalingGroupArn: string;
+
   private readonly userDataLines = new Array<string>();
   private readonly autoScalingGroup: CfnAutoScalingGroup;
   private readonly securityGroup: ec2.ISecurityGroup;
@@ -432,6 +443,11 @@ export class AutoScalingGroup extends AutoScalingGroupBase implements
     this.autoScalingGroup = new CfnAutoScalingGroup(this, 'ASG', asgProps);
     this.osType = machineImage.os.type;
     this.autoScalingGroupName = this.autoScalingGroup.autoScalingGroupName;
+    this.autoScalingGroupArn = Stack.of(this).formatArn({
+      service: 'autoscaling',
+      resource: 'autoScalingGroup:*:autoScalingGroupName',
+      resourceName: this.autoScalingGroupName
+    });
 
     this.applyUpdatePolicies(props);
   }
@@ -707,6 +723,12 @@ export interface IAutoScalingGroup extends IResource {
    */
   readonly autoScalingGroupName: string;
 
+  /**
+   * The arn of the AutoScalingGroup
+   * @attribute
+   */
+  readonly autoScalingGroupArn: string;
+
   /**
    * Send a message to either an SQS queue or SNS topic when instances launch or terminate
    */
diff --git a/packages/@aws-cdk/aws-ecs/lib/cluster.ts b/packages/@aws-cdk/aws-ecs/lib/cluster.ts
index cf9812bbe887a..eb2c78062a24d 100644
--- a/packages/@aws-cdk/aws-ecs/lib/cluster.ts
+++ b/packages/@aws-cdk/aws-ecs/lib/cluster.ts
@@ -401,13 +401,7 @@ class ImportedCluster extends Resource implements ICluster {
     this.clusterArn = props.clusterArn !== undefined ? props.clusterArn : Stack.of(this).formatArn({
       service: 'ecs',
       resource: 'cluster',
-      resourceName: props.clusterName,
-    });
-
-    this.clusterArn = props.clusterArn !== undefined ? props.clusterArn : Stack.of(this).formatArn({
-      service: 'ecs',
-      resource: 'cluster',
-      resourceName: props.clusterName,
+      resourceName: props.clusterName
     });
 
     let i = 1;
diff --git a/packages/@aws-cdk/aws-ecs/lib/drain-hook/instance-drain-hook.ts b/packages/@aws-cdk/aws-ecs/lib/drain-hook/instance-drain-hook.ts
index 87fd362a0eccb..87ca922817a5e 100644
--- a/packages/@aws-cdk/aws-ecs/lib/drain-hook/instance-drain-hook.ts
+++ b/packages/@aws-cdk/aws-ecs/lib/drain-hook/instance-drain-hook.ts
@@ -69,20 +69,24 @@ export class InstanceDrainHook extends cdk.Construct {
       heartbeatTimeoutSec: drainTimeSeconds,
     });
 
-    // FIXME: These should probably be restricted usefully in some way, but I don't exactly
-    // know how.
+    // Describe actions cannot be restricted and restrict the CompleteLifecycleAction to the ASG arn
+    // https://docs.aws.amazon.com/autoscaling/ec2/userguide/control-access-using-iam.html
     fn.addToRolePolicy(new iam.PolicyStatement()
       .addActions(
-        'autoscaling:CompleteLifecycleAction',
         'ec2:DescribeInstances',
         'ec2:DescribeInstanceAttribute',
         'ec2:DescribeInstanceStatus',
-        'ec2:DescribeHosts',
+        'ec2:DescribeHosts'
       )
       .addAllResources());
 
-    // FIXME: These should be restricted to the ECS cluster probably, but I don't exactly
-    // know how.
+    // Restrict to the ASG
+    fn.addToRolePolicy(new iam.PolicyStatement()
+      .addActions(
+        'autoscaling:CompleteLifecycleAction'
+      )
+      .addResource(props.autoScalingGroup.autoScalingGroupArn));
+
     fn.addToRolePolicy(new iam.PolicyStatement()
       .addActions(
         'ecs:ListContainerInstances',
@@ -93,5 +97,15 @@ export class InstanceDrainHook extends cdk.Construct {
         'ecs:ListTasks',
         'ecs:DescribeTasks')
       .addAllResources());
+
+    // Restrict to the ECS Cluster
+    fn.addToRolePolicy(new iam.PolicyStatement()
+      .addActions(
+        'ecs:ListContainerInstances',
+        'ecs:SubmitContainerStateChange',
+        'ecs:SubmitTaskStateChange',
+        'ecs:UpdateContainerInstancesState',
+        'ecs:ListTasks')
+      .addResource(props.cluster.clusterArn));
   }
 }