Repository: cdklabs/cdk-enterprise-iac (language: python)

Issue:
bug: Resource Extractor not correctly processing Fn::Join for extracted resources
### What is the problem?

When using the resource extractor, if you try to dynamically build a string that would be used, the resource extractor injects a Dictionary/Object with Numbers instead of accurate CloudFormation Fn::Join statements.

### Reproduction Steps

Dynamically build a string and pass it to any resource, like this:

```python
iam.ArnPrincipal(
                "arn:{}:sts::{}:assumed-role/{}/{}".format(
                    Aws.PARTITION,
                    Aws.ACCOUNT_ID,
                    lambda.role.role_name,
                    lambda.function_name,
                )
            )
)
```




### What did you expect to happen?

A normal Fn::Join statement is generated , which looks like

```json
{
  "Fn::Join": [
    "",
    [
       "foo",
       "bar"
    ]
  ]
}
```


### What actually happened?

Incorrect cloudformation was generated with bogus object that contained numbers

```yaml
            Principal:
                AWS:
                    - 
                        'Fn::Join':
                            '1':
                                '5': 'role'
                    - 
                        'Fn::Join':
                            '1':
                                '5': 'function'
```

### cdk-enterprise-iac version

0.0.218

### Language

Typescript, Python

### Other information

_No response_


Discussion / hints:

Fix patch:
diff --git a/src/patches/resource-extractor/cfnStore.ts b/src/patches/resource-extractor/cfnStore.ts
index d37f5a7..fe5049d 100644
--- a/src/patches/resource-extractor/cfnStore.ts
+++ b/src/patches/resource-extractor/cfnStore.ts
@@ -3,7 +3,7 @@ Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 SPDX-License-Identifier: Apache-2.0
 */
 import * as cp from 'child_process';
-import { CfnResource } from 'aws-cdk-lib';
+import { CfnResource, Fn } from 'aws-cdk-lib';
 import { CloudFormationStackArtifact } from 'aws-cdk-lib/cx-api';
 import { CloudFormation } from 'aws-sdk';
 import { Flattener } from './flattener';
@@ -21,6 +21,7 @@ export class CfnStore {
   public readonly templates: Json = {};
   public readonly flatTemplates: FlatJson;
   public readonly extractedStackExports: FlatJson = {};
+  public readonly fnJoins: Json = {};
 
   constructor(props: CfnStoreProps) {
     /** Save CloudFormation templates for future lookup */
@@ -167,4 +168,41 @@ export class CfnStore {
     }
     return exports;
   }
+
+  /**
+   * Rebuilds an `Fn::Join` statement from the Flat Json path. This is
+   * necessary because CDK does not correctly inject values back into
+   * `Fn::Join` blocks. The rebuilt object is injected back into the spot where
+   * the original `Fn::Join` was at. By using this method, we can get around
+   * the problem where `addPropertyOverride` does not work for Fn::Join lists.
+   *
+   * @param fnJoinFlatJsonPath
+   * @returns string of Fn::Join (may be string or Json string depending on
+   * if CloudFormation needs to reference values or if the entire string can
+   * be hardcoded back)
+   */
+  public rebuildFnJoin(fnJoinFlatJsonPath: string): string {
+    // Find the items that this fnJoin references
+    const items = Object.keys(this.flatTemplates).filter((key) =>
+      key.includes(`${fnJoinFlatJsonPath}.1`)
+    );
+
+    // Sort the items so that they are rebuilt in the right order
+    items.sort();
+
+    // Get the actual values from the flat templates map
+    const listOfValues: string[] = [];
+    items.forEach((item) => {
+      if (Object.keys(this.fnJoins).includes(item)) {
+        listOfValues.push(this.fnJoins[item]);
+      } else if (item.split('.').slice(-1)[0] == 'Ref') {
+        listOfValues.push(Fn.ref(this.flatTemplates[item]));
+      } else if (!item.includes('Fn::GetAtt.1')) {
+        listOfValues.push(this.flatTemplates[item]);
+      }
+    });
+
+    // Return the rebuilt Join statement
+    return Fn.join(this.flatTemplates[`${fnJoinFlatJsonPath}.0`], listOfValues);
+  }
 }
