fix(custom-resources): retry CloudFormation response PUT in AwsCustomResource handler

packages/@aws-cdk/custom-resource-handlers/lib/aws-logs/log-retention-handler/index.ts

@@ -1,6 +1,7 @@
 
 // eslint-disable-next-line import/no-extraneous-dependencies
 import * as Logs from '@aws-sdk/client-cloudwatch-logs';
+import { DEFAULT_RESPONSE_RETRY_OPTIONS, httpRequest, withRetries } from '../../utils';
 
 let FAKE_SLEEP = false;
 
@@ -166,17 +167,7 @@ export async function handler(event: LogRetentionEvent, context: AWSLambda.Conte
       },
     };
 
-    return new Promise((resolve, reject) => {
-      try {
-        // eslint-disable-next-line @typescript-eslint/no-require-imports
-        const request = require('https').request(requestOptions, resolve);
-        request.on('error', reject);
-        request.write(responseBody);
-        request.end();
-      } catch (e) {
-        reject(e);
-      }
-    });
+    return withRetries(DEFAULT_RESPONSE_RETRY_OPTIONS, httpRequest)(requestOptions, responseBody);
   }
 }
 

packages/@aws-cdk/custom-resource-handlers/lib/custom-resources/aws-custom-resource-handler/utils.ts

@@ -2,6 +2,7 @@
 
 import type { AwsCredentialIdentityProvider } from '@smithy/types';
 import type { AwsSdkCall } from './construct-types';
+import { DEFAULT_RESPONSE_RETRY_OPTIONS, httpRequest, withRetries } from '../../utils';
 
 type Event = AWSLambda.CloudFormationCustomResourceEvent;
 
@@ -82,16 +83,7 @@ export function respond(
     },
   };
 
-  return new Promise((resolve, reject) => {
-    try {
-      const request = require('https').request(requestOptions, resolve);
-      request.on('error', reject);
-      request.write(responseBody);
-      request.end();
-    } catch (e) {
-      reject(e);
-    }
-  });
+  return withRetries(DEFAULT_RESPONSE_RETRY_OPTIONS, httpRequest)(requestOptions, responseBody);
 }
 
 /**

packages/@aws-cdk/custom-resource-handlers/lib/utils.ts

@@ -0,0 +1,98 @@
+import * as https from 'https';
+
+/**
+ * Shared helpers for sending a Custom Resource response back to CloudFormation
+ * over the pre-signed S3 response URL.
+ *
+ * This mirrors the behavior of the custom resource provider framework runtime:
+ * the response PUT is retried with exponential backoff, and only a successful
+ * (< 400) HTTP response is treated as success. It lives here so it can be
+ * shared by the bundled handlers in this package instead of being
+ * reimplemented per handler.
+ *
+ * `withRetries` mirrors:
+ * https://github.com/aws/aws-cdk/blob/3b1df7422f1e922849e94ec2a90928e6f2a05163/packages/aws-cdk-lib/custom-resources/lib/provider-framework/runtime/util.ts#L24-L40
+ * `httpRequest` mirrors `defaultHttpRequest` from:
+ * https://github.com/aws/aws-cdk/blob/3b1df7422f1e922849e94ec2a90928e6f2a05163/packages/aws-cdk-lib/custom-resources/lib/provider-framework/runtime/outbound.ts#L19-L36
+ *
+ * These cannot be imported directly: `aws-cdk-lib` depends on this package, so
+ * importing from it would create a circular dependency (see the
+ * `copied-from-aws-cdk-lib/` directory, which copies code for the same reason).
+ *
+ * NOTE: this module can only be consumed by handlers that are minified and
+ * bundled by the custom-resources-framework (`minifyAndBundle: true`), because
+ * esbuild inlines the import. Handlers that are copied verbatim
+ * (`minifyAndBundle: false`, e.g. the nodejs-entrypoint handler) cannot import
+ * it and must keep their own self-contained copy.
+ */
+
+export interface RetryOptions {
+  /** How many retries (will at least try once) */
+  readonly attempts: number;
+  /** Sleep base, in ms */
+  readonly sleep: number;
+}
+
+/**
+ * Default retry options for sending a Custom Resource response to CloudFormation.
+ *
+ * Shared by the bundled handlers so they retry the response PUT consistently
+ * (5 attempts, exponential backoff from a 1s base) instead of each PUT being
+ * a single un-retried attempt.
+ */
+export const DEFAULT_RESPONSE_RETRY_OPTIONS: RetryOptions = {
+  attempts: 5,
+  sleep: 1000,
+};
+
+/**
+ * Wraps an async function so it is retried with exponential backoff (and
+ * jitter) on failure, throwing the last error once the attempts are exhausted.
+ */
+export function withRetries<A extends Array<any>, B>(options: RetryOptions, fn: (...xs: A) => Promise<B>): (...xs: A) => Promise<B> {
+  return async (...xs: A) => {
+    let attempts = options.attempts;
+    let ms = options.sleep;
+    while (true) {
+      try {
+        return await fn(...xs);
+      } catch (e) {
+        if (attempts-- <= 0) {
+          throw e;
+        }
+        await sleep(Math.floor(Math.random() * ms));
+        ms *= 2;
+      }
+    }
+  };
+}
+
+async function sleep(ms: number): Promise<void> {
+  return new Promise((ok) => setTimeout(ok, ms));
+}
+
+/**
+ * Performs a single HTTP request (used to PUT the response to the CloudFormation
+ * pre-signed S3 response URL). Rejects on a network error or a non-successful
+ * (>= 400) HTTP response so that a caller wrapping it in `withRetries` can retry,
+ * instead of treating any received response as success.
+ */
+export async function httpRequest(options: https.RequestOptions, requestBody: string): Promise<void> {
+  return new Promise<void>((resolve, reject) => {
+    try {
+      const request = https.request(options, (response) => {
+        response.resume(); // Consume the response but don't care about it
+        if (!response.statusCode || response.statusCode >= 400) {
+          reject(new Error(`Unsuccessful HTTP response: ${response.statusCode}`));
+        } else {
+          resolve();
+        }
+      });
+      request.on('error', reject);
+      request.write(requestBody);
+      request.end();
+    } catch (e) {
+      reject(e);
+    }
+  });
+}

