RustのカスタムLambdaランタイムでバイナリレスポンスを扱う

2022-10-18

thumbnail

このブログ投稿ではRustで書かれたAWS Lambdaカスタムランタイムでバイナリレスポンスを扱う方法を紹介します。

背景

Lambdaインテグレーションを使えばAmazon API Gateway REST API (REST API)から動的なレスポンスを返すことができます。 私はREST API用に動的なバイナリデータを生成するAWS Lambda (Lambda)関数を実装することにしました。 Rustを学習中なので、Lambda関数をRustで実装することにしました*。 このブログはNon-Proxyインテグレーション用のLambda関数を扱いますのでご留意ください。

* 私はPythonとNode.js (JavaScript)はバイナリデータを処理するのはあまり得意でないだろうと考えています。 Goはバイナリデータ処理によい選択かもしれず、AWSも公式なGoのLambdaランタイムを提供していますが、とにかくRustを学びたいわけです。

Rust用のLambdaランタイム

AWSはRustに熱心な様子ですが、公式なRust用のLambdaランタイムは提供していません。 ということでカスタムLambdaランタイムをRustで実装するか、コンテナを作成するかしなければなりません。 Lambda関数をRustでどうやって実装するかを探していると、おそらくaws-lambda-rust-runtimeというライブラリに出くわします。 このライブラリはRust用のカスタムLambdaランタイムを実装する際に面倒なことを代わりに全部やってくれます。

aws-lambda-rust-runtimeのおかげで、Lambda関数をRustで実装するのはかなり簡単です。 aws-lambda-rust-runtimeのGitHubレポジトリに簡単なチュートリアルがあります。 チュートリアルをここでは繰り返しませんが、抜粋(実際に記述するRustコード)を以下に示します。

use lambda_runtime::{service_fn, LambdaEvent, Error};
use serde_json::{json, Value};

#[tokio::main]
async fn main() -> Result<(), Error> {
    let func = service_fn(func);
    lambda_runtime::run(func).await?;
    Ok(())
}

async fn func(event: LambdaEvent<Value>) -> Result<Value, Error> {
    let (event, _context) = event.into_parts();
    let first_name = event["firstName"].as_str().unwrap_or("world");

    Ok(json!({ "message": format!("Hello, {}!", first_name) }))
}

aws-lambda-rust-runtimeでバイナリデータを扱う

aws-lambda-rust-runtimeは出力がJSONにエンコード可能な限りうまく機能します。

Vec<u8>を返すのはうまくいかない

コアモジュールlambda_runtimeのrun関数はserde::Serializeを実装する値を返すサービスならなんでも受け付けます。 バイナリを出力するのに、まず私はサービス関数にVec<u8>を返させれば良いと考えてやってみました。 すると、lambda_runtimeはBLOBの代わりに配列のJSON形式を生成しました。 例えば、[0x61, 0x62, 0x63]をVec<u8>としてlambda_runtimeに与えると、以下のような出力を、

[
  97,
  98,
  99
]

abc*の代わりに得ました。

* 0x61, 0x62, 0x63はASCIIでそれぞれ'a', 'b', 'c'を表します。

lambda_runtimeはどのように結果を扱っているのか?

lambda_runtime::runはあらゆるサービス関数の出力をJSONに変換しますが、この動作はlambda-runtime/src/requests.rs#L77-L85にハードコードされています(serde_json::to_vec(&self.body)に注目)。

    fn into_req(self) -> Result<Request<Body>, Error> {
        let uri = format!("/2018-06-01/runtime/invocation/{}/response", self.request_id);
        let uri = Uri::from_str(&uri)?;
        let body = serde_json::to_vec(&self.body)?;
        let body = Body::from(body);

        let req = build_request().method(Method::POST).uri(uri).body(body)?;
        Ok(req)
    }

解決策?

バイナリレスポンスに関する議論がGitHubに見つかりましたが、それはlambda_httpモジュールを使うことを提案しています。 私が調べた限り、lambda_httpはLambdaのProxyインテグレーション用に設計されています。 なので単純にBLOBを生成することはできません。

