Zenn
DRESS CODE TECH BLOGPublicationへの投稿
🫧

Node.jsで動いているLambdaをRustに移行してみた

かわうそ
に公開
AWS
Node.js
Rust
Lambda
AWS CDK
tech

TL;DR

  • 移行前のLambda
    • パッケージタイプ：Zip
    • ランタイム：Node.js 20.x
  • 移行後のLambda
    • パッケージタイプ：Image
    • ランタイム：Rust
  • あくまで試しにやってみたときの作業ログ程度の内容

移行対象

  • Cognitoの認証時に設定しているPre token generation Lambda triggerが対象
  • Pre token generation Lambda trigger内ではBackendのAPIにアクセスし必要な情報を取得
  • 取得した情報を元にJWTのPayloadを加工

参考

移行前のコード（Node.js）

handler.ts
export const handler = async (
  event: PreTokenGenerationV2TriggerEvent,
  context: Context,
) => {
  // Backend の API を呼び出すために必要なキーをSecretManager等から取得
  const apiKey = await getApiKey(); 

  const response = await axios.post(
    `${process.env.API_ENDPOINT}/auth/cognito-jwt-payload`,
    event,
    {
      headers: {
        Authorization: apiKey,
        'content-type': 'application/json',
      },
    },
  );

  // API から受け取ったデータをそのままレスポンスに設定
  event.response = response.data;

  return event;
};

簡単に説明すると、
Cognitoから受け取った情報（Event）をそのままBackendのAPIに送信して、Backend側でこの情報（Event）を加工して返し、受け取ったLambdaはそのままResponseに詰めて終了するコードです。

移行後のコード（Rust）

セットアップ

まずはcargo lambdaをインストール
https://www.cargo-lambda.info/guide/installation.html

Lambda関数のテンプレートを作成
Pre token generation Lambda trigger用のLambdaを選択

$ cargo lambda new sample-rust-lambda 
> Is this function an HTTP function? No
? Event type that this function receives  
^ codepipeline_cloudwatch::CodePipelineInstanceEvent
  codepipeline_job::CodePipelineJobEvent
  cognito::CognitoEvent
> cognito::CognitoEventUserPoolsPreTokenGenV2
  config::ConfigEvent
  connect::ConnectEvent
v documentdb::DocumentDbEvent
[use arrows (↑↓) to move, tab to auto-complete, enter to submit.
Leave this input empty if you want to use a predefined example]

以下のプロジェクトが出来上がる

$ tree sample-rust-lambda 
sample-rust-lambda
├── Cargo.toml
├── README.md
└── src
    ├── http_handler.rs
    └── main.r

Lambdaの実装

ここでポイントがあるので簡単に解説しておきます。

  • TypeScriptとRustの型の厳密さが違うため、Cognitoから受け取った情報（Event）をそのまま渡す、そのままResponseに詰め直す、みたいなことができない
    • なので、ここの変換部分をしっかり実装し直す必要があった
    • （これぞ、静的型付け言語ですわねっ）
  • ログに関してはAWSの公式ドキュメントに記載されている通り、tracing_subscriberを導入
    • 構造化されたJSONログがほしかったので、導入しています
  • 今後、RustでLambdaが複数作成されていくことを予想して、workspaceを導入
    • ひとまず、SecretsManagerから取得する部分だけSharedパッケージとして定義
$ tree sample-rust-lambda 
sample-rust-lambda
├── Cargo.lock
├── Cargo.toml
├── Dockerfile
└── src
    ├── cognito-pre-token-generator
    │   ├── Cargo.toml
    │   └── src
    │       ├── event_handler.rs
    │       ├── helpers
    │       │   ├── create_new_response.rs
    │       │   ├── post_cognito_jwt_payload.rs
    │       │   └── post_cognito_jwt_payload_dto.rs
    │       ├── helpers.rs
    │       └── main.rs
    └── shared
        ├── Cargo.toml
        └── src
            ├── get_api_key.rs
            └── lib.rs
main.rs
use lambda_runtime::{run, service_fn, Error};

mod event_handler;
use event_handler::function_handler;

mod helpers;

