fix(core): policy validation trace incorrect for larger constructs (#26466)

Ran into an issue where the construct trace was incorrect in larger projects, specifically where there are constructs that contain multiple constructs.

To get the construct trace tree we first construct a new tree that only contains the path to the terminal construct (the one with the trace). We drop all the other constructs in the tree that don't relate. For example, if we had a tree like:

```
->App
-->MyStage
--->MyStack1
---->MyConstruct
----->Resource
--->MyStack2
---->MyConstruct
----->Resource
--->MyStack3
---->MyConstruct
----->Resource
```

And we want to get a new tree for
`/App/MyStage/MyStack2/MyConstruct/Resource` it should look like this:

```
->App
-->MyStage
--->MyStack2
---->MyConstruct
----->Resource
```

We weren't correctly generating the new tree correctly and would always end up with the tree being the first item in the list.

I've updated one of the tests, and also tested this on a more complex application.


----

*By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*

Author: corymhall

packages/aws-cdk-lib/core/lib/private/tree-metadata.ts

@@ -10,7 +10,9 @@ import { ISynthesisSession } from '../stack-synthesizers';
 import { IInspectable, TreeInspector } from '../tree';
 
 const FILE_PATH = 'tree.json';
-
+type Mutable<T> = {
+  -readonly [P in keyof T]: Mutable<T[P]>;
+};
 /**
  * Construct that is automatically attached to the top-level `App`.
  * This generates, as part of synthesis, a file containing the construct tree and the metadata for each node in the tree.
@@ -109,21 +111,34 @@ export class TreeMetadata extends Construct {
    * tree that leads to a specific construct so drop any nodes not in that path
    *
    * @param node Node the current tree node
-   * @param child Node the previous tree node and the current node's child node
-   * @returns Node the new tree
+   * @returns Node the root node of the new tree
    */
-  private renderTreeWithChildren(node: Node, child?: Node): Node {
-    if (node.parent) {
-      return this.renderTreeWithChildren(node.parent, node);
-    } else if (child) {
-      return {
-        ...node,
-        children: {
-          [child.id]: child,
-        },
+  private renderTreeWithChildren(node: Node): Node {
+    /**
+     * @param currentNode - The current node being evaluated
+     * @param currentNodeChild - The previous node which should be the only child of the current node
+     * @returns The node with all children removed except for the path to the current node
+     */
+    function renderTreeWithSingleChild(currentNode: Mutable<Node>, currentNodeChild: Mutable<Node>) {
+      currentNode.children = {
+        [currentNodeChild.id]: currentNodeChild,
       };
+      if (currentNode.parent) {
+        currentNode.parent = renderTreeWithSingleChild(currentNode.parent, currentNode);
+      }
+      return currentNode;
     }
-    return node;
+
+    const currentNode = node.parent ? renderTreeWithSingleChild(node.parent, node) : node;
+    // now that we have the new tree we need to return the root node
+    let root = currentNode;
+    do {
+      if (root.parent) {
+        root = root.parent;
+      }
+    } while (root.parent);
+
+    return root;
   }
 
   /**

packages/aws-cdk-lib/core/lib/validation/private/construct-tree.ts

@@ -147,14 +147,26 @@ export class ConstructTree {
     return node;
   }
 
+  /**
+   * @param node - the root node of the tree
+   * @returns the terminal node in the tree
+   */
+  private lastChild(node: Node): Node {
+    if (node.children) {
+      return this.lastChild(this.getChild(node.children));
+    }
+    return node;
+  }
+
   /**
    * Get a ConstructTrace from the cache for a given construct
    *
    * Construct the stack trace of constructs. This will start with the
    * root of the tree and go down to the construct that has the violation
    */
   public getTrace(node: Node, locations?: string[]): ConstructTrace | undefined {
-    const trace = this._traceCache.get(node.path);
+    const lastChild = this.lastChild(node);
+    const trace = this._traceCache.get(lastChild.path);
     if (trace) {
       return trace;
     }
@@ -177,7 +189,9 @@ export class ConstructTree {
       libraryVersion: node.constructInfo?.version,
       location: thisLocation ?? "Run with '--debug' to include location info",
     };
-    this._traceCache.set(constructTrace.path, constructTrace);
+    // set the cache for the last child path. If the last child path is different then
+    // we have a different tree and need to retrieve the trace again
+    this._traceCache.set(lastChild.path, constructTrace);
     return constructTrace;
   }
 

packages/aws-cdk-lib/core/test/validation/trace.test.ts

@@ -113,6 +113,12 @@ class MyL2Resource extends Resource implements IMyL2Resource {
   public readonly constructPath: string;
   constructor(scope: Construct, id: string) {
     super(scope, id);
+    new core.CfnResource(this, 'Resource1', {
+      type: 'AWS::CDK::TestResource',
+      properties: {
+        testProp1: 'testValue',
+      },
+    });
     const resource = new core.CfnResource(this, 'Resource', {
       type: 'AWS::CDK::TestResource',
       properties: {
@@ -127,6 +133,7 @@ class MyConstruct extends Construct {
   public readonly constructPath: string;
   constructor(scope: Construct, id: string) {
     super(scope, id);
+    new MyL2Resource(this, 'MyL2Resource1');
     const myResource = new MyL2Resource(this, 'MyL2Resource');
     this.constructPath = myResource.constructPath;
   }
@@ -136,6 +143,8 @@ class MyStack extends core.Stack {
   public readonly constructPath: string;
   constructor(scope: Construct, id: string, props?: core.StackProps) {
     super(scope, id, props);
+    new MyConstruct(this, 'MyConstruct2');
+    new MyConstruct(this, 'MyConstruct3');
     const myConstruct = new MyConstruct(this, 'MyConstruct');
     this.constructPath = myConstruct.constructPath;
   }