From 67ad9fe24a1e03b2c2884c16bedb2447d69a7c3a Mon Sep 17 00:00:00 2001
From: hemige <hemige@firmwide.corp.gs.com>
Date: Sat, 16 Nov 2024 05:24:59 +0000
Subject: [PATCH] feat(custom-resource): support security group

---
 .../aws-custom-resource.ts                    |  8 +++++
 .../aws-custom-resource.test.ts               | 33 +++++++++++++++++++
 2 files changed, 41 insertions(+)

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..c25de916ca7fc 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[];
 }
 
 /**
@@ -500,6 +507,7 @@ export class AwsCustomResource extends Construct implements iam.IGrantable {
       functionName: props.functionName,
       vpc: props.vpc,
       vpcSubnets: props.vpcSubnets,
+      securityGroups: props.securityGroups,
     });
     this.grantPrincipal = provider.grantPrincipal;
 
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..c0653bd67c239 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
@@ -1317,6 +1317,39 @@ test('vpcSubnets without vpc results in an error', () => {
   })).toThrow('Cannot configure \'vpcSubnets\' without configuring a VPC');
 });
 
+test('can specify subnets', () => {
+  // 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.each([
   [undefined, true],
   [true, true],