feat: inject user agent in Lambda runtime API calls (#363)

nmoutschen · web-flow · commit c948939caf0d · 2021-11-15T07:46:56.000-08:00
diff --git a/lambda-runtime/src/requests.rs b/lambda-runtime/src/requests.rs
@@ -4,6 +4,8 @@ use hyper::Body;
 use serde::Serialize;
 use std::str::FromStr;
 
+const USER_AGENT: &str = concat!("aws-lambda-rust/", env!("CARGO_PKG_VERSION"));
+
 pub(crate) trait IntoRequest {
     fn into_req(self) -> Result<Request<Body>, Error>;
 }
@@ -20,6 +22,7 @@ impl IntoRequest for NextEventRequest {
     fn into_req(self) -> Result<Request<Body>, Error> {
         let req = Request::builder()
             .method(Method::GET)
+            .header("User-Agent", USER_AGENT)
             .uri(Uri::from_static("/2018-06-01/runtime/invocation/next"))
             .body(Body::empty())?;
         Ok(req)
@@ -57,6 +60,10 @@ fn test_next_event_request() {
     let req = req.into_req().unwrap();
     assert_eq!(req.method(), Method::GET);
     assert_eq!(req.uri(), &Uri::from_static("/2018-06-01/runtime/invocation/next"));
+    assert!(match req.headers().get("User-Agent") {
+        Some(header) => header.to_str().unwrap().starts_with("aws-lambda-rust/"),
+        None => false,
+    });
 }
 
 // /runtime/invocation/{AwsRequestId}/response
@@ -75,7 +82,11 @@ where
         let body = serde_json::to_vec(&self.body)?;
         let body = Body::from(body);
 
-        let req = Request::builder().method(Method::POST).uri(uri).body(body)?;
+        let req = Request::builder()
+            .header("User-Agent", USER_AGENT)
+            .method(Method::POST)
+            .uri(uri)
+            .body(body)?;
         Ok(req)
     }
 }
@@ -90,6 +101,10 @@ fn test_event_completion_request() {
     let expected = Uri::from_static("/2018-06-01/runtime/invocation/id/response");
     assert_eq!(req.method(), Method::POST);
     assert_eq!(req.uri(), &expected);
+    assert!(match req.headers().get("User-Agent") {
+        Some(header) => header.to_str().unwrap().starts_with("aws-lambda-rust/"),
+        None => false,
+    });
 }
 
 // /runtime/invocation/{AwsRequestId}/error
@@ -108,6 +123,7 @@ impl<'a> IntoRequest for EventErrorRequest<'a> {
         let req = Request::builder()
             .method(Method::POST)
             .uri(uri)
+            .header("User-Agent", USER_AGENT)
             .header("lambda-runtime-function-error-type", "unhandled")
             .body(body)?;
         Ok(req)
@@ -127,6 +143,10 @@ fn test_event_error_request() {
     let expected = Uri::from_static("/2018-06-01/runtime/invocation/id/error");
     assert_eq!(req.method(), Method::POST);
     assert_eq!(req.uri(), &expected);
+    assert!(match req.headers().get("User-Agent") {
+        Some(header) => header.to_str().unwrap().starts_with("aws-lambda-rust/"),
+        None => false,
+    });
 }
 
 // /runtime/init/error
@@ -140,6 +160,7 @@ impl IntoRequest for InitErrorRequest {
         let req = Request::builder()
             .method(Method::POST)
             .uri(uri)
+            .header("User-Agent", USER_AGENT)
             .header("lambda-runtime-function-error-type", "unhandled")
             .body(Body::empty())?;
         Ok(req)
@@ -153,4 +174,8 @@ fn test_init_error_request() {
     let expected = Uri::from_static("/2018-06-01/runtime/init/error");
     assert_eq!(req.method(), Method::POST);
     assert_eq!(req.uri(), &expected);
+    assert!(match req.headers().get("User-Agent") {
+        Some(header) => header.to_str().unwrap().starts_with("aws-lambda-rust/"),
+        None => false,
+    });
 }