diff --git a/src/patches/resource-extractor/resourceExtractor.ts b/src/patches/resource-extractor/resourceExtractor.ts
index 90d0da3..8354594 100644
--- a/src/patches/resource-extractor/resourceExtractor.ts
+++ b/src/patches/resource-extractor/resourceExtractor.ts
@@ -317,6 +317,16 @@ export class ResourceExtractor implements IAspect {
       const newValue = props.foundRefNode.stack.resolve(props.importValue);
       if (props.foundRefNode instanceof CfnOutput) {
         props.foundRefNode.value = newValue;
+      } else if (props.flattenedKey.includes('Fn::Join')) {
+        // Get Fn::Join path and save new value to fnJoins table
+        const fnJoinFlatJsonPath =
+          props.flattenedKey.split('.Fn::Join')[0] + '.Fn::Join';
+        const fnJoinOverridePath = propertyOverridePath.split('.Fn::Join')[0];
+        this.cfn.fnJoins[props.flattenedKey] = newValue;
+
+        // Rebuild Fn::Join object
+        const joined = this.cfn.rebuildFnJoin(fnJoinFlatJsonPath);
+        props.foundRefNode.addPropertyOverride(fnJoinOverridePath, joined);
       } else {
         props.foundRefNode.addPropertyOverride(propertyOverridePath, newValue);
       }
@@ -329,6 +339,16 @@ export class ResourceExtractor implements IAspect {
       );
       if (props.foundRefNode instanceof CfnOutput) {
         props.foundRefNode.value = newValue;
+      } else if (props.flattenedKey.includes('Fn::Join')) {
+        // Get Fn::Join path and save new value to fnJoins table
+        const fnJoinFlatJsonPath =
+          props.flattenedKey.split('.Fn::Join')[0] + '.Fn::Join';
+        const fnJoinOverridePath = propertyOverridePath.split('.Fn::Join')[0];
+        this.cfn.fnJoins[props.flattenedKey] = newValue;
+
+        // Rebuild Fn::Join object
+        const joined = this.cfn.rebuildFnJoin(fnJoinFlatJsonPath);
+        props.foundRefNode.addPropertyOverride(fnJoinOverridePath, joined);
       } else {
         props.foundRefNode.addPropertyOverride(propertyOverridePath, newValue);
       }
@@ -343,6 +363,16 @@ export class ResourceExtractor implements IAspect {
         `dummy-value-for-${importValue}`;
       if (props.foundRefNode instanceof CfnOutput) {
         props.foundRefNode.value = newValue;
+      } else if (props.flattenedKey.includes('Fn::Join')) {
+        // Get Fn::Join path and save new value to fnJoins table
+        const fnJoinFlatJsonPath =
+          props.flattenedKey.split('.Fn::Join')[0] + '.Fn::Join';
+        const fnJoinOverridePath = propertyOverridePath.split('.Fn::Join')[0];
+        this.cfn.fnJoins[props.flattenedKey] = newValue;
+
+        // Rebuild Fn::Join object
+        const joined = this.cfn.rebuildFnJoin(fnJoinFlatJsonPath);
+        props.foundRefNode.addPropertyOverride(fnJoinOverridePath, joined);
       } else {
         props.foundRefNode.addPropertyOverride(propertyOverridePath, newValue);
       }


Test patch:
diff --git a/test/patches/resource-extractor/resourceExtractor.test.ts b/test/patches/resource-extractor/resourceExtractor.test.ts
index 48959a5..e47c3e6 100644
--- a/test/patches/resource-extractor/resourceExtractor.test.ts
+++ b/test/patches/resource-extractor/resourceExtractor.test.ts
@@ -2,7 +2,14 @@
 Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 SPDX-License-Identifier: Apache-2.0
 */
-import { App, Aspects, CfnElement, RemovalPolicy, Stack } from 'aws-cdk-lib';
+import {
+  App,
+  Aspects,
+  Aws,
+  CfnElement,
+  RemovalPolicy,
+  Stack,
+} from 'aws-cdk-lib';
 import { Match, Template, Annotations } from 'aws-cdk-lib/assertions';
 import {
   Instance,
@@ -148,6 +155,69 @@ describe('Extracting resources from stack', () => {
     });
   });
 
+  test('Dynamically build string with exported values', () => {
+    const role = new Role(stack, 'TestRole', {
+      assumedBy: new ServicePrincipal('lambda.amazonaws.com'),
+    });
+    const testFunc = new Function(stack, 'TestFunction', {
+      code: Code.fromInline("print('hello_world')"),
+      handler: 'index.handler',
+      runtime: Runtime.PYTHON_3_9,
+      environment: {
+        JOINED: `arn:${role.roleName}:foo:${role.roleArn}:bar:${Aws.REGION}`,
+      },
+    });
+    testFunc.grantInvoke(role);
+
+    const synthedApp = app.synth();
+
+    Aspects.of(app).add(
+      new ResourceExtractor({
+        extractDestinationStack: extractedStack,
+        stackArtifacts: synthedApp.stacks,
+        valueShareMethod: ResourceExtractorShareMethod.CFN_OUTPUT,
+        resourceTypesToExtract,
+      })
+    );
+    app.synth({ force: true });
+
+    const extractedTemplate = Template.fromStack(extractedStack);
+    const appTemplate = Template.fromStack(stack);
+
+    // Extracted stack has IAM resources
+    extractedTemplate.resourceCountIs('AWS::IAM::Role', 2);
+    extractedTemplate.resourceCountIs('AWS::IAM::Policy', 1);
+    // Non-IAM resources present in app stack
+    appTemplate.resourceCountIs('AWS::Lambda::Function', 1);
+    // Lambda function has correct environment variable that has been joined
+    // using Fn::Join since we need to import values from Outputs
+    appTemplate.hasResourceProperties('AWS::Lambda::Function', {
+      Environment: {
+        Variables: {
+          JOINED: {
+            'Fn::Join': [
+              '',
+              [
+                'arn:',
+                {
+                  'Fn::ImportValue': Match.anyValue(),
+                },
+                ':foo:',
+                {
+                  'Fn::ImportValue': Match.anyValue(),
+                },
+                ':bar:',
+                {
+                  Ref: 'AWS::Region',
+                },
+              ],
+            ],
+          },
+        },
+      },
+    });
+  });
+
   test('Only appropriate DependsOn values are removed', () => {
     const func = new Function(stack, 'TestLambda', {
       code: Code.fromInline(`def handler(event, context)\n    print(event)`),
@@ -552,6 +622,52 @@ describe('Sharing Methods - SSM_PARAMETER', () => {
     });
   });
 
+  test('Dynamically build string with exported values', () => {
+    const role = new Role(stack, 'TestRole', {
+      assumedBy: new ServicePrincipal('lambda.amazonaws.com'),
+    });
+    const testFunc = new Function(stack, 'TestFunction', {
+      code: Code.fromInline("print('hello_world')"),
+      handler: 'index.handler',
+      runtime: Runtime.PYTHON_3_9,
+      environment: {
+        JOINED: `arn:${role.roleName}:foo:${role.roleArn}:bar`,
+      },
+    });
+    testFunc.grantInvoke(role);
+
+    const synthedApp = app.synth();
+    Aspects.of(app).add(
+      new ResourceExtractor({
+        extractDestinationStack: extractedStack,
+        stackArtifacts: synthedApp.stacks,
+        valueShareMethod: ResourceExtractorShareMethod.SSM_PARAMETER,
+        resourceTypesToExtract,
+      })
+    );
+    app.synth({ force: true });
+
+    const extractedTemplate = Template.fromStack(extractedStack);
+    const appTemplate = Template.fromStack(stack);
+
+    // Extracted stack has IAM resources
+    extractedTemplate.resourceCountIs('AWS::IAM::Role', 2);
+    extractedTemplate.resourceCountIs('AWS::IAM::Policy', 1);
+    // Non-IAM resources present in app stack
+    appTemplate.resourceCountIs('AWS::Lambda::Function', 1);
+    // Lambda function has correct environment variable that has been joined
+    // using the Fn.Join function, which combines the result of all strings
+    appTemplate.hasResourceProperties('AWS::Lambda::Function', {
+      Environment: {
+        Variables: {
+          JOINED: Match.stringLikeRegexp(
+            'arn:dummy-value-for-(.*?):foo:dummy-value-for-(.*?):bar'
+          ),
+        },
+      },
+    });
+  });
+
   test('Exports work correctly', () => {
     const secondStack = new Stack(app, 'Stack2', { env });
     const role = new Role(stack, 'ExportedRole', {
@@ -640,6 +756,52 @@ describe('Sharing Methods - API_LOOKUP', () => {
     });
   });
 
+  test('Dynamically build string with exported values', () => {
+    const role = new Role(stack, 'TestRole', {
+      assumedBy: new ServicePrincipal('lambda.amazonaws.com'),
+    });
+    const testFunc = new Function(stack, 'TestFunction', {
+      code: Code.fromInline("print('hello_world')"),
+      handler: 'index.handler',
+      runtime: Runtime.PYTHON_3_9,
+      environment: {
+        JOINED: `arn:${role.roleName}:foo:${role.roleArn}:bar`,
+      },
+    });
+    testFunc.grantInvoke(role);
+
+    const synthedApp = app.synth();
+    Aspects.of(app).add(
+      new ResourceExtractor({
+        extractDestinationStack: extractedStack,
+        stackArtifacts: synthedApp.stacks,
+        valueShareMethod: ResourceExtractorShareMethod.API_LOOKUP,
+        resourceTypesToExtract,
+      })
+    );
+    app.synth({ force: true });
+
+    const extractedTemplate = Template.fromStack(extractedStack);
+    const appTemplate = Template.fromStack(stack);
+
+    // Extracted stack has IAM resources
+    extractedTemplate.resourceCountIs('AWS::IAM::Role', 2);
+    extractedTemplate.resourceCountIs('AWS::IAM::Policy', 1);
+    // Non-IAM resources present in app stack
+    appTemplate.resourceCountIs('AWS::Lambda::Function', 1);
+    // Lambda function has correct environment variable that has been joined
+    // using the Fn.Join function, which combines the result of all strings
+    appTemplate.hasResourceProperties('AWS::Lambda::Function', {
+      Environment: {
+        Variables: {
+          JOINED: Match.stringLikeRegexp(
+            'arn:dummy-value-for-(.*?):foo:dummy-value-for-(.*?):bar'
+          ),
+        },
+      },
+    });
+  });
+
   test('Exports work correctly', () => {
     const secondStack = new Stack(app, 'Stack2', { env });
     const role = new Role(stack, 'ExportedRole', {