そこで以下の2つの回避策を思いつきました。

  1. Base64エンコードされたバイナリをJSONオブジェクトのフィールド値として埋め込み、インテグレーションレスポンスのマッピングテンプレートでそれを取り出し、CONVERT_TO_BINARYを適用 (この方法も単純なBLOBを生成しません。)
  2. lambda_runtimeをバイナリ出力を扱えるように改造

簡単な方法は最初のものだったはずですが、 Rustを練習するよい機会だったので2番目の方法*を採りました。

* 後に、私の努力は必要なかったことが判明します。 もっと単純で簡単な方法を知りたい方は節「もっと単純な解決策」まで読み飛ばしてもらって結構です。

lambda_runtimeの改造

改造の主な目的はバイナリレスポンスをサポートすることですが、2番目の目的も設定しました。JSONシリアライズ可能なオブジェクトを返す既存のプログラムも確実に動作し続けるということ(後方互換性)です。

この節では以下を紹介します。

  1. 関数のレスポンスの出口となるIntoRequest Trait
  2. 私の回避策: IntoBytesとRawBytesの導入

IntoRequest Trait

IntoRequest(lambda-runtime/src/requests.rs#L8-L10)はサービス関数のすべての成功出力をシリアライズし、Lambda関数の呼び出しレスポンス用エンドポイント*に送るリクエストオブジェクトを作成します。

pub(crate) trait IntoRequest {
    fn into_req(self) -> Result<Request<Body>, Error>;
}

IntoRequestはlambda_runtimeにおいて重要な役割を果たしていますが、関数の出力用に皆さんが直接実装*2 することはありません。 関数の結果とIntoRequestの間には「ブリッジ」があります。EventCompletionRequest<T> (Tは関数が出力する型)です(lambda-runtime/src/requests.rs#L68-L71)。

pub(crate) struct EventCompletionRequest<'a, T> {
    pub(crate) request_id: &'a str,
    pub(crate) body: T,
}

lambda_runtimeは関数の出力をEventCompletionRequest<T>でラップしIntoRequestとして処理します(lambda-runtime/src/lib.rs#L164-L168)。

                                EventCompletionRequest {
                                    request_id,
                                    body: response,
                                }
                                .into_req()

IntoRequestはTがserde::SerializeであるようなEventCompletionRequest<T>に対して実装されています(lambda-runtime/src/requests.rs#L73-L75)。

impl<'a, T> IntoRequest for EventCompletionRequest<'a, T>
where
    T: for<'serialize> Serialize,

これがサービス関数が出力するSerializeがJSONオブジェクトに変換される理由です。

* IntoRequestが関数のレスポンスを表しているのに名前に"Request"という単語が含まれていて混乱するかもしれませんが、ここでの"Request"はカスタムLambdaランタイムの呼び出しレスポンス用エンドポイントに送るリクエストという意味です。

*2 そもそもIntoRequestはエクスポートされていません。

IntoBytesとRawBytesの導入

IntoRequest::into_reqの中のserde_json::to_vec呼び出しをどうにかして一般化しなければなりません。 これはサービス出力からバイト列(Vec<u8>)への変換と捉えることができます。 ならばIntoRequestをEventCompletionRequest<'a, Vec<u8>>用に特殊化してはどうでしょうか? 残念ながら、lambda_runtime::runをSerializeとVec<u8>の両方をサービス出力として受け付けるようにすることはできないのでうまくいきません。

pub async fn run<A, B, F>(handler: F) -> Result<(), Error>
where
    F: Service<LambdaEvent<A>>,
    F::Future: Future<Output = Result<B, F::Error>>,
    F::Error: fmt::Debug + fmt::Display,
    A: for<'de> Deserialize<'de>,
    B: Serialize | Vec<u8>, // エラー: こういうことはできない!

Serializeにも生のバイト列にも解釈できるような新しい型が必要です。 では、新しいIntoBytes Traitを導入し以下のとおりIntoRequestをEventCompletionRequest<'a, Serialize>ではなくEventCompletionRequest<'a, IntoBytes>用に特殊化してはどうでしょうか?

pub trait IntoBytes {
    fn to_bytes(self) -> Result<Vec<u8>, Error>;
}

impl<'a, T> IntoRequest for EventCompletionRequest<'a, T>
where
    T: IntoBytes,
{
    fn into_req(self) -> Result<Request<Body>, Error> {
        let uri = format!("/2018-06-01/runtime/invocation/{}/response", self.request_id);
        let uri = Uri::from_str(&uri)?;
        let body = self.body.to_bytes()?;
        let body = Body::from(body);

        let req = build_request().method(Method::POST).uri(uri).body(body)?;
        Ok(req)
    }
}

すると、lambda_runtime::runのシグネチャも書き換えなければなりません。 以下はどうでしょうか?

pub async fn run<A, B, F>(handler: F) -> Result<(), Error>
where
    F: Service<LambdaEvent<A>>,
    F::Future: Future<Output = Result<B, F::Error>>,
    F::Error: fmt::Debug + fmt::Display,
    A: for<'de> Deserialize<'de>,
    B: IntoBytes, // Serializeを受け付けない

上記のようにすると、後方互換性が失われます(サービス関数は単純にSerializeを返すことができなくなります)。 この問題を回避するため、IntoBytesBridgeという別のStructを導入しました。

pub struct IntoBytesBridge<T>(pub T);

IntoBytesは以下のようにIntoBytesBridge<Serialize>に対して特殊化します。

impl<T> IntoBytes for IntoBytesBridge<T>
where
    T: Serialize,
{
    fn to_bytes(self) -> Result<Vec<u8>, Error> {
        let bytes = serde_json::to_vec(&self.0)?;
        Ok(bytes)
    }
}

IntoBytesBridgeのおかげで、lambda_runtime::runのシグネチャは次のように書き換えることができます。

pub async fn run<A, B, F>(handler: F) -> Result<(), Error>
where
    F: Service<LambdaEvent<A>>,
    F::Future: Future<Output = Result<B, F::Error>>,
    F::Error: fmt::Debug + fmt::Display,
    A: for<'de> Deserialize<'de>,
    IntoBytesBridge<B>: IntoBytes, // IntoBytesをIntoBytesBridge<B>に対して実装していることを要求する

これはSerializeに対しても機能します。IntoBytesがIntoBytesBridge<Serialize>に対して実装されているからです(上述のコード参照)。

今度は、生のバイト列を出力するという意図を伝えるために新しいデータ型RawBytesを導入し、IntoBytesをIntoBytesBridge<RawBytes>に対して特殊化します。

pub struct RawBytes(pub Vec<u8>);

impl IntoBytes for IntoBytesBridge<RawBytes> {
    fn to_bytes(self) -> Result<Vec<u8>, Error> {
        Ok(self.0.0)
    }
}

これでサービス関数の出力をRawBytesでラップすれば生のバイト列を出力することができます。 以下の例はJSONの配列[97, 98, 99]ではなく生の文字列abcを出力します。

use lambda_runtime::{service_fn, LambdaEvent, Error, RawBytes};
use serde_json::Value;

#[tokio::main]
async fn main() -> Result<(), Error> {
    let func = service_fn(func);
    lambda_runtime::run(func).await?;
    Ok(())
}

async fn func(event: LambdaEvent<Value>) -> Result<RawBytes, Error> {
    Ok(RawBytes(vec![0x61, 0x62, 0x63]))
}

Lambdaインテグレーションの制限?

前節で開発した新機能を自分のREST APIを実装するのに試してみましたが、おかしなことが起こりました。 サービスレスポンスとして返したバイト列とは少し違うバイト列をAPIが出力してしまうのです。 APIがときどき特定の3バイト(0xE, 0xBF, 0xBD)を生成していることに気づきました。 問題をしっかり観察すると、その3バイトは最上位ビットが1の(言い換えると、0x80とのビットANDが0でない)バイトを置き換えていることが分かりました。 (0xEF, 0xBF, 0xBD)という3バイトは実際のところ置換文字(U+FFFD)のUTF-8表現であり、それはつまり、有効なUTF-8列を期待している誰かが私のAPI出力の不明なバイトを置き換えたということです。 では誰に責任があるのでしょうか? Lambdaの呼び出しレスポンス用のエンドポイント? それともAmazon API GatewayのLambdaインテグレーション? それともまさかRustの依存関係?

AWSのCLIコマンドaws lambda invokeでLambda関数を直接呼び出したときには、私が期待するバイト列が得られました。 ということはLambdaの呼び出しレスポンスのエンドポイントとRustの依存関係は無実で、Lambdaインテグレーションが原因ということです。 LambdaのProxyインテグレーションが出力にJSON(有効なUTF-8列の部分集合)を期待しているのは知っていましたが、 LambdaのNon-Proxyインテグレーションも出力にUTF-8列を期待している*とは知りませんでした。

* 公式な情報源ではまだ確認できていません。

回避策

私の回避策はサービス出力をBase64にエンコードしcontentHandlingにCONVERT_TO_BINARYを指定するというものです。 しかしひとつ疑問が湧きます。バイナリをBase64テキストにエンコードしなければならないのなら、lambda_runtimeの拡張は本当に必要だったのでしょうか? この疑問は次節で解説するもっと単純な解決策につながります。

もっと単純な解決策

前節では、サービスが任意のバイト列を出力することをLambdaのNon-Proxyインテグレーションが許さないことを学びました。 なので、結局サービス出力はBase64エンコードしなければならないわけです。 ということはlambda_runtimeを改造したことのメリットはほとんど失われたことになります。 節「解決策?」では、lambda_runtimeに対する改造を必要としない別の方法を提案しました。

  1. Base64エンコードされたバイナリをJSONオブジェクトのフィールド値として埋め込み、インテグレーションレスポンスのマッピングテンプレートでそれを取り出し、CONVERT_TO_BINARYを適用

しかし、フィールドの値を取り出すためだけにマッピングテンプレートを書くのも少し面倒だと思います。 では、Base64エンコードしたStringを直接出力し、それをJSONにシリアライズするというのはどうでしょうか? SerializeはStringに対して実装されているので、サービス関数はなんの改造もなくStringを出力できます。

use lambda_runtime::{service_fn, LambdaEvent, Error};
use serde_json::Value;

#[tokio::main]
async fn main() -> Result<(), Error> {
    let func = service_fn(func);
    lambda_runtime::run(func).await?;
    Ok(())
}

async fn func(event: LambdaEvent<Value>) -> Result<String, Error> {
    Ok(base64::encode([0x61, 0x62, 0x63]))
}

私は最初これはダメだろうなと思っていました。なぜなら出力が余計なダブルクオーテーション(")で囲まれてしまい無効なBase64テキストになってしまうからです。

ところが、うまくいったのです!

出力をaws lambda invokeで取り出すと、実際にダブルクオーテーションで取り囲まれていました。 しかしLambdaインテグレーションはどうしてかその扱い方を認識しており、所望のBase64テキストを正しくデコードしてくれたのです。

まとめ

このブログでは、aws-lambda-rust-runtimeがRust用のカスタムLambdaランタイムを実装するのにとても役立つことを確認しました。 それからaws-lambda-rust-runtimeがバイナリデータを扱えるようにするための改造を紹介しました。 しかし、Base64エンコードしたStringを単純に返すのが、Amazon API GatewayのLambdaインテグレーションでバイナリ出力を扱うのに最も簡単な方法だということが分かりました。

あまり役に立たなくなってしまいましたが、私がaws-lambda-rust-runtimeに行った改造はGitHubのフォークで確認できます。

補足

CDKでRustのLambdaランタイムをビルドする

AWS Cloud Development Kit (CDK)でRustのLambda関数をビルドするのにrust.aws-cdk-lambdaというモジュールがとても役に立ちました。

ツイートする