chore(mixins-preview): improved find bucket policy (#36144)

### Reason for this change

A helper function to reflect on the construct tree to find the closet Bucket Policy for a Bucket.

### Description of changes

Implemented a generic function to find two related resources that are "close" to each other. Close is defined here as:
- the closest (transitive) child matching the predicate
- the above to the closest parent

Than added a specific function just for bucket. For now this is hand-written and not publicly exposed. If we can identify these kind of connection patterns, we can likely code gen this for all resources.

### Describe any new or updated permissions being added

n/a

### Description of how you validated changes

Unit tests.

### Checklist
- [x] My code adheres to the [CONTRIBUTING GUIDE](https://github.com/aws/aws-cdk/blob/main/CONTRIBUTING.md) and [DESIGN GUIDELINES](https://github.com/aws/aws-cdk/blob/main/docs/DESIGN_GUIDELINES.md)

----

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

Author: mrgrain

packages/@aws-cdk/mixins-preview/lib/mixins/private/reflections.ts

@@ -0,0 +1,90 @@
+import type { IConstruct } from 'constructs';
+import { CfnResource } from 'aws-cdk-lib/core';
+import type { IBucketRef, CfnBucketPolicy } from 'aws-cdk-lib/aws-s3';
+
+/**
+ * Finds the closest related resource in the construct tree.
+ * Searches children first (depth-first), then ancestors (breadth-first).
+ *
+ * @param primary - The construct to search a related resource for
+ * @param relatedCfnResourceType - The CloudFormation resource type to search for
+ * @param isConnected - Predicate to determine if a candidate is related to the primary
+ * @returns The closest matching resource, or undefined if none found
+ */
+function findClosestRelatedResource<TPrimary extends IConstruct, TRelated extends CfnResource>(
+  primary: TPrimary,
+  relatedCfnResourceType: string,
+  isConnected: (primary: TPrimary, candidate: TRelated) => boolean,
+): TRelated | undefined {
+  let closestMatch: TRelated | undefined;
+  let closestDistance = Infinity;
+  const visited = new Set<IConstruct>();
+
+  const isRelatedResource = (child: IConstruct): child is TRelated => {
+    return CfnResource.isCfnResource(child) && child.cfnResourceType === relatedCfnResourceType;
+  };
+
+  const checkCandidate = (candidate: IConstruct, distance: number) => {
+    // Check if candidate is the right type, connected to primary, and closer than current match
+    if (isRelatedResource(candidate) && isConnected(primary, candidate) && distance < closestDistance) {
+      closestMatch = candidate;
+      closestDistance = distance;
+    }
+  };
+
+  // Search all descendants depth-first
+  const searchChildren = (parent: IConstruct, distance: number) => {
+    // Stop searching if we're already farther than the closest match
+    if (distance >= closestDistance) {
+      return;
+    }
+
+    // Check each child and recursively search its descendants
+    for (const child of parent.node.children) {
+      // Skip if already visited
+      if (visited.has(child)) {
+        continue;
+      }
+      visited.add(child);
+
+      checkCandidate(child, distance);
+      searchChildren(child, distance + 1);
+    }
+  };
+
+  searchChildren(primary, 1);
+
+  // Search ancestors and their descendants breadth-first
+  let ancestor = primary.node.scope;
+  let ancestorDistance = 1;
+
+  // Walk up the tree while we have ancestors and haven't exceeded closest distance
+  while (ancestor && ancestorDistance < closestDistance) {
+    // Check all siblings and their descendants at this ancestor level
+    for (const sibling of ancestor.node.children) {
+      searchChildren(sibling, ancestorDistance);
+    }
+    ancestor = ancestor.node.scope;
+    ancestorDistance++;
+  }
+
+  return closestMatch;
+}
+
+/**
+ * Attempts to find an existing bucket policy for the specified S3 bucket.
+ * Finds the closest matching policy to the specified bucket.
+ *
+ * @param bucket - The S3 bucket reference to search for an associated bucket policy
+ * @returns The bucket policy if found, undefined otherwise
+ */
+export function tryFindBucketPolicyForBucket(bucket: IBucketRef): CfnBucketPolicy | undefined {
+  return findClosestRelatedResource<IBucketRef, CfnBucketPolicy>(
+    bucket,
+    'AWS::S3::BucketPolicy',
+    (b: any, policy) => {
+      const possibleRefs = new Set([b.ref, b.bucketName, b.bucketArn, b.bucketRef?.bucketName, b.bucketRef?.bucketArn].filter(Boolean));
+      return possibleRefs.has(policy.bucket) || String(policy.bucket).includes(b.node.id);
+    },
+  );
+}

packages/@aws-cdk/mixins-preview/lib/services/aws-s3/bucket.ts

@@ -3,6 +3,7 @@ import * as s3 from 'aws-cdk-lib/aws-s3';
 import { CfnResource, CustomResource, Tags } from 'aws-cdk-lib/core';
 import { AutoDeleteObjectsProvider } from '../../custom-resource-handlers/aws-s3/auto-delete-objects-provider';
 import type { IMixin } from '../../core';
+import { tryFindBucketPolicyForBucket } from '../../mixins/private/reflections';
 
 const AUTO_DELETE_OBJECTS_RESOURCE_TYPE = 'Custom::S3AutoDeleteObjects';
 const AUTO_DELETE_OBJECTS_TAG = 'aws-cdk:auto-delete-objects';
@@ -29,7 +30,7 @@ export class AutoDeleteObjects implements IMixin {
     });
 
     // Get or create bucket policy
-    let policy = construct.node.tryFindChild('Policy') as s3.CfnBucketPolicy | undefined;
+    let policy = tryFindBucketPolicyForBucket(construct);
     if (!policy) {
       policy = new s3.CfnBucketPolicy(construct, 'Policy', {
         bucket: ref.bucketName,

packages/@aws-cdk/mixins-preview/test/mixins/reflections.test.ts

@@ -0,0 +1,242 @@
+import { Construct } from 'constructs';
+import { Stack, App, Resource } from 'aws-cdk-lib/core';
+import * as s3 from 'aws-cdk-lib/aws-s3';
+import { tryFindBucketPolicyForBucket } from '../../lib/mixins/private/reflections';
+
+class CustomBucket extends Resource implements s3.IBucketRef {
+  public readonly bucketRef: s3.BucketReference;
+
+  constructor(scope: Construct, id: string, bucketName: string) {
+    super(scope, id);
+    this.bucketRef = {
+      bucketName,
+      bucketArn: `arn:aws:s3:::${bucketName}`,
+    };
+  }
+}
+
+describe('find bucket policy', () => {
+  let app: App;
+  let stack: Stack;
+
+  beforeEach(() => {
+    app = new App();
+    stack = new Stack(app, 'TestStack');
+  });
+
+  describe('find in construct tree', () => {
+    test('returns undefined when no bucket policy exists', () => {
+      const bucket = new s3.CfnBucket(stack, 'Bucket');
+      expect(tryFindBucketPolicyForBucket(bucket)).toBeUndefined();
+    });
+
+    test('finds bucket policy as direct child of bucket', () => {
+      const bucket = new s3.CfnBucket(stack, 'Bucket');
+      const policy = new s3.CfnBucketPolicy(bucket, 'Policy', {
+        bucket: bucket.ref,
+        policyDocument: {},
+      });
+
+      expect(tryFindBucketPolicyForBucket(bucket)).toBe(policy);
+    });
+
+    test('finds bucket policy as transitive child of bucket', () => {
+      const bucket = new s3.CfnBucket(stack, 'Bucket');
+      const intermediate = new Construct(bucket, 'Intermediate');
+      const policy = new s3.CfnBucketPolicy(intermediate, 'Policy', {
+        bucket: bucket.ref,
+        policyDocument: {},
+      });
+
+      expect(tryFindBucketPolicyForBucket(bucket)).toBe(policy);
+    });
+
+    test('finds bucket policy as sibling (child of parent)', () => {
+      const parent = new Construct(stack, 'Parent');
+      const bucket = new s3.CfnBucket(parent, 'Bucket');
+      const policy = new s3.CfnBucketPolicy(parent, 'Policy', {
+        bucket: bucket.ref,
+        policyDocument: {},
+      });
+
+      expect(tryFindBucketPolicyForBucket(bucket)).toBe(policy);
+    });
+
+    test('finds bucket policy in parent hierarchy', () => {
+      const grandparent = new Construct(stack, 'Grandparent');
+      const parent = new Construct(grandparent, 'Parent');
+      const bucket = new s3.CfnBucket(parent, 'Bucket');
+      const policy = new s3.CfnBucketPolicy(grandparent, 'Policy', {
+        bucket: bucket.ref,
+        policyDocument: {},
+      });
+
+      expect(tryFindBucketPolicyForBucket(bucket)).toBe(policy);
+    });
+
+    test('finds cousin bucket policies', () => {
+      const grandparent = new Construct(stack, 'Grandparent');
+      const parent = new Construct(grandparent, 'Parent');
+      const auncle = new Construct(grandparent, 'Auncle');
+
+      const bucket = new s3.CfnBucket(parent, 'Bucket');
+      const policy = new s3.CfnBucketPolicy(auncle, 'Policy', {
+        bucket: bucket.ref,
+        policyDocument: {},
+      });
+
+      expect(tryFindBucketPolicyForBucket(bucket)).toBe(policy);
+    });
+
+    test('prefers closest child over parent policy', () => {
+      const parent = new Construct(stack, 'Parent');
+      const bucket = new s3.CfnBucket(parent, 'Bucket');
+      const childPolicy = new s3.CfnBucketPolicy(bucket, 'ChildPolicy', {
+        bucket: bucket.ref,
+        policyDocument: {},
+      });
+      new s3.CfnBucketPolicy(parent, 'ParentPolicy', {
+        bucket: bucket.ref,
+        policyDocument: {},
+      });
+
+      expect(tryFindBucketPolicyForBucket(bucket)).toBe(childPolicy);
+    });
+
+    test('prefers closer transitive child over distant one', () => {
+      const bucket = new s3.CfnBucket(stack, 'Bucket');
+      const child1 = new Construct(bucket, 'Child1');
+      const closerPolicy = new s3.CfnBucketPolicy(child1, 'CloserPolicy', {
+        bucket: bucket.ref,
+        policyDocument: {},
+      });
+      const child2 = new Construct(bucket, 'Child2');
+      const intermediate = new Construct(child2, 'Intermediate');
+      new s3.CfnBucketPolicy(intermediate, 'FartherPolicy', {
+        bucket: bucket.ref,
+        policyDocument: {},
+      });
+
+      expect(tryFindBucketPolicyForBucket(bucket)).toBe(closerPolicy);
+    });
+
+    test('prefers closer parent over distant parent', () => {
+      const grandparent = new Construct(stack, 'Grandparent');
+      const parent = new Construct(grandparent, 'Parent');
+      const bucket = new s3.CfnBucket(parent, 'Bucket');
+      const closerPolicy = new s3.CfnBucketPolicy(parent, 'CloserPolicy', {
+        bucket: bucket.ref,
+        policyDocument: {},
+      });
+      new s3.CfnBucketPolicy(grandparent, 'FartherPolicy', {
+        bucket: bucket.ref,
+        policyDocument: {},
+      });
+
+      expect(tryFindBucketPolicyForBucket(bucket)).toBe(closerPolicy);
+    });
+
+    test('ignores unrelated bucket policies', () => {
+      const bucket1 = new s3.CfnBucket(stack, 'Bucket1');
+      const bucket2 = new s3.CfnBucket(stack, 'Bucket2');
+      new s3.CfnBucketPolicy(bucket2, 'Policy', {
+        bucket: bucket2.ref,
+        policyDocument: {},
+      });
+
+      expect(tryFindBucketPolicyForBucket(bucket1)).toBeUndefined();
+    });
+  });
+
+  describe('matches different buckets', () => {
+    test('matches L2 Bucket with policy using bucketName', () => {
+      const bucket = new s3.Bucket(stack, 'Bucket');
+      const policy = new s3.CfnBucketPolicy(bucket, 'Policy', {
+        bucket: bucket.bucketName,
+        policyDocument: {},
+      });
+
+      expect(tryFindBucketPolicyForBucket(bucket)).toBe(policy);
+    });
+
+    test('matches L2 Bucket with policy using bucketRef.bucketName', () => {
+      const bucket = new s3.Bucket(stack, 'Bucket');
+      const policy = new s3.CfnBucketPolicy(bucket, 'Policy', {
+        bucket: bucket.bucketRef.bucketName,
+        policyDocument: {},
+      });
+
+      expect(tryFindBucketPolicyForBucket(bucket)).toBe(policy);
+    });
+
+    test('matches L2 Bucket with policy using bucketArn', () => {
+      const bucket = new s3.Bucket(stack, 'Bucket');
+      const policy = new s3.CfnBucketPolicy(bucket, 'Policy', {
+        bucket: bucket.bucketArn,
+        policyDocument: {},
+      });
+
+      expect(tryFindBucketPolicyForBucket(bucket)).toBe(policy);
+    });
+
+    test('matches L2 Bucket', () => {
+      const bucket = new s3.Bucket(stack, 'Bucket');
+      const policy = new s3.CfnBucketPolicy(bucket, 'Policy', {
+        bucket: bucket,
+        policyDocument: {},
+      });
+
+      expect(tryFindBucketPolicyForBucket(bucket)).toBe(policy);
+    });
+
+    test('matches L1 Bucket', () => {
+      const bucket = new s3.CfnBucket(stack, 'Bucket');
+      const policy = new s3.CfnBucketPolicy(bucket, 'Policy', {
+        bucket: bucket,
+        policyDocument: {},
+      });
+
+      expect(tryFindBucketPolicyForBucket(bucket)).toBe(policy);
+    });
+
+    test('matches custom IBucketRef implementation', () => {
+      const bucket = new CustomBucket(stack, 'CustomBucket', 'my-bucket');
+      const policy = new s3.CfnBucketPolicy(bucket, 'Policy', {
+        bucket: bucket,
+        policyDocument: {},
+      });
+
+      expect(tryFindBucketPolicyForBucket(bucket)).toBe(policy);
+    });
+
+    test('matches L2 BucketPolicy', () => {
+      const bucket = new s3.Bucket(stack, 'Bucket');
+      const policy = new s3.BucketPolicy(bucket, 'Policy', {
+        bucket: bucket,
+      });
+
+      expect(tryFindBucketPolicyForBucket(bucket)).toBe(policy.node.defaultChild);
+    });
+
+    test('ignores policy for different bucket name', () => {
+      const bucket = new s3.CfnBucket(stack, 'Bucket');
+      bucket.bucketName = 'my-bucket';
+      new s3.CfnBucketPolicy(bucket, 'Policy', {
+        bucket: 'other-bucket',
+        policyDocument: {},
+      });
+
+      expect(tryFindBucketPolicyForBucket(bucket)).toBeUndefined();
+    });
+
+    test('matches using attrArn', () => {
+      const bucket = new s3.CfnBucket(stack, 'Bucket');
+      const policy = new s3.CfnBucketPolicy(bucket, 'Policy', {
+        bucket: bucket.attrArn,
+        policyDocument: {},
+      });
+
+      expect(tryFindBucketPolicyForBucket(bucket)).toBe(policy);
+    });
+  });
+});

packages/aws-cdk-lib/core/lib/cfn-element.ts

@@ -177,14 +177,14 @@ export abstract class CfnElement extends Construct {
 }
 
 /**
- * Base class for referencable CloudFormation constructs which are not Resources
+ * Base class for referenceable CloudFormation constructs which are not Resources
  *
  * These constructs are things like Conditions and Parameters, can be
  * referenced by taking the `.ref` attribute.
  *
  * Resource constructs do not inherit from CfnRefElement because they have their
  * own, more specific types returned from the .ref attribute. Also, some
- * resources aren't referencable at all (such as BucketPolicies or GatewayAttachments).
+ * resources aren't referenceable at all (such as BucketPolicies or GatewayAttachments).
  */
 export abstract class CfnRefElement extends CfnElement {
   /**