#[tokio::main]
async fn main() -> Result<(), Error> {
    tracing_subscriber::fmt().json().init();

    run(service_fn(function_handler)).await
}

※ 参照している関数は割愛しています（リクエストあれば公開します笑）

event_handler.rs
use aws_lambda_events::event::cognito::CognitoEventUserPoolsPreTokenGenV2;
use lambda_runtime::{tracing, Error as LambdaError, LambdaEvent};
use shared::get_api_key::get_api_key;

use crate::helpers::{
    create_new_response::create_new_response, post_cognito_jwt_payload::post_cognito_jwt_payload,
};

pub(crate) async fn function_handler(
    event: LambdaEvent<CognitoEventUserPoolsPreTokenGenV2>,
) -> Result<CognitoEventUserPoolsPreTokenGenV2, LambdaError> {
    // リクエストのペイロードを取得
    let payload = event.payload;
    tracing::info!("Payload: {:?}", payload);

    let context = event.context;
    tracing::info!("Context: {:?}", context);

    // Backend の API を呼び出すために必要なキーをSecretManager等から取得
    let api_key_id = std::env::var("API_KEY_ID").unwrap();
    let api_key = get_api_key(api_key_id).await?;

    tracing::info!("User attributes: {:?}", payload.request.user_attributes);
    let user_name = payload.request.user_attributes.get("sub").unwrap();

    // Backend の API にリクエスト（serde_jsonで定義したオブジェクトが返ってくる）
    let api_endpoint = std::env::var("API_ENDPOINT").unwrap();
    let response = post_cognito_jwt_payload(api_endpoint, api_key, user_name.to_string()).await?;

    // 新しいレスポンスを作成（型を詰め直す）
    let new_response = create_new_response(response);

    let new_event = CognitoEventUserPoolsPreTokenGenV2 {
        cognito_event_user_pools_header: payload.cognito_event_user_pools_header,
        request: payload.request,
        response: new_response,
    };

    Ok(new_event)
}

CDKのコード

特別、変わったことはありませんが、せっかくなので記載しておきます。
今回はDockerImageから作成する方法でデプロイしました。

new DockerImageFunction(this, "PreTokenGenerationLambdaByRustFn", {
  code: DockerImageCode.fromImageAsset("lambda/lambda-handlers-by-rust"),
  functionName: "cognito-pre-token-generator3",
  description:
    "Generate custom claim for JWT before issue token. Including tenant id, organization id, dresscode user id",
  environment: {
    API_KEY_ID: "***",
    API_ENDPOINT: "***",
  },
  logRetention: RetentionDays.ONE_MONTH,
  memorySize: 128,
  timeout: Duration.seconds(8),
  tracing: Tracing.ACTIVE,
  architecture: Architecture.X86_64, // ARM64でもOK
});

Dockerfile

意外と手こずる部分があったので、Dockerfileも記載しておきます。
（バージョンとか少し古いです←他のサンプルをコピペしてます）

FROM --platform=linux/amd64 rust:1.83 AS builder

WORKDIR /tmp/mnt

COPY . .

RUN cargo build --release

FROM --platform=linux/amd64 public.ecr.aws/lambda/provided:al2023.2024.11.22.14 AS rust-lambda

COPY --from=builder \
  /tmp/mnt/target/release/cognito-pre-token-generator \
  ${LAMBDA_RUNTIME_DIR}/bootstrap

CMD [ "lambda-handler" ]

おわりに

Rustってプログラミング言語として実装してて楽しいんですよね〜

静的型付け言語としてメモリ安全性とパフォーマンスも担保されてるし、型の厳密さも守れている。
さらに、所有権やライフタイムといった独自概念も持ち合わせて、素晴らしい言語ですよね笑

関数型と合わせてRustを普及していきたい←その前にまずはRustマスターしなきゃw

DRESS CODE TECH BLOG
DRESS CODE TECH BLOGPublication

DRESS CODEのProduct & Technologyチームによるテックブログです！ プロダクトマネジメントからモデリング、アーキテクチャ、フロントエンド、バックエンド、SRE、セキュリティなどなど様々なテーマで情報を発信しています！

Discussion