Improve extensions API support

- Set request content-types for JSON payloads - Print the returned status codes in case of errors Signed-off-by: David Calavera <david.calavera@gmail.com>
awslabs · Feb 1, 2023 · 2f322ae · 2f322ae
1 parent cccce9d
commit 2f322ae
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/lambda-extension/src/extension.rs b/lambda-extension/src/extension.rs
@@ -253,7 +253,8 @@ where
             )?;
             let res = client.call(req).await?;
             if res.status() != http::StatusCode::OK {
-                return Err(ExtensionError::boxed("unable to initialize the logs api"));
+                let err = format!("unable to initialize the logs api: {}", res.status());
+                return Err(ExtensionError::boxed(err));
             }
             trace!("Registered extension with Logs API");
         }
@@ -288,7 +289,8 @@ where
             )?;
             let res = client.call(req).await?;
             if res.status() != http::StatusCode::OK {
-                return Err(ExtensionError::boxed("unable to initialize the telemetry api"));
+                let err = format!("unable to initialize the telemetry api: {}", res.status());
+                return Err(ExtensionError::boxed(err));
             }
             trace!("Registered extension with Telemetry API");
         }
@@ -422,7 +424,8 @@ async fn register<'a>(
     let req = requests::register_request(&name, events)?;
     let res = client.call(req).await?;
     if res.status() != http::StatusCode::OK {
-        return Err(ExtensionError::boxed("unable to register the extension"));
+        let err = format!("unable to register the extension: {}", res.status());
+        return Err(ExtensionError::boxed(err));
     }
 
     let header = res

diff --git a/lambda-extension/src/requests.rs b/lambda-extension/src/requests.rs
@@ -7,6 +7,8 @@ use serde::Serialize;
 const EXTENSION_NAME_HEADER: &str = "Lambda-Extension-Name";
 pub(crate) const EXTENSION_ID_HEADER: &str = "Lambda-Extension-Identifier";
 const EXTENSION_ERROR_TYPE_HEADER: &str = "Lambda-Extension-Function-Error-Type";
+const CONTENT_TYPE_HEADER_NAME: &str = "Content-Type";
+const CONTENT_TYPE_HEADER_VALUE: &str = "application/json";
 
 pub(crate) fn next_event_request(extension_id: &str) -> Result<Request<Body>, Error> {
     let req = build_request()
@@ -24,6 +26,7 @@ pub(crate) fn register_request(extension_name: &str, events: &[&str]) -> Result<
         .method(Method::POST)
         .uri("/2020-01-01/extension/register")
         .header(EXTENSION_NAME_HEADER, extension_name)
+        .header(CONTENT_TYPE_HEADER_NAME, CONTENT_TYPE_HEADER_VALUE)
         .body(Body::from(serde_json::to_string(&events)?))?;
 
     Ok(req)
@@ -73,6 +76,7 @@ pub(crate) fn subscribe_request(
         .method(Method::PUT)
         .uri(api.uri())
         .header(EXTENSION_ID_HEADER, extension_id)
+        .header(CONTENT_TYPE_HEADER_NAME, CONTENT_TYPE_HEADER_VALUE)
         .body(Body::from(serde_json::to_string(&data)?))?;
 
     Ok(req)