feat(lambda): add docker platform support for lambda (#21405)

Based on [this](#16770) PR
Add the missing part to add platform support when using lambda `fromImageAsset`

As we are not allowed to specify `platform` flag for `DockerImageAsset`, users deploying cdk on x86_64 platform will not be able to bundle lambda.DockerImageFunction for the new arm64 architecture. Similarly, users deploying cdk on arm64 architecture like Mac M1 will not be able to bundle images for AWS Fargate, which is x86_64 only.

# builder experience with aws-lambda
For x86_64 users deploying Lambda functions with container runtime on Lambda Graviton2(arm64) from local container image assets with multi-architecture docker base images. Specify the platform to ensure the image would be built for ARM64 architecture.

```
 new DockerImageFunction(stack, 'Lambda', {
      code: DockerImageCode.fromImageAsset(path.join(__dirname, 'docker-arm64-handler')),
      architecture: Architecture.ARM_64,
    });
```

Fixes: #12472, #20907
----

### All Submissions:

* [x] Have you followed the guidelines in our [Contributing guide?](https://github.com/aws/aws-cdk/blob/main/CONTRIBUTING.md)

### Adding new Unconventional Dependencies:

* [ ] This PR adds new unconventional dependencies following the process described [here](https://github.com/aws/aws-cdk/blob/main/CONTRIBUTING.md/#adding-new-unconventional-dependencies)

### New Features

* [x] Have you added the new feature to an [integration test](https://github.com/aws/aws-cdk/blob/main/INTEGRATION_TESTS.md)?
	* [x] Did you use `yarn integ` to deploy the infrastructure and generate the snapshot (i.e. `yarn integ` without `--dry-run`)?

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

packages/@aws-cdk/aws-lambda/README.md

@@ -81,6 +81,16 @@ new lambda.DockerImageFunction(this, 'ECRFunction', {
 The props for these docker image resources allow overriding the image's `CMD`, `ENTRYPOINT`, and `WORKDIR`
 configurations as well as choosing a specific tag or digest. See their docs for more information.
 
+To deploy a `DockerImageFunction` on Lambda `arm64` architecture, specify `Architecture.ARM_64` in `architecture`.
+This will bundle docker image assets for `arm64` architecture with `--platform linux/arm64` even if build within an `x86_64` host.
+
+```ts
+new DockerImageFunction(this, 'AssetFunction', {
+  code: DockerImageCode.fromImageAsset(path.join(__dirname, 'docker-arm64-handler')),
+  architecture: Architecture.ARM_64,
+});
+```
+
 ## Execution Role
 
 Lambda functions assume an IAM role during execution. In CDK by default, Lambda

packages/@aws-cdk/aws-lambda/lib/image-function.ts

@@ -1,5 +1,7 @@
 import * as ecr from '@aws-cdk/aws-ecr';
+import { Platform } from '@aws-cdk/aws-ecr-assets';
 import { Construct } from 'constructs';
+import { Architecture } from './architecture';
 import { AssetImageCode, AssetImageCodeProps, EcrImageCode, EcrImageCodeProps, Code } from './code';
 import { Function, FunctionOptions } from './function';
 import { Handler } from './handler';
@@ -41,8 +43,12 @@ export abstract class DockerImageCode {
    */
   public static fromImageAsset(directory: string, props: AssetImageCodeProps = {}): DockerImageCode {
     return {
-      _bind() {
-        return new AssetImageCode(directory, props);
+      _bind(architecture?: Architecture) {
+        return new AssetImageCode(directory, {
+          // determine the platform from `architecture`.
+          ...architecture?.dockerPlatform ? { platform: Platform.custom(architecture.dockerPlatform) } : {},
+          ...props,
+        });
       },
     };
   }
@@ -51,7 +57,7 @@ export abstract class DockerImageCode {
    * Produce a `Code` instance from this `DockerImageCode`.
    * @internal
    */
-  public abstract _bind(): Code;
+  public abstract _bind(architecture?: Architecture): Code;
 }
 
 /**
@@ -63,7 +69,7 @@ export class DockerImageFunction extends Function {
       ...props,
       handler: Handler.FROM_IMAGE,
       runtime: Runtime.FROM_IMAGE,
-      code: props.code._bind(),
+      code: props.code._bind(props.architecture),
     });
   }
 }

packages/@aws-cdk/aws-lambda/test/docker-arm64-handler/Dockerfile

@@ -0,0 +1,6 @@
+FROM public.ecr.aws/lambda/python:latest
+
+ARG FUNCTION_DIR="/var/task"
+COPY index.py ${FUNCTION_DIR}
+
+CMD [ "index.handler" ]

packages/@aws-cdk/aws-lambda/test/docker-arm64-handler/index.py

@@ -0,0 +1,6 @@
+import json, platform
+def handler(event, context):
+  return {
+    'statusCode': 200,
+    'body': json.dumps( f'Hello CDK from Lambda({platform.platform()})!')
+  }

packages/@aws-cdk/aws-lambda/test/integ.lambda.docker-arm64.ts

@@ -0,0 +1,19 @@
+import * as path from 'path';
+import { App, Stack } from '@aws-cdk/core';
+import * as integ from '@aws-cdk/integ-tests';
+import { Architecture, DockerImageCode, DockerImageFunction } from '../lib';
+
+const app = new App();
+
+const stack = new Stack(app, 'lambda-ecr-docker-arm64');
+
+new DockerImageFunction(stack, 'MyLambda', {
+  code: DockerImageCode.fromImageAsset(path.join(__dirname, 'docker-arm64-handler')),
+  architecture: Architecture.ARM_64,
+});
+
+new integ.IntegTest(app, 'lambda-docker-arm64', {
+  testCases: [stack],
+});
+
+app.synth();

packages/@aws-cdk/aws-lambda/test/lambda-platform.test.ts

@@ -0,0 +1,44 @@
+import * as path from 'path';
+import { Template } from '@aws-cdk/assertions';
+import * as cdk from '@aws-cdk/core';
+import { Architecture, DockerImageCode, DockerImageFunction } from '../lib';
+
+describe('lambda platform', () => {
+  test('can choose lambda architecture arm64', () => {
+    // GIVEN
+    const app = new cdk.App();
+    const stack = new cdk.Stack(app, 'stack');
+
+    // WHEN
+    new DockerImageFunction(stack, 'Lambda', {
+      code: DockerImageCode.fromImageAsset(path.join(__dirname, 'docker-arm64-handler')),
+      architecture: Architecture.ARM_64,
+    });
+
+    // THEN
+    Template.fromStack(stack).hasResourceProperties('AWS::Lambda::Function', {
+      Architectures: [
+        'arm64',
+      ],
+    });
+  });
+
+  test('can choose lambda architecture x86_64', () => {
+    // GIVEN
+    const app = new cdk.App();
+    const stack = new cdk.Stack(app, 'stack');
+
+    // WHEN
+    new DockerImageFunction(stack, 'Lambda', {
+      code: DockerImageCode.fromImageAsset(path.join(__dirname, 'docker-arm64-handler')),
+      architecture: Architecture.X86_64,
+    });
+
+    // THEN
+    Template.fromStack(stack).hasResourceProperties('AWS::Lambda::Function', {
+      Architectures: [
+        'x86_64',
+      ],
+    });
+  });
+});

packages/@aws-cdk/aws-lambda/test/lambda.docker-arm64.integ.snapshot/cdk.out

@@ -0,0 +1 @@
+{"version":"20.0.0"}

packages/@aws-cdk/aws-lambda/test/lambda.docker-arm64.integ.snapshot/integ.json

@@ -0,0 +1,11 @@
+{
+  "version": "20.0.0",
+  "testCases": {
+    "lambda-docker-arm64/DefaultTest": {
+      "stacks": [
+        "lambda-ecr-docker-arm64"
+      ],
+      "assertionStack": "lambdadockerarm64DefaultTestDeployAssert07D408EF"
+    }
+  }
+}

packages/@aws-cdk/aws-lambda/test/lambda.docker-arm64.integ.snapshot/integ.lambda.docker-arm64.expected.json

@@ -0,0 +1,74 @@
+{
+  "Resources": {
+    "MyLambdaServiceRole4539ECB6": {
+      "Type": "AWS::IAM::Role",
+      "Properties": {
+        "AssumeRolePolicyDocument": {
+          "Statement": [
+            {
+              "Action": "sts:AssumeRole",
+              "Effect": "Allow",
+              "Principal": {
+                "Service": "lambda.amazonaws.com"
+              }
+            }
+          ],
+          "Version": "2012-10-17"
+        },
+        "ManagedPolicyArns": [
+          {
+            "Fn::Join": [
+              "",
+              [
+                "arn:",
+                {
+                  "Ref": "AWS::Partition"
+                },
+                ":iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
+              ]
+            ]
+          }
+        ]
+      }
+    },
+    "MyLambdaCCE802FB": {
+      "Type": "AWS::Lambda::Function",
+      "Properties": {
+        "Code": {
+          "ImageUri": {
+            "Fn::Join": [
+              "",
+              [
+                {
+                  "Ref": "AWS::AccountId"
+                },
+                ".dkr.ecr.",
+                {
+                  "Ref": "AWS::Region"
+                },
+                ".",
+                {
+                  "Ref": "AWS::URLSuffix"
+                },
+                "/aws-cdk/assets:166017f6f6f219800dee266c149e208fe18dea1788d822b1783afbc008c25db7"
+              ]
+            ]
+          }
+        },
+        "Role": {
+          "Fn::GetAtt": [
+            "MyLambdaServiceRole4539ECB6",
+            "Arn"
+          ]
+        },
+        "Architectures": [
+          "arm64"
+        ],
+        "PackageType": "Image"
+      },
+      "DependsOn": [
+        "MyLambdaServiceRole4539ECB6"
+      ]
+    }
+  }
+}

packages/@aws-cdk/aws-lambda/test/lambda.docker-arm64.integ.snapshot/lambda-ecr-docker-arm64.assets.json

@@ -0,0 +1,33 @@
+{
+  "version": "20.0.0",
+  "files": {
+    "e57a612d901fd500f4eae1beabe146b3d7a4f81dd10a81199b9c884d1e227a1e": {
+      "source": {
+        "path": "lambda-ecr-docker-arm64.template.json",
+        "packaging": "file"
+      },
+      "destinations": {
+        "current_account-current_region": {
+          "bucketName": "cdk-hnb659fds-assets-${AWS::AccountId}-${AWS::Region}",
+          "objectKey": "e57a612d901fd500f4eae1beabe146b3d7a4f81dd10a81199b9c884d1e227a1e.json",
+          "assumeRoleArn": "arn:${AWS::Partition}:iam::${AWS::AccountId}:role/cdk-hnb659fds-file-publishing-role-${AWS::AccountId}-${AWS::Region}"
+        }
+      }
+    }
+  },
+  "dockerImages": {
+    "027b9b499ce9e488d4c3cfa41abdbdc6afe203989a5bd77258f471da03f3f040": {
+      "source": {
+        "directory": "asset.027b9b499ce9e488d4c3cfa41abdbdc6afe203989a5bd77258f471da03f3f040",
+        "platform": "linux/arm64"
+      },
+      "destinations": {
+        "current_account-current_region": {
+          "repositoryName": "cdk-hnb659fds-container-assets-${AWS::AccountId}-${AWS::Region}",
+          "imageTag": "027b9b499ce9e488d4c3cfa41abdbdc6afe203989a5bd77258f471da03f3f040",
+          "assumeRoleArn": "arn:${AWS::Partition}:iam::${AWS::AccountId}:role/cdk-hnb659fds-image-publishing-role-${AWS::AccountId}-${AWS::Region}"
+        }
+      }
+    }
+  }
+}

packages/@aws-cdk/aws-lambda/test/lambda.docker-arm64.integ.snapshot/lambda-ecr-docker-arm64.template.json

@@ -0,0 +1,74 @@
+{
+ "Resources": {
+  "MyLambdaServiceRole4539ECB6": {
+   "Type": "AWS::IAM::Role",
+   "Properties": {
+    "AssumeRolePolicyDocument": {
+     "Statement": [
+      {
+       "Action": "sts:AssumeRole",
+       "Effect": "Allow",
+       "Principal": {
+        "Service": "lambda.amazonaws.com"
+       }
+      }
+     ],
+     "Version": "2012-10-17"
+    },
+    "ManagedPolicyArns": [
+     {
+      "Fn::Join": [
+       "",
+       [
+        "arn:",
+        {
+         "Ref": "AWS::Partition"
+        },
+        ":iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
+       ]
+      ]
+     }
+    ]
+   }
+  },
+  "MyLambdaCCE802FB": {
+   "Type": "AWS::Lambda::Function",
+   "Properties": {
+    "Code": {
+     "ImageUri": {
+      "Fn::Join": [
+       "",
+       [
+        {
+         "Ref": "AWS::AccountId"
+        },
+        ".dkr.ecr.",
+        {
+         "Ref": "AWS::Region"
+        },
+        ".",
+        {
+         "Ref": "AWS::URLSuffix"
+        },
+        "/aws-cdk/assets:027b9b499ce9e488d4c3cfa41abdbdc6afe203989a5bd77258f471da03f3f040"
+       ]
+      ]
+     }
+    },
+    "Role": {
+     "Fn::GetAtt": [
+      "MyLambdaServiceRole4539ECB6",
+      "Arn"
+     ]
+    },
+    "Architectures": [
+     "arm64"
+    ],
+    "PackageType": "Image"
+   },
+   "DependsOn": [
+    "MyLambdaServiceRole4539ECB6"
+   ]
+  }
+ }
+}

packages/@aws-cdk/aws-lambda/test/lambda.docker-arm64.integ.snapshot/lambdadockerarm64DefaultTestDeployAssert07D408EF.template.json

@@ -0,0 +1 @@
+{}