From 551c489b6aa59a31b8533951ca56b958da805a62 Mon Sep 17 00:00:00 2001
From: hemige <hemige@firmwide.corp.gs.com>
Date: Tue, 19 Nov 2024 20:10:05 +0000
Subject: [PATCH] feat(custom-resource): support security group

---
 .../integ.aws-custom-resource-vpc.ts          |  6 ++++
 .../aws-cdk-lib/custom-resources/README.md    |  4 ++-
 .../aws-custom-resource.ts                    |  7 ++++
 .../aws-custom-resource.test.ts               | 33 +++++++++++++++++++
 4 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/packages/@aws-cdk-testing/framework-integ/test/custom-resources/test/aws-custom-resource/integ.aws-custom-resource-vpc.ts b/packages/@aws-cdk-testing/framework-integ/test/custom-resources/test/aws-custom-resource/integ.aws-custom-resource-vpc.ts
index ff979dae61cfb..9882091d8d4f7 100644
--- a/packages/@aws-cdk-testing/framework-integ/test/custom-resources/test/aws-custom-resource/integ.aws-custom-resource-vpc.ts
+++ b/packages/@aws-cdk-testing/framework-integ/test/custom-resources/test/aws-custom-resource/integ.aws-custom-resource-vpc.ts
@@ -14,6 +14,11 @@ import { AwsCustomResource, AwsCustomResourcePolicy, PhysicalResourceId } from '
 const app = new cdk.App();
 const stack = new cdk.Stack(app, 'aws-cdk-customresources-vpc');
 const vpc = new ec2.Vpc(stack, 'Vpc', { restrictDefaultSecurityGroup: false });
+const securityGroup = new ec2.SecurityGroup(stack, 'SecurityGroup', {
+  vpc,
+  securityGroupName: 'custom security group',
+  disableInlineRules: true,
+});
 new AwsCustomResource(stack, 'DescribeVpcAttribute', {
   onUpdate: {
     service: 'EC2',
@@ -28,6 +33,7 @@ new AwsCustomResource(stack, 'DescribeVpcAttribute', {
   timeout: cdk.Duration.minutes(3),
   vpc: vpc,
   vpcSubnets: { subnetType: ec2.SubnetType.PRIVATE_WITH_EGRESS },
+  securityGroups: [securityGroup],
 });
 
 new IntegTest(app, 'CustomResourceVpc', {
diff --git a/packages/aws-cdk-lib/custom-resources/README.md b/packages/aws-cdk-lib/custom-resources/README.md
index 21a9efda9d22e..0130282d6a353 100644
--- a/packages/aws-cdk-lib/custom-resources/README.md
+++ b/packages/aws-cdk-lib/custom-resources/README.md
@@ -584,11 +584,12 @@ In both the cases, you will get a synth time error if you attempt to use it in c
 
 ### Customizing the Lambda function implementing the custom resource
 
-Use the `role`, `timeout`, `memorySize`, `logGroup`, `functionName` and `removalPolicy` properties to customize
+Use the `role`, `timeout`, `memorySize`, `logGroup`, `functionName`, `securityGroups` and `removalPolicy` properties to customize
 the Lambda function implementing the custom resource:
 
 ```ts
 declare const myRole: iam.Role;
+declare const sg: ec2.SecurityGroup
 new cr.AwsCustomResource(this, 'Customized', {
   role: myRole, // must be assumable by the `lambda.amazonaws.com` service principal
   timeout: Duration.minutes(10), // defaults to 2 minutes
@@ -601,6 +602,7 @@ new cr.AwsCustomResource(this, 'Customized', {
   policy: cr.AwsCustomResourcePolicy.fromSdkCalls({
     resources: cr.AwsCustomResourcePolicy.ANY_RESOURCE,
   }),
+  securityGroups: [sg], 
 });
 ```
 
diff --git a/packages/aws-cdk-lib/custom-resources/lib/aws-custom-resource/aws-custom-resource.ts b/packages/aws-cdk-lib/custom-resources/lib/aws-custom-resource/aws-custom-resource.ts
index 9000a3fb1d70f..f468ad833ecf0 100644
--- a/packages/aws-cdk-lib/custom-resources/lib/aws-custom-resource/aws-custom-resource.ts
+++ b/packages/aws-cdk-lib/custom-resources/lib/aws-custom-resource/aws-custom-resource.ts
@@ -417,6 +417,13 @@ export interface AwsCustomResourceProps {
    * @default - the Vpc default strategy if not specified
    */
   readonly vpcSubnets?: ec2.SubnetSelection;
+
+  /**
+   * A list of IDs of security groups that the lambda function should use
+   *
+   * @default - a new security group will be created in the specified VPC
+   */
+  readonly securityGroups?: ec2.ISecurityGroup[];
 }
 
 /**
diff --git a/packages/aws-cdk-lib/custom-resources/test/aws-custom-resource/aws-custom-resource.test.ts b/packages/aws-cdk-lib/custom-resources/test/aws-custom-resource/aws-custom-resource.test.ts
index a000ffed675f7..20634cfe5408b 100644
--- a/packages/aws-cdk-lib/custom-resources/test/aws-custom-resource/aws-custom-resource.test.ts
+++ b/packages/aws-cdk-lib/custom-resources/test/aws-custom-resource/aws-custom-resource.test.ts
@@ -1207,6 +1207,39 @@ test('can specify VPC', () => {
   });
 });
 
+test('can specify security group', () => {
+  // GIVEN
+  const stack = new cdk.Stack();
+  const vpc = new ec2.Vpc(stack, 'TestVpc');
+  const securityGroups = [
+    new ec2.SecurityGroup(stack, 'Sg1', {
+      vpc: vpc,
+      allowAllOutbound: false,
+      description: 'my security group',
+    }),
+  ];
+
+  // WHEN
+  new AwsCustomResource(stack, 'AwsSdk', {
+    onCreate: {
+      service: 'service',
+      action: 'action',
+      physicalResourceId: PhysicalResourceId.of('id'),
+    },
+    policy: AwsCustomResourcePolicy.fromSdkCalls({ resources: AwsCustomResourcePolicy.ANY_RESOURCE }),
+    vpc,
+    vpcSubnets: { subnetType: ec2.SubnetType.PRIVATE_WITH_EGRESS },
+    securityGroups,
+  });
+
+  // THEN
+  Template.fromStack(stack).hasResourceProperties('AWS::Lambda::Function', {
+    VpcConfig: {
+      SecurityGroupIds: stack.resolve(securityGroups.map(sg => sg.securityGroupId)),
+    },
+  });
+});
+
 test('specifying public subnets results in a synthesis error', () => {
   // GIVEN
   const stack = new cdk.Stack();