From 589e2d82362870fd2c58033f388f7f4375e1510b Mon Sep 17 00:00:00 2001
From: Marcelo Luiz Onhate <onhate@gmail.com>
Date: Mon, 11 Mar 2024 17:19:39 -0300
Subject: [PATCH] feat: use LoggingConfig.LogGroupName when lambda is
 configured with custom log groups

---
 .../api/evaluate-cloudformation-template.ts   |  4 ++
 .../lib/api/logs/find-cloudwatch-logs.ts      | 71 +++++++++++--------
 2 files changed, 46 insertions(+), 29 deletions(-)

diff --git a/packages/aws-cdk/lib/api/evaluate-cloudformation-template.ts b/packages/aws-cdk/lib/api/evaluate-cloudformation-template.ts
index 329bf2ec9ba59..a8e255a616f74 100644
--- a/packages/aws-cdk/lib/api/evaluate-cloudformation-template.ts
+++ b/packages/aws-cdk/lib/api/evaluate-cloudformation-template.ts
@@ -314,6 +314,10 @@ export class EvaluateCloudFormationTemplate {
     return cfnExpression;
   }
 
+  public getResourceProperty(logicalId: string, propertyName: string) {
+    return this.template.Resources?.[logicalId]?.Properties?.[propertyName];
+  }
+
   private references(logicalId: string, templateElement: any): boolean {
     if (typeof templateElement === 'string') {
       return logicalId === templateElement;
diff --git a/packages/aws-cdk/lib/api/logs/find-cloudwatch-logs.ts b/packages/aws-cdk/lib/api/logs/find-cloudwatch-logs.ts
index bfff256221a52..dd98c0077a547 100644
--- a/packages/aws-cdk/lib/api/logs/find-cloudwatch-logs.ts
+++ b/packages/aws-cdk/lib/api/logs/find-cloudwatch-logs.ts
@@ -1,22 +1,12 @@
 import * as cxapi from '@aws-cdk/cx-api';
 import { CloudFormation } from 'aws-sdk';
-import { Mode, SdkProvider, ISDK } from '../aws-auth';
+import { ISDK, Mode, SdkProvider } from '../aws-auth';
 import { Deployments } from '../deployments';
 import { EvaluateCloudFormationTemplate, LazyListStackResources } from '../evaluate-cloudformation-template';
 
 // resource types that have associated CloudWatch Log Groups that should _not_ be monitored
 const IGNORE_LOGS_RESOURCE_TYPES = ['AWS::EC2::FlowLog', 'AWS::CloudTrail::Trail', 'AWS::CodeBuild::Project'];
 
-// Resource types that will create a CloudWatch log group with a specific name if one is not provided.
-// The keys are CFN resource types, and the values are the name of the physical name property of that resource
-// and the service name that is used in the automatically created CloudWatch log group.
-const RESOURCE_TYPES_WITH_IMPLICIT_LOGS: { [cfnResourceType: string]: { [key: string]: string } } = {
-  'AWS::Lambda::Function': {
-    PhysicalNamePropertyName: 'FunctionName',
-    LogGroupServiceName: 'lambda',
-  },
-};
-
 /**
  * Configuration needed to monitor CloudWatch Log Groups
  * found in a given CloudFormation Stack
@@ -84,38 +74,61 @@ function isReferencedFromIgnoredResource(
   logGroupResource: CloudFormation.StackResourceSummary,
   evaluateCfnTemplate: EvaluateCloudFormationTemplate,
 ): boolean {
-  let foundReference = false;
   const resourcesReferencingLogGroup = evaluateCfnTemplate.findReferencesTo(logGroupResource.LogicalResourceId);
   for (const reference of resourcesReferencingLogGroup) {
     if (IGNORE_LOGS_RESOURCE_TYPES.includes(reference.Type)) {
-      foundReference = true;
+      return true;
     }
   }
-  return foundReference;
+  return false;
 }
 
+type CloudWatchLogsResolver = (
+  resource: CloudFormation.StackResourceSummary,
+  evaluateCfnTemplate: EvaluateCloudFormationTemplate
+) => string | undefined;
+
+const cloudWatchLogsResolvers: Record<string, CloudWatchLogsResolver> = {
+  'AWS::Logs::LogGroup': (resource, evaluateCfnTemplate) => {
+    if (isReferencedFromIgnoredResource(resource, evaluateCfnTemplate)) {
+      return undefined;
+    }
+    return resource.PhysicalResourceId?.toString();
+  },
+
+  // Resource types that will create a CloudWatch log group with a specific name if one is not provided.
+  // The keys are CFN resource types, and the values are the name of the physical name property of that resource
+  // and the service name that is used in the automatically created CloudWatch log group.
+  'AWS::Lambda::Function': (resource, evaluateCfnTemplate) => {
+    const loggingConfig = evaluateCfnTemplate.getResourceProperty(resource.LogicalResourceId, 'LoggingConfig');
+    if (loggingConfig?.LogGroupName) {
+      return loggingConfig.LogGroupName;
+    }
+
+    return `/aws/lambda/${resource.PhysicalResourceId}`;
+  },
+};
+
 /**
  * Find all CloudWatch Log Groups in the deployed template.
- * This will find both explicitely created Log Groups (excluding those associated with ignored resources)
- * as well as Log Groups created implicitely (i.e. Lambda Functions)
+ * This will find both explicitly created Log Groups (excluding those associated with ignored resources)
+ * and Log Groups created implicitly (i.e. Lambda Functions)
  */
 function findAllLogGroupNames(
   stackResources: CloudFormation.StackResourceSummary[],
   evaluateCfnTemplate: EvaluateCloudFormationTemplate,
 ): string[] {
-  return stackResources.reduce((logGroupNames: string[], resource) => {
-    let logGroupName;
-    if (resource.ResourceType === 'AWS::Logs::LogGroup') {
-      if (!isReferencedFromIgnoredResource(resource, evaluateCfnTemplate)) {
-        logGroupName = resource.PhysicalResourceId;
+  const logGroupNames: string[] = [];
+
+  for (const resource of stackResources) {
+    const logGroupResolver = cloudWatchLogsResolvers[resource.ResourceType];
+    if (logGroupResolver) {
+      const logGroupName = logGroupResolver(resource, evaluateCfnTemplate);
+      if (logGroupName) {
+        logGroupNames.push(logGroupName);
       }
-    } else if (RESOURCE_TYPES_WITH_IMPLICIT_LOGS[resource.ResourceType]) {
-      const servicePart = RESOURCE_TYPES_WITH_IMPLICIT_LOGS[resource.ResourceType].LogGroupServiceName;
-      logGroupName = `/aws/${servicePart}/${resource.PhysicalResourceId}`;
-    }
-    if (logGroupName) {
-      logGroupNames.push(logGroupName);
     }
-    return logGroupNames;
-  }, []);
+  }
+
+  return logGroupNames;
 }