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

sn99 · 2024-02-19T11:34:26Z

sn99
Feb 19, 2024

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 ?