What is the idiomatic way to rewrite some standard JS lambda code in Rust?

[sn99]

Let's say I have the following JS code:

exports.handler = async function (event, context, callback) {
    const authorizer_context = event.requestContext.authorizer;
    const field1 = authorizer_context.field1;
    const field2 = authorizer_context.field2;
    
        let requestBody
        if (event.isBase64Encoded) {
            const bodyBuffer = Buffer.from(event.body, 'base64')
            requestBody = JSON.parse(bodyBuffer.toString('utf-8'))
        } else {
            requestBody = JSON.parse(event.body)
        }

       // Change requestBody to responseBody based on field1 and field2

        const response = {
                "isBase64Encoded": false,
                "statusCode": 200,
                "body": JSON.stringify(responseBody),
                "headers": {
                    'Access-Control-Allow-Origin': '*',
                    'some-header': field1,
                }
        }

    callback(null, response)
}

I have the following Rust conversion using Axum 0.6.x:

// required imports

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
    let cors = CorsLayer::new().allow_origin(Any);

    let state = Arc::new(AppState::new().await?);
    let app = Router::new()
        .route("/endpoint", post(my_fn))
        .layer(cors)
        .with_state(state);

    run(app).await
}

#[derive(Debug, Serialize, Deserialize)]
pub struct SomeRequest {
    body: String,
}

#[derive(Deserialize, Serialize, Debug)]
pub struct SomeResponse {
    pub body: String,
    pub headers: HeaderMap,
}

impl IntoResponse for SomeResponse {
    fn into_response(self) -> Response<String> {
        let mut response: Response<String> = Response::default();
        *response.headers_mut() = self.headers;
        *response.body_mut().push_str(&self.body);
        response
    }
}

pub async fn my_fn(
    State(state): State<Arc<AppState>>,
    field1: ExtractField1,
    field2: ExtractField2,
    Json(event): Json<SomeRequest>,
) -> Result<SomeResponse> {
     // Change event body based on ExtractField1 and ExtractField2
       
    let mut headers = HeaderMap::new();

    headers.insert(
        "some-header",
        field1.0,
    );
    
    Ok(SessionCreateResponse { body: event.body, headers })
}

macro_rules! from_request_parts {
    ($($struct:ident, $item:expr)*) => {
        $(
            #[async_trait]
            impl<S> FromRequestParts<S> for $struct
                where
                S: Send + Sync,
            {
                type Rejection = SomeError;

                async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
                    let request_context = parts
                        .extensions
                        .get::<RequestContext>()
                        .ok_or(SomeError::AuthorizerContextNotFound)?;

                    const ITEM: &str = $item;

                    let item_value = match request_context {
                        RequestContext::ApiGatewayV1(context) => context
                            .authorizer
                            .get(ITEM)
                            .ok_or(SomeError::AuthorizerContextNotFound),
                        RequestContext::ApiGatewayV2(context) => context
                            .authorizer
                            .as_ref()
                            .and_then(|authorizer| authorizer.lambda.get(ITEM))
                            .ok_or(SomeError::AuthorizerContextNotFound),
                        RequestContext::WebSocket(context) => context
                            .authorizer
                            .as_ref()
                            .and_then(|authorizer| authorizer.get(ITEM))
                            .ok_or(SomeError::AuthorizerContextNotFound),
                        RequestContext::Alb(_) => Err(SomeError::AuthorizerContextNotFound),
                    }?;

                    if let serde_json::Value::String(item_string) = item_value {
                        let item = item_string.to_string();
                        Ok(Self(item))
                    } else {
                        warn!("{} is not a string", ITEM);
                        Err(SomeError::AuthorizerContextNotFound)
                    }
                }
            }
        )*
    }
}

#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct ExtractField1(pub(crate) String);

#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct ExtractField2(pub(crate) String);

from_request_parts!(
    ExtractField1, "field1"
    ExtractField2, "field2"
);

Is there a better way to rewrite this? Can I include field1/2 in RequestBody itself or just keep the header as part of the struct instead of making IntoResponse ? Do I need to specify isBase64Encoded or is it automatically handled, both for reading and writing?

awslabs/aws-sdk-rust#1069

[calavera]

I don't know if there is an easier way, but we should probably make it easier to extract things from the context. We can always build a common interface so you don't have to match all the payloads.

[sn99]

My other question will then be about soundness of this code, would the above be correct rewrite?

[calavera]

I've made some improvements on how you can get the authorizer in the main branch. The changes are not in any release yet, but you can point your dependencies to the GitHub repository if you want to use them.

Your code is sound, but it's probably too complicated to maintain. Take a look at the examples in https://github.com/awslabs/aws-lambda-rust-runtime/blob/main/examples/http-axum-apigw-authorizer/src/main.rs

Do I need to specify isBase64Encoded or is it automatically handled, both for reading and writing?

it's only handled for writing.

You don't need to implement IntoResponse unless you want a very specific response object. You can also return a http::Response directly. Something like this works too:

fn axum_handler() -> http::Result<http::Response> {
  http::Response::builder()
    .header("som-header", "some-value")
    .body("response body")
}

[calavera]

btw, I don't know if you're using Axum for any specific reason, but you don't need Axum to process http requests.

[sn99]

That makes me ask another question:

Do I also have to encode decode the same way that js does?

        if (event.isBase64Encoded) {
            const bodyBuffer = Buffer.from(event.body, 'base64')
            requestBody = JSON.parse(bodyBuffer.toString('utf-8'))
        } else {
            requestBody = JSON.parse(event.body)
        }

[calavera]

It has been so long that I didn't even remember how this works. You don't have to decode the payload manually if your receiving a Request, the runtime does it for you. You can see the transformation here: https://docs.rs/lambda_http/latest/lambda_http/ext/request/trait.RequestPayloadExt.html#examples

If your request body is a JSON object, you can also turn it into a structure directly, here's an example: https://docs.rs/lambda_http/latest/lambda_http/ext/request/trait.RequestPayloadExt.html#examples

Basically you can do something like this (I have not really tested this code):

use RequestPayloadExt;

#[derive(Default, Deserialize, Serialize)]
struct JsonPayload {
  field_1: String,
  field_2: String,
}

fn axum_handler(request: http::Request) -> http::Result<http::Response> {
  let request_body: JsonPayload = request.payload()
    .unwrap_or_else(|_parse_err| None)
    .unwrap_or_default();

    http::Response::builder()
    .header("som-header", "some-value")
    .body(serde_json::to_string(&request_body))
}