packages/@aws-cdk/custom-resource-handlers/test/custom-resources/aws-custom-resource-handler/utils.test.ts

@@ -1,5 +1,8 @@
+import * as https from 'https';
 import type { AwsSdkCall } from '../../../lib/custom-resources/aws-custom-resource-handler/construct-types';
-import { getCredentials } from '../../../lib/custom-resources/aws-custom-resource-handler/utils';
+import { getCredentials, respond } from '../../../lib/custom-resources/aws-custom-resource-handler/utils';
+
+jest.mock('https');
 
 // Mock the @aws-sdk/credential-providers import
 const mockFromTemporaryCredentials = jest.fn();
@@ -184,3 +187,84 @@ describe('getCredentials with External ID support', () => {
     expect(result).toBe(mockCredentials);
   });
 });
+
+describe('respond', () => {
+  function makeEvent(): AWSLambda.CloudFormationCustomResourceEvent {
+    return {
+      ResponseURL: 'https://cfn.example.com/response?token=abc',
+      StackId: '<StackId>',
+      RequestId: '<RequestId>',
+      LogicalResourceId: '<LogicalResourceId>',
+      ResourceType: '<ResourceType>',
+      ServiceToken: '<ServiceToken>',
+      ResourceProperties: { ServiceToken: '<ServiceToken>' },
+      RequestType: 'Create',
+    } as any;
+  }
+
+  beforeEach(() => {
+    // make backoff sleeps instantaneous and deterministic
+    jest.spyOn(Math, 'random').mockReturnValue(0);
+    jest.spyOn(console, 'log').mockImplementation(() => undefined);
+    (https.request as jest.Mock).mockReset();
+  });
+
+  afterEach(() => {
+    jest.restoreAllMocks();
+  });
+
+  test('PUTs the response to the CloudFormation response URL on success', async () => {
+    // GIVEN
+    let capturedOptions: any;
+    let capturedBody: string | undefined;
+    (https.request as jest.Mock).mockImplementation((options: any, cb: any) => {
+      capturedOptions = options;
+      cb({ statusCode: 200, resume: jest.fn() });
+      return { on: jest.fn(), write: (b: string) => { capturedBody = b; }, end: jest.fn() };
+    });
+
+    // WHEN
+    await respond(makeEvent(), 'SUCCESS', 'reason', 'physical-id', { Foo: 'Bar' }, true);
+
+    // THEN
+    expect(https.request).toHaveBeenCalledTimes(1);
+    expect(capturedOptions.method).toEqual('PUT');
+    expect(capturedOptions.hostname).toEqual('cfn.example.com');
+    expect(capturedOptions.path).toEqual('/response?token=abc');
+    const parsedBody = JSON.parse(capturedBody!);
+    expect(parsedBody.Status).toEqual('SUCCESS');
+    expect(parsedBody.PhysicalResourceId).toEqual('physical-id');
+    expect(parsedBody.Data).toEqual({ Foo: 'Bar' });
+  });
+
+  test('retries a transient failure and then succeeds', async () => {
+    // GIVEN: first attempt returns a 500, second attempt returns a 200
+    (https.request as jest.Mock)
+      .mockImplementationOnce((_options: any, cb: any) => {
+        cb({ statusCode: 500, resume: jest.fn() });
+        return { on: jest.fn(), write: jest.fn(), end: jest.fn() };
+      })
+      .mockImplementationOnce((_options: any, cb: any) => {
+        cb({ statusCode: 200, resume: jest.fn() });
+        return { on: jest.fn(), write: jest.fn(), end: jest.fn() };
+      });
+
+    // WHEN / THEN
+    await expect(respond(makeEvent(), 'SUCCESS', 'reason', 'physical-id', {}, false)).resolves.toBeUndefined();
+    expect(https.request).toHaveBeenCalledTimes(2);
+  });
+
+  test('rejects after exhausting retries when every attempt fails', async () => {
+    // GIVEN: every attempt returns a 500
+    (https.request as jest.Mock).mockImplementation((_options: any, cb: any) => {
+      cb({ statusCode: 500, resume: jest.fn() });
+      return { on: jest.fn(), write: jest.fn(), end: jest.fn() };
+    });
+
+    // WHEN / THEN
+    // RESPONSE_RETRY_OPTIONS.attempts is 5 => 1 initial call + 5 retries = 6 invocations
+    await expect(respond(makeEvent(), 'SUCCESS', 'reason', 'physical-id', {}, false))
+      .rejects.toThrow('Unsuccessful HTTP response: 500');
+    expect(https.request).toHaveBeenCalledTimes(6);
+  });
+});

packages/@aws-cdk/custom-resource-handlers/test/custom-resources/aws-custom-resource/aws-sdk-v3-handler.test.ts

@@ -66,11 +66,13 @@ jest.mock('@aws-sdk/credential-providers', () => {
 jest.mock('https', () => {
   return {
     ...jest.requireActual('https'),
-    request: (_: any, callback: () => void) => {
+    request: (_: any, callback: (res: any) => void) => {
       return {
         on: () => undefined,
         write: () => true,
-        end: callback,
+        // Invoke the response callback with a successful response so the
+        // handler's status-code check (>= 400 => retry/reject) passes.
+        end: () => callback({ statusCode: 200, resume: () => undefined }),
       };
     },
   };

packages/@aws-cdk/custom-resource-handlers/test/utils.test.ts

@@ -0,0 +1,85 @@
+import * as https from 'https';
+import { httpRequest, withRetries } from '../lib/utils';
+
+jest.mock('https');
+
+describe('withRetries', () => {
+  beforeEach(() => {
+    // make backoff sleeps instantaneous and deterministic
+    jest.spyOn(Math, 'random').mockReturnValue(0);
+  });
+
+  afterEach(() => {
+    jest.restoreAllMocks();
+  });
+
+  test('returns the result without retrying when the function succeeds', async () => {
+    const fn = jest.fn().mockResolvedValue('ok');
+
+    const result = await withRetries({ attempts: 5, sleep: 1 }, fn)();
+
+    expect(result).toEqual('ok');
+    expect(fn).toHaveBeenCalledTimes(1);
+  });
+
+  test('retries and eventually succeeds', async () => {
+    const fn = jest.fn()
+      .mockRejectedValueOnce(new Error('transient'))
+      .mockRejectedValueOnce(new Error('transient'))
+      .mockResolvedValue('ok');
+
+    const result = await withRetries({ attempts: 5, sleep: 1 }, fn)();
+
+    expect(result).toEqual('ok');
+    expect(fn).toHaveBeenCalledTimes(3);
+  });
+
+  test('throws the last error after exhausting all attempts', async () => {
+    const fn = jest.fn().mockRejectedValue(new Error('always fails'));
+
+    // attempts: 3 => 1 initial call + 3 retries = 4 invocations
+    await expect(withRetries({ attempts: 3, sleep: 1 }, fn)()).rejects.toThrow('always fails');
+    expect(fn).toHaveBeenCalledTimes(4);
+  });
+});
+
+describe('httpRequest', () => {
+  afterEach(() => {
+    jest.restoreAllMocks();
+    (https.request as jest.Mock).mockReset();
+  });
+
+  test('resolves on a successful (< 400) response', async () => {
+    (https.request as jest.Mock).mockImplementation((_options: any, cb: any) => {
+      cb({ statusCode: 200, resume: jest.fn() });
+      return { on: jest.fn(), write: jest.fn(), end: jest.fn() };
+    });
+
+    await expect(httpRequest({}, 'body')).resolves.toBeUndefined();
+  });
+
+  test('rejects on a >= 400 response', async () => {
+    (https.request as jest.Mock).mockImplementation((_options: any, cb: any) => {
+      cb({ statusCode: 500, resume: jest.fn() });
+      return { on: jest.fn(), write: jest.fn(), end: jest.fn() };
+    });
+
+    await expect(httpRequest({}, 'body')).rejects.toThrow('Unsuccessful HTTP response: 500');
+  });
+
+  test('rejects on a network error', async () => {
+    (https.request as jest.Mock).mockImplementation((_options: any, _cb: any) => {
+      return {
+        on: (event: string, handler: (e: Error) => void) => {
+          if (event === 'error') {
+            handler(new Error('socket hang up'));
+          }
+        },
+        write: jest.fn(),
+        end: jest.fn(),
+      };
+    });
+
+    await expect(httpRequest({}, 'body')).rejects.toThrow('socket hang up');
+  });
+});