feat: add architecture prop to RustExtension (#77)

* feat: add `architecture` prop to `RustExtension`

* fix readme import

* add test

Author: garysassano

README.md

@@ -73,9 +73,11 @@ Define a `RustExtension` that get's deployed as a layer to use it with any other
 
 ```ts
 import { RustExtension, RustFunction } from 'cargo-lambda-cdk';
+import { Architecture } from 'aws-cdk-lib/aws-lambda';
 
 const extensionLayer = new RustExtension(this, 'Rust extension', {
   manifestPath: 'path/to/package/directory/with/Cargo.toml',
+  architecture: Architecture.ARM_64,
 });
 
 new RustFunction(this, 'Rust function', {

src/extension.ts

@@ -1,4 +1,8 @@
-import { LayerVersion, LayerVersionOptions } from 'aws-cdk-lib/aws-lambda';
+import {
+  LayerVersion,
+  LayerVersionOptions,
+  Architecture,
+} from 'aws-cdk-lib/aws-lambda';
 import { Construct } from 'constructs';
 import { Bundling } from './bundling';
 import { getManifestPath } from './cargo';
@@ -9,10 +13,10 @@ import { BundlingOptions } from './types';
  */
 export interface RustExtensionProps extends LayerVersionOptions {
   /**
- * Bundling options
- *
- * @default - use default bundling options
- */
+   * Bundling options
+   *
+   * @default - use default bundling options
+   */
   readonly bundling?: BundlingOptions;
 
   /**
@@ -57,24 +61,38 @@ export interface RustExtensionProps extends LayerVersionOptions {
    * temporary directory.
    */
   readonly gitForceClone?: boolean;
+
+  /**
+   * The system architecture of the lambda extension
+   *
+   * @default - Architecture.X86_64
+   */
+  readonly architecture?: Architecture;
 }
 
 /**
  * A Lambda extension written in Rust
  */
 export class RustExtension extends LayerVersion {
-  constructor(scope: Construct, resourceName: string, props?: RustExtensionProps) {
+  constructor(
+    scope: Construct,
+    resourceName: string,
+    props?: RustExtensionProps,
+  ) {
     const manifestPath = getManifestPath(props || {});
     const bundling = props?.bundling ?? {};
+    const architecture = props?.architecture ?? Architecture.X86_64;
 
     super(scope, resourceName, {
       ...props,
+      compatibleArchitectures: [architecture],
       code: Bundling.bundle({
         ...bundling,
         manifestPath,
         binaryName: props?.binaryName,
         lambdaExtension: true,
+        architecture,
       }),
     });
   }
-}
+}

test/extension.test.ts

@@ -1,6 +1,8 @@
 import { join } from 'path';
 import { env } from 'process';
 import { App, Stack } from 'aws-cdk-lib';
+import { Template } from 'aws-cdk-lib/assertions';
+import { Architecture } from 'aws-cdk-lib/aws-lambda';
 import { RustExtension, cargoLambdaVersion } from '../src/index';
 
 const forcedDockerBundling = !!env.FORCE_DOCKER_RUN || !cargoLambdaVersion();
@@ -48,4 +50,27 @@ describe('CargoLambda.RustExtension', () => {
       app.synth();
     });
   });
-});
+
+  describe('With ARM64 architecture', () => {
+    const app = new App();
+    const stack = new Stack(app);
+    const testSource = join(__dirname, 'fixtures/single-package');
+
+    new RustExtension(stack, 'rust extension arm64', {
+      manifestPath: testSource,
+      architecture: Architecture.ARM_64,
+    });
+
+    test('bundle extension with ARM64 architecture', () => {
+      // Synthesize the stack to a CloudFormation template
+      const template = Template.fromStack(stack);
+
+      // Verify that the Lambda layer has ARM64 in compatibleArchitectures
+      template.hasResourceProperties('AWS::Lambda::LayerVersion', {
+        CompatibleArchitectures: ['arm64'],
+      });
+
+      app.synth();
+    });
+  });
+});