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

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,37 @@ 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);
+      }
+      // we just reset the parent, but the currentNodeChild still has the original
+      // parent. Replace that as well
+      currentNode.children[currentNodeChild.id].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;
   }