Repository: awslabs/aws-lambda-rust-runtime (language: rust)

Issue:
Lambda for Cognito Post Confirmation event response causes error for Account Creation
Hi!
I recently implemented a lambda for AWS Cognito integration with the [user-pool-lambda-post-confirmation](https://docs.aws.amazon.com/cognito/latest/developerguide/user-pool-lambda-post-confirmation.html#aws-lambda-triggers-post-confirmation-example) event.

The sdk event for that is this: [event/cognito/mod.rs#L87](https://github.com/awslabs/aws-lambda-rust-runtime/blob/main/lambda-events/src/event/cognito/mod.rs#L87). 

Looks like this event is not correct. I think maybe the docs are out of date? The docs page show an empty json for Response ([html#cognito-user-pools-lambda-trigger-syntax-post-confirmation](https://docs.aws.amazon.com/cognito/latest/developerguide/user-pool-lambda-post-confirmation.html#cognito-user-pools-lambda-trigger-syntax-post-confirmation)).
```
{ "request": { "userAttributes": {"string": "string",. . .},"clientMetadata": {"string": "string",. . .}},"response": {}}
```

The [CognitoEventUserPoolsPostConfirmationResponse](https://github.com/awslabs/aws-lambda-rust-runtime/blob/main/lambda-events/src/event/cognito/mod.rs#L256) is also an empty object.

We found when logging into the app via socials for the first time (login implemented with Amplify SDK for iOS), an error is thrown and login fails. The error message mentions the lambda
```
Location:applocation://?error_description=Invalid+version+in+Lambda+response.+Version+should+be+1+&state=e3811fde-b3fe-4507-9992-fa2b228ed194&error=invalid_request
```

The JS example in the lower part of the docs shows simply returning the input even which contains the version field. After updating the code to return the event user pools header (I saw it had this same version field), it worked.
```rust
type EventResponse = CognitoEventUserPoolsHeader<CognitoEventUserPoolsPostConfirmationTriggerSource>;


async fn function_handler(
    event: LambdaEvent<CognitoEventUserPoolsPostConfirmation>
) -> Result<EventResponse> {
   let payload = event.payload;
   let event_response = payload.cognito_event_user_pools_header;
   .... // details
   Ok(event_response)
}
```


Discussion / hints:
I checked with the Cognito team and it looks like the documentation might not be super clear. Cognito expects that the payload that the lambda invocation returns is similar to the payload that it receives, this is key in the docs:

> Amazon Cognito passes event information to your Lambda function. The function then returns the same event object to Amazon Cognito, with any changes in the response.

As you point out, the example in the doc returns the same event payload that it receives.

I think this could be an opportunity for us to have some examples of how to work with Cognito and Rust Lambda functions, I'm not an expert on it though, but any help will be very appreciated.

Fix patch:
diff --git a/examples/basic-cognito-post-confirmation/.gitignore b/examples/basic-cognito-post-confirmation/.gitignore
new file mode 100644
index 00000000..c41cc9e3
--- /dev/null
+++ b/examples/basic-cognito-post-confirmation/.gitignore
@@ -0,0 +1,1 @@
+/target
\ No newline at end of file
diff --git a/examples/basic-cognito-post-confirmation/Cargo.toml b/examples/basic-cognito-post-confirmation/Cargo.toml
new file mode 100644
index 00000000..7d2e7ab4
--- /dev/null
+++ b/examples/basic-cognito-post-confirmation/Cargo.toml
@@ -0,0 +1,24 @@
+[package]
+name = "basic-cognito-post-confirmation"
+version = "0.1.0"
+edition = "2021"
+
+# Starting in Rust 1.62 you can use `cargo add` to add dependencies 
+# to your project.
+#
+# If you're using an older Rust version,
+# download cargo-edit(https://github.com/killercup/cargo-edit#installation) 
+# to install the `add` subcommand.
+#
+# Running `cargo add DEPENDENCY_NAME` will
+# add the latest version of a dependency to the list,
+# and it will keep the alphabetic ordering for you.
+
+[dependencies]
+aws-config = "1.5.0"
+aws-sdk-ses = "1.28.0"
+aws_lambda_events = { path = "../../lambda-events", default-features = false, features = ["cognito"] }
+
+lambda_runtime = { path = "../../lambda-runtime" }
+tokio = { version = "1", features = ["macros"] }
+
diff --git a/examples/basic-cognito-post-confirmation/README.md b/examples/basic-cognito-post-confirmation/README.md
new file mode 100644
index 00000000..4391e16c
--- /dev/null
+++ b/examples/basic-cognito-post-confirmation/README.md
@@ -0,0 +1,15 @@
+# Cognito Post Confirmation Request example
+
+This example shows how to write a Lambda function in Rust to process Cognito's Post Confirmation requests.
+
+This is a translation of the example in the AWS Docs to Rust: https://docs.aws.amazon.com/cognito/latest/developerguide/user-pool-lambda-post-confirmation.html#aws-lambda-triggers-post-confirmation-example
+
+## Build & Deploy
+
+1. Install [cargo-lambda](https://github.com/cargo-lambda/cargo-lambda#installation)
+2. Build the function with `cargo lambda build --release`
+3. Deploy the function to AWS Lambda with `cargo lambda deploy`
+
+## Build for ARM 64
+
+Build the function with `cargo lambda build --release --arm64`
diff --git a/examples/basic-cognito-post-confirmation/src/main.rs b/examples/basic-cognito-post-confirmation/src/main.rs
new file mode 100644
index 00000000..178bd3a9
--- /dev/null
+++ b/examples/basic-cognito-post-confirmation/src/main.rs
@@ -0,0 +1,60 @@
+use aws_config::BehaviorVersion;
+use aws_lambda_events::event::cognito::CognitoEventUserPoolsPostConfirmation;
+use aws_sdk_ses::{
+    types::{Body, Content, Destination, Message},
+    Client,
+};
+use lambda_runtime::{run, service_fn, tracing, Error, LambdaEvent};
+
+const SOURCE_EMAIL: &str = "<source_email>";
+
+async fn function_handler(
+    client: &aws_sdk_ses::Client,
+    event: LambdaEvent<CognitoEventUserPoolsPostConfirmation>,
+) -> Result<CognitoEventUserPoolsPostConfirmation, Error> {
+    let payload = event.payload;
+
+    if let Some(email) = payload.request.user_attributes.get("email") {
+        let body = if let Some(name) = payload.request.user_attributes.get("name") {
+            format!("Welcome {name}, you have been confirmed.")
+        } else {
+            "Welcome, you have been confirmed.".to_string()
+        };
+        send_post_confirmation_email(client, email, "Cognito Identity Provider registration completed", &body).await?;
+    }
+
+    // Cognito always expect a response with the same shape as
+    // the event when it handles Post Confirmation triggers.
+    Ok(payload)
+}
+
+async fn send_post_confirmation_email(client: &Client, email: &str, subject: &str, body: &str) -> Result<(), Error> {
+    let destination = Destination::builder().to_addresses(email).build();
+    let subject = Content::builder().data(subject).build()?;
+    let body = Content::builder().data(body).build()?;
+
+    let message = Message::builder()
+        .body(Body::builder().text(body).build())
+        .subject(subject)
+        .build();
+
+    client
+        .send_email()
+        .source(SOURCE_EMAIL)
+        .destination(destination)
+        .message(message)
+        .send()
+        .await?;
+
+    Ok(())
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Error> {
+    tracing::init_default_subscriber();
+
+    let config = aws_config::load_defaults(BehaviorVersion::latest()).await;
+    let client = Client::new(&config);
+
+    run(service_fn(|event| function_handler(&client, event))).await
+}
diff --git a/lambda-events/src/event/cognito/mod.rs b/lambda-events/src/event/cognito/mod.rs
index 6eb1c001..1b31abcc 100644
--- a/lambda-events/src/event/cognito/mod.rs
+++ b/lambda-events/src/event/cognito/mod.rs
@@ -84,13 +84,18 @@ pub enum CognitoEventUserPoolsPreAuthenticationTriggerSource {
 /// allowing the Lambda to send custom messages or add custom logic.
 #[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
 #[serde(rename_all = "camelCase")]
-pub struct CognitoEventUserPoolsPostConfirmation {
+pub struct CognitoEventUserPoolsPostConfirmation<T = CognitoEventUserPoolsPostConfirmationResponse>
+where
+    T: DeserializeOwned,
+    T: Serialize,
+{
     #[serde(rename = "CognitoEventUserPoolsHeader")]
     #[serde(flatten)]
     pub cognito_event_user_pools_header:
         CognitoEventUserPoolsHeader<CognitoEventUserPoolsPostConfirmationTriggerSource>,
     pub request: CognitoEventUserPoolsPostConfirmationRequest,
-    pub response: CognitoEventUserPoolsPostConfirmationResponse,
+    #[serde(bound = "")]
+    pub response: T,
 }
 
 #[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize, Default)]
@@ -254,6 +259,7 @@ pub struct CognitoEventUserPoolsPostConfirmationRequest {
 /// `CognitoEventUserPoolsPostConfirmationResponse` contains the response portion of a PostConfirmation event
 #[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
 pub struct CognitoEventUserPoolsPostConfirmationResponse {}
+
 /// `CognitoEventUserPoolsPreTokenGenRequest` contains request portion of PreTokenGen event
 #[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
 #[serde(rename_all = "camelCase")]


Test patch: