From d9346bc9eedd2054bde00be930e80fa28bdd3a6e Mon Sep 17 00:00:00 2001
From: sakurai-ryo <sakurairyo.10do@gmail.com>
Date: Mon, 30 Dec 2024 05:12:32 +0900
Subject: [PATCH] fix(cli): unhandled nextToken returned by listImagesCommand
 in garbage collector for ECR (#32679)

### Issue # (if applicable)

Closes #32498

### Reason for this change
When `listImagesCommand` returns nextToken in the `readRepoInBatches` function, nextToken is not passed as an argument for the subsequent `listImagesCommand` execution, causing `listImagesCommand` to continue executing.
https://github.com/aws/aws-cdk/blob/v2.173.4/packages/aws-cdk/lib/api/garbage-collection/garbage-collector.ts#L621

According to the `listImagesCommand` documentation, if maxResults is not specified, a maximum of 100 images will be returned, so this bug requires at least 100 images in the asset repository.
https://docs.aws.amazon.com/AWSJavaScriptSDK/v3/latest/Package/-aws-sdk-client-ecr/Interface/ListImagesCommandInput/

#### Reproduction Steps
The following bash script and Dockerfile saved locally and executed, will push 120 container images to the asset repository.

```bash
#!/usr/bin/env bash

set -eu

ACCOUNT_ID="your account id"
REGION="your region"
REPO_NAME="cdk-hnb659fds-container-assets-${ACCOUNT_ID}-${REGION}"
IMAGE_NAME="test-image"
AWS_PROFILE="your AWS profile"

echo "Logging in to ECR..."
aws ecr get-login-password --region "${REGION}" --profile "${AWS_PROFILE}" \
| docker login --username AWS --password-stdin "${ACCOUNT_ID}.dkr.ecr.${REGION}.amazonaws.com"

for i in $(seq 1 120); do
  hash=$(head -c 32 /dev/urandom | xxd -p -c 64)
  echo "Building and pushing image with tag: ${hash}"
  touch "${i}.txt"

  docker build \
    --build-arg BUILD_NO="${i}" \
    -t "${IMAGE_NAME}:${i}" \
    .

  docker tag "${IMAGE_NAME}:${i}" \
    "${ACCOUNT_ID}.dkr.ecr.${REGION}.amazonaws.com/${REPO_NAME}:${hash}"

  docker push \
    "${ACCOUNT_ID}.dkr.ecr.${REGION}.amazonaws.com/${REPO_NAME}:${hash}"

  rm "${i}.txt"

  sleep 0.01
done

echo "Done!"
```

```dockerfile
FROM scratch

ARG BUILD_NO
ENV BUILD_NO=${BUILD_NO}

COPY ${BUILD_NO}.txt /
```

You can reproduce this bug by running the following command after the images have been pushed.
```bash
$ cdk gc aws://{account id}/{region} --type ecr --unstable=gc --created-buffer-days 0 --action full --confirm=true
```

### Description of changes
Fix the problem of correctly handling nextToken when executing `listImagesCommand` in the `readRepoInBatches` function.

### Describe any new or updated permissions being added
Nothing.

### Description of how you validated changes
Verifying that this bug has been fixed using the CLI integration tests is difficult, so only unit tests are added.

### 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*
---
 .../garbage-collection/garbage-collector.ts   |  1 +
 .../test/api/garbage-collection.test.ts       | 85 +++++++++++++++++++
 2 files changed, 86 insertions(+)

diff --git a/packages/aws-cdk/lib/api/garbage-collection/garbage-collector.ts b/packages/aws-cdk/lib/api/garbage-collection/garbage-collector.ts
index ba22174ebeaf1..34908cfec305b 100644
--- a/packages/aws-cdk/lib/api/garbage-collection/garbage-collector.ts
+++ b/packages/aws-cdk/lib/api/garbage-collection/garbage-collector.ts
@@ -621,6 +621,7 @@ export class GarbageCollector {
       while (batch.length < batchSize) {
         const response = await ecr.listImages({
           repositoryName: repo,
+          nextToken: continuationToken,
         });
 
         // No images in the repository
diff --git a/packages/aws-cdk/test/api/garbage-collection.test.ts b/packages/aws-cdk/test/api/garbage-collection.test.ts
index a64c4484efc15..d6a55dcf951ff 100644
--- a/packages/aws-cdk/test/api/garbage-collection.test.ts
+++ b/packages/aws-cdk/test/api/garbage-collection.test.ts
@@ -571,6 +571,91 @@ describe('ECR Garbage Collection', () => {
       imageTag: expect.stringContaining(`1-${ECR_ISOLATED_TAG}`),
     });
   });
+
+  test('listImagesCommand returns nextToken', async () => {
+    // This test is to ensure that the garbage collector can handle paginated responses from the ECR API
+    // If not handled correctly, the garbage collector will continue to make requests to the ECR API
+    mockTheToolkitInfo({
+      Outputs: [
+        {
+          OutputKey: 'BootstrapVersion',
+          OutputValue: '999',
+        },
+      ],
+    });
+
+    prepareDefaultEcrMock();
+    ecrClient.on(ListImagesCommand).resolves({ // default response
+      imageIds: [
+        {
+          imageDigest: 'digest1',
+          imageTag: 'abcde',
+        },
+        {
+          imageDigest: 'digest2',
+          imageTag: 'fghij',
+        },
+      ],
+      nextToken: 'nextToken',
+    }).on(ListImagesCommand, { // response when nextToken is provided
+      repositoryName: 'REPO_NAME',
+      nextToken: 'nextToken',
+    }).resolves({
+      imageIds: [
+        {
+          imageDigest: 'digest3',
+          imageTag: 'klmno',
+        },
+      ],
+    });
+    ecrClient.on(BatchGetImageCommand).resolvesOnce({
+      images: [
+        { imageId: { imageDigest: 'digest1' } },
+        { imageId: { imageDigest: 'digest2' } },
+      ],
+    }).resolvesOnce({
+      images: [
+        { imageId: { imageDigest: 'digest3' } },
+      ],
+    });
+    ecrClient.on(DescribeImagesCommand).resolvesOnce({
+      imageDetails: [
+        {
+          imageDigest: 'digest1',
+          imageTags: ['abcde'],
+          imagePushedAt: daysInThePast(100),
+          imageSizeInBytes: 1_000_000_000,
+        },
+        { imageDigest: 'digest2', imageTags: ['fghij'], imagePushedAt: daysInThePast(10), imageSizeInBytes: 300_000_000 },
+      ],
+    }).resolvesOnce({
+      imageDetails: [
+        { imageDigest: 'digest3', imageTags: ['klmno'], imagePushedAt: daysInThePast(2), imageSizeInBytes: 100 },
+      ],
+    });
+    prepareDefaultCfnMock();
+
+    garbageCollector = gc({
+      type: 'ecr',
+      rollbackBufferDays: 0,
+      action: 'full',
+    });
+    await garbageCollector.garbageCollect();
+
+    expect(ecrClient).toHaveReceivedCommandTimes(DescribeImagesCommand, 2);
+    expect(ecrClient).toHaveReceivedCommandTimes(ListImagesCommand, 4);
+
+    // no tagging
+    expect(ecrClient).toHaveReceivedCommandTimes(PutImageCommand, 0);
+
+    expect(ecrClient).toHaveReceivedCommandWith(BatchDeleteImageCommand, {
+      repositoryName: 'REPO_NAME',
+      imageIds: [
+        { imageDigest: 'digest2' },
+        { imageDigest: 'digest3' },
+      ],
+    });
+  });
 });
 
 describe('CloudFormation API calls', () => {