Use conditional dependencies for documents that are conditional

source/lib/ssm-doc-rate-limit.test.ts

@@ -0,0 +1,222 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+import { CfnDocument } from 'aws-cdk-lib/aws-ssm';
+import { Aspects, CfnCondition, CfnParameter, Fn, Stack } from 'aws-cdk-lib';
+import SsmDocRateLimit from './ssm-doc-rate-limit';
+import { WaitProvider } from './wait-provider';
+import { Template } from 'aws-cdk-lib/assertions';
+
+describe('SSM doc rate limit aspect', function () {
+  const stack = new Stack();
+  const serviceToken = 'my-token';
+  const waitProvider = WaitProvider.fromServiceToken(stack, 'WaitProvider', serviceToken);
+  Aspects.of(stack).add(new SsmDocRateLimit(waitProvider));
+  const content = {};
+  const numDocuments = 12;
+  for (let i = 0; i < numDocuments; ++i) {
+    new CfnDocument(stack, `Document${i}`, { content });
+  }
+  const template = Template.fromStack(stack);
+
+  it('matches snapshot', function () {
+    expect(template).toMatchSnapshot();
+  });
+
+  const documents = template.findResources('AWS::SSM::Document', { Properties: { Content: content } });
+  const documentLogicalIds = Object.getOwnPropertyNames(documents);
+
+  const expectedBatchSize = 5;
+
+  const createWaits = template.findResources('Custom::Wait', {
+    Properties: {
+      CreateIntervalSeconds: 1,
+      UpdateIntervalSeconds: 1,
+      DeleteIntervalSeconds: 0,
+      ServiceToken: serviceToken,
+    },
+  });
+  const createWaitLogicalIds = Object.getOwnPropertyNames(createWaits);
+
+  it('has the correct number of create resources', function () {
+    expect(createWaitLogicalIds).toHaveLength(Math.ceil(numDocuments / expectedBatchSize));
+  });
+
+  const deleteWaits = template.findResources('Custom::Wait', {
+    Properties: {
+      CreateIntervalSeconds: 0,
+      UpdateIntervalSeconds: 0,
+      DeleteIntervalSeconds: 0.5,
+      ServiceToken: serviceToken,
+    },
+  });
+  const deleteWaitLogicalIds = Object.getOwnPropertyNames(deleteWaits);
+
+  it('has the correct number of delete resources', function () {
+    expect(deleteWaitLogicalIds).toHaveLength(Math.ceil(numDocuments / expectedBatchSize));
+  });
+
+  it('create resources form dependency chain', function () {
+    expect(formDependencyChain(createWaits)).toStrictEqual(true);
+  });
+
+  it('delete resources form dependency chain', function () {
+    expect(formDependencyChain(deleteWaits)).toStrictEqual(true);
+  });
+
+  it('documents depend on create and delete depends on documents', function () {
+    const documentSets: string[][] = [];
+    deleteWaitLogicalIds.forEach(function (logicalId: string) {
+      documentSets.push(
+        (deleteWaits[logicalId].DependsOn as Array<string>).filter(function (value: string) {
+          return documentLogicalIds.includes(value);
+        })
+      );
+    });
+    const remainingDocuments = { ...documents };
+    documentSets.forEach(function (documentSet: string[]) {
+      // all documents depend on the same create resource
+      const expectedCreateResource = documents[documentSet[0]].DependsOn[0];
+      documentSet.forEach(function (value: string) {
+        delete remainingDocuments[value];
+        expect(documents[value].DependsOn).toHaveLength(1);
+        expect(documents[value].DependsOn[0]).toStrictEqual(expectedCreateResource);
+      });
+    });
+    // all documents in a set
+    expect(Object.getOwnPropertyNames(remainingDocuments)).toHaveLength(0);
+  });
+});
+
+describe('SSM doc rate limit aspect with conditional documents', function () {
+  const stack = new Stack();
+  const serviceToken = 'my-token';
+  const waitProvider = WaitProvider.fromServiceToken(stack, 'WaitProvider', serviceToken);
+  Aspects.of(stack).add(new SsmDocRateLimit(waitProvider));
+  const content = {};
+  const numDocuments = 12;
+  for (let i = 0; i < numDocuments; ++i) {
+    const param = new CfnParameter(stack, `Parameter${i}`);
+    const condition = new CfnCondition(stack, `Condition${i}`, { expression: Fn.conditionEquals(param, 'asdf') });
+    const doc = new CfnDocument(stack, `Document${i}`, { content });
+    doc.cfnOptions.condition = condition;
+  }
+  const template = Template.fromStack(stack);
+
+  it('matches snapshot', function () {
+    expect(template).toMatchSnapshot();
+  });
+
+  const documents = template.findResources('AWS::SSM::Document', { Properties: { Content: content } });
+  const documentLogicalIds = Object.getOwnPropertyNames(documents);
+
+  const expectedBatchSize = 5;
+
+  const createWaits = template.findResources('Custom::Wait', {
+    Properties: {
+      CreateIntervalSeconds: 1,
+      UpdateIntervalSeconds: 1,
+      DeleteIntervalSeconds: 0,
+      ServiceToken: serviceToken,
+    },
+  });
+  const createWaitLogicalIds = Object.getOwnPropertyNames(createWaits);
+
+  it('has the correct number of create resources', function () {
+    expect(createWaitLogicalIds).toHaveLength(Math.ceil(numDocuments / expectedBatchSize));
+  });
+
+  const deleteWaits = template.findResources('Custom::Wait', {
+    Properties: {
+      CreateIntervalSeconds: 0,
+      UpdateIntervalSeconds: 0,
+      DeleteIntervalSeconds: 0.5,
+      ServiceToken: serviceToken,
+    },
+  });
+  const deleteWaitLogicalIds = Object.getOwnPropertyNames(deleteWaits);
+
+  it('has the correct number of delete resources', function () {
+    expect(deleteWaitLogicalIds).toHaveLength(Math.ceil(numDocuments / expectedBatchSize));
+  });
+
+  it('create resources form dependency chain', function () {
+    expect(formDependencyChain(createWaits)).toStrictEqual(true);
+  });
+
+  it('delete resources form dependency chain', function () {
+    expect(formDependencyChain(deleteWaits)).toStrictEqual(true);
+  });
+
+  const dummyResources = template.findResources('AWS::CloudFormation::WaitConditionHandle');
+  const dummyResourceLogicalIds = Object.getOwnPropertyNames(dummyResources);
+
+  it('documents depend on create and delete depends on documents', function () {
+    const documentSets: string[][] = [];
+    deleteWaitLogicalIds.forEach(function (logicalId: string) {
+      const documentSet: string[] = [];
+      const dependencies = deleteWaits[logicalId].DependsOn as Array<string>;
+      dependencies.forEach(function (value: string) {
+        if (dummyResourceLogicalIds.includes(value)) {
+          const dummyResource = dummyResources[value];
+          Object.entries(dummyResource.Metadata).forEach(function (meta: [string, unknown]) {
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            documentSet.push((meta[1] as { [_: string]: any })['Fn::If'][1].Ref);
+          });
+        }
+      });
+      documentSet.push(
+        ...dependencies.filter(function (value: string) {
+          return documentLogicalIds.includes(value);
+        })
+      );
+      documentSets.push(documentSet);
+    });
+    const remainingDocuments = { ...documents };
+    documentSets.forEach(function (documentSet: string[]) {
+      // all documents depend on the same create resource
+      const expectedCreateResource = documents[documentSet[0]].DependsOn[0];
+      documentSet.forEach(function (value: string) {
+        delete remainingDocuments[value];
+        expect(documents[value].DependsOn).toHaveLength(1);
+        expect(documents[value].DependsOn[0]).toStrictEqual(expectedCreateResource);
+      });
+    });
+    // all documents in a set
+    expect(Object.getOwnPropertyNames(remainingDocuments)).toHaveLength(0);
+  });
+});
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+type Resources = { [_: string]: { [_: string]: any } };
+
+// do the resources depend on each other in a serial manner
+// this isn't foolproof, but it should be enough for simple cases
+function formDependencyChain(resources: Resources): boolean {
+  const logicalIds = Object.getOwnPropertyNames(resources);
+  let dependencyChainFound = false;
+  // if so, there will be a resource which starts a chain that contains all the other resources
+  logicalIds.forEach(function (logicalId: string | undefined) {
+    const resourcesRemaining = { ...resources };
+    while (logicalId) {
+      let dependencies = resourcesRemaining[logicalId].DependsOn;
+      // only check dependencies of the same resource type
+      if (dependencies) {
+        dependencies = (dependencies as Array<string>).filter(function (value: string) {
+          return logicalIds.includes(value);
+        });
+      }
+      delete resourcesRemaining[logicalId];
+      if (dependencies && dependencies.length != 0) {
+        expect(dependencies).toHaveLength(1);
+        logicalId = dependencies[0];
+      } else {
+        logicalId = undefined;
+      }
+    }
+    // if there are no resources left, this resource is the terminal resource
+    if (Object.getOwnPropertyNames(resourcesRemaining).length === 0) {
+      dependencyChainFound = true;
+    }
+  });
+  return dependencyChainFound;
+}

source/lib/ssm-doc-rate-limit.ts

@@ -1,6 +1,6 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
-import { CfnCustomResource, CustomResource, IAspect, Stack } from 'aws-cdk-lib';
+import { CfnCustomResource, CfnWaitConditionHandle, CustomResource, Fn, IAspect, Stack } from 'aws-cdk-lib';
 import { CfnDocument } from 'aws-cdk-lib/aws-ssm';
 import { Construct, IConstruct } from 'constructs';
 import { createHash, Hash } from 'crypto';
@@ -13,6 +13,7 @@ export default class SsmDocRateLimit implements IAspect {
   private previousCreateWaitResource: CustomResource | undefined;
   private currentDeleteWaitResource: CustomResource | undefined;
   private previousDeleteWaitResource: CustomResource | undefined;
+  private currentDummyResource: CfnWaitConditionHandle | undefined;
 
   private hash: Hash;
 
@@ -51,6 +52,12 @@ export default class SsmDocRateLimit implements IAspect {
     }
   }
 
+  initDummyResource(scope: Construct): void {
+    if (!this.currentDummyResource) {
+      this.currentDummyResource = new CfnWaitConditionHandle(scope, `Gate${this.waitResourceIndex - 1}`);
+    }
+  }
+
   visit(node: IConstruct): void {
     if (node instanceof CfnDocument) {
       const scope = Stack.of(node);
@@ -67,7 +74,20 @@ export default class SsmDocRateLimit implements IAspect {
       updateWaitResourceHash(this.currentDeleteWaitResource, digest);
 
       node.addDependency(this.currentCreateWaitResource.node.defaultChild as CfnCustomResource);
-      this.currentDeleteWaitResource.node.addDependency(node);
+
+      if (node.cfnOptions.condition) {
+        this.initDummyResource(scope);
+        if (!this.currentDummyResource) {
+          throw new Error('Dummy resource not initialized!');
+        }
+        this.currentDummyResource.addMetadata(
+          `${node.logicalId}Ready`,
+          Fn.conditionIf(node.cfnOptions.condition.logicalId, Fn.ref(node.logicalId), '')
+        );
+        this.currentDeleteWaitResource.node.addDependency(this.currentDummyResource);
+      } else {
+        this.currentDeleteWaitResource.node.addDependency(node);
+      }
 
       ++this.documentIndex;
 
@@ -77,6 +97,7 @@ export default class SsmDocRateLimit implements IAspect {
         this.previousDeleteWaitResource = this.currentDeleteWaitResource;
         this.currentCreateWaitResource = undefined;
         this.currentDeleteWaitResource = undefined;
+        this.currentDummyResource = undefined;
         this.hash = createHash('sha256');
       }
     }

source/playbooks/AFSBP/test/__snapshots__/afsbp_stack.test.ts.snap

@@ -1131,9 +1131,7 @@ def parse_event(event, _):
     "DeletWait0": {
       "DeletionPolicy": "Delete",
       "DependsOn": [
-        "ControlAFSBPEC21",
-        "ControlAFSBPLambda1",
-        "ControlAFSBPRDS1",
+        "Gate0",
       ],
       "Properties": {
         "CreateIntervalSeconds": 0,
@@ -1147,6 +1145,38 @@ def parse_event(event, _):
       "Type": "Custom::Wait",
       "UpdateReplacePolicy": "Delete",
     },
+    "Gate0": {
+      "Metadata": {
+        "ControlAFSBPEC21Ready": {
+          "Fn::If": [
+            "EnableEC21Condition",
+            {
+              "Ref": "ControlAFSBPEC21",
+            },
+            "",
+          ],
+        },
+        "ControlAFSBPLambda1Ready": {
+          "Fn::If": [
+            "EnableLambda1Condition",
+            {
+              "Ref": "ControlAFSBPLambda1",
+            },
+            "",
+          ],
+        },
+        "ControlAFSBPRDS1Ready": {
+          "Fn::If": [
+            "EnableRDS1Condition",
+            {
+              "Ref": "ControlAFSBPRDS1",
+            },
+            "",
+          ],
+        },
+      },
+      "Type": "AWS::CloudFormation::WaitConditionHandle",
+    },
   },
 }
 `;

source/playbooks/CIS120/test/__snapshots__/cis_stack.test.ts.snap

@@ -1481,9 +1481,7 @@ def parse_event(event, _):
     "DeletWait0": {
       "DeletionPolicy": "Delete",
       "DependsOn": [
-        "ControlCIS13",
-        "ControlCIS15",
-        "ControlCIS21",
+        "Gate0",
       ],
       "Properties": {
         "CreateIntervalSeconds": 0,
@@ -1497,6 +1495,38 @@ def parse_event(event, _):
       "Type": "Custom::Wait",
       "UpdateReplacePolicy": "Delete",
     },
+    "Gate0": {
+      "Metadata": {
+        "ControlCIS13Ready": {
+          "Fn::If": [
+            "Enable13Condition",
+            {
+              "Ref": "ControlCIS13",
+            },
+            "",
+          ],
+        },
+        "ControlCIS15Ready": {
+          "Fn::If": [
+            "Enable15Condition",
+            {
+              "Ref": "ControlCIS15",
+            },
+            "",
+          ],
+        },
+        "ControlCIS21Ready": {
+          "Fn::If": [
+            "Enable21Condition",
+            {
+              "Ref": "ControlCIS21",
+            },
+            "",
+          ],
+        },
+      },
+      "Type": "AWS::CloudFormation::WaitConditionHandle",
+    },
     "RemapCIS4245EB49A0": {
       "Properties": {
         "Description": "Remap the CIS 4.2 finding to CIS 4.1 remediation",