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

Issue:
Allow any path parts in `AWS_LAMBDA_RUNTIME_API` to remain intact
Hey folks, this project is great! I have been at work working on implementing full Lambda custom runtime API support in OpenJS Architect, and wanted to raise some behavior that appears to be specific to the Rust Lambda runtime library.

A goal of Architect's local execution model is to be as fast (or faster) than real Lambda; as such, we do not rely on Docker. To accomplish this, we run handler functions in short-lived child processes on the local machine and do our own emulation (instead of relying on a lambCI or SAM Docker images. Since each invocation shares the same host ports, this means we cannot run the local RIE server, as there would be port conflict.

So what I've been working on is building our own simple take on the Lambda runtime API, but shared across multiple Lambdas – but there’s a bit of a problem. The current runtime API design assumes it is not necessary for the Next invocation request to need any identification, as the each Lambda in question is contacting `localhost:9001` – so knowledge of which Lambda is getting the invocation data is implied. In our case, again, the Lambda runtime API will be shared across multiple Lambda invocations (and we'll be responsible for maintaining state), so I need a way to identify which Lambda is requesting the invocation.

I had hoped to set `AWS_LAMBDA_RUNTIME_API` to `localhost:{port}/{lambdaID}/`, resulting in the Lambda runtime API client hitting `/{lambdaID}/2018-06-01/runtime/invocation/next`. I believe this [should work in the Go Lambda client](https://github.com/aws/aws-lambda-go/blob/main/lambda/runtime_api_client.go#L39), for example. Unfortunately, the implementation in the Rust client as it stands today strips out everything except the host name and port, ensuring that every request is only ever sent to `/2018-06-01/runtime/invocation/next`. And without headers or query params, I have no way of identifying which Lambda is requesting its next invocation.

It would be great if you could follow the Lambda Go runtime approach, leaving the contents of the environment variable intact.

## Repro steps
- Build the [basic example function](https://github.com/awslabs/aws-lambda-rust-runtime#example-function)
- Set up a simple HTTP server on port `9001` to simply print request metadata
- Execute a Lambda with `AWS_LAMBDA_RUNTIME_API` set to `localhost:9001/foo/`
- `GET /2018-06-01/runtime/invocation/next` will be fired, and not `GET /foo/2018-06-01/runtime/invocation/next`


Discussion / hints:
Thanks for the detailed description, Ryan! Also your project looks really cool! Hoping to read more about it when I get some free time.

As per your issue, this seems reasonable to me. As far as I can tell the reason we're not using the path is that we kind of just ignore it currently. If you look [here](https://github.com/awslabs/aws-lambda-rust-runtime/blob/6da8fe9d08e5f24b1790efeb45100b0cd1b6caed/lambda-runtime-api-client/src/lib.rs#L61-L66) you'll see that we get the `scheme` and `authority` from the base uri, but then ignore the `path_and_query` which is then just used from the request. I think we would just need to merge the two paths, but I'm not quite sure if this is the proper behavior. I see that this would probably work on the Go client, but I think it could potentially lead to some wonky stuff. 

Let me read through some documentation to verify this isn't dangerous behavior. If not I think this would be simple enough to implement.
Thanks! Yeah, I noticed that code you linked while poking around. I'm not sure about query params – I could make use of those as an alternative to the path parts, but ultimately sanitizing query params is probably a good thing as not doing so could potentially lead to some extremely weird behavior. I hope that otherwise passing through the path parts should be safe and low/no impact, as locally it would be assumed that one controls the runtime API endpoints anyway.

A bit of a side question (but semi-related) that you may or may not know: in production Lambda, what is providing service on `localhost:9001`, and by what means is it supplied with the triggering event prior to the bootstrap hitting the `next` endpoint? (I feel like in one of the SAM Docker images I saw an RIE build on the filesystem while poking around.) This isn't necessary for us to resolve this issue, just esoterica I've noodling on since working on this implementation.

Fix patch:
diff --git a/lambda-runtime-api-client/src/lib.rs b/lambda-runtime-api-client/src/lib.rs
index 9703ecd2..83d789cc 100644
--- a/lambda-runtime-api-client/src/lib.rs
+++ b/lambda-runtime-api-client/src/lib.rs
@@ -4,7 +4,7 @@
 
 //! This crate includes a base HTTP client to interact with
 //! the AWS Lambda Runtime API.
-use http::{uri::Scheme, Request, Response, Uri};
+use http::{uri::PathAndQuery, uri::Scheme, Request, Response, Uri};
 use hyper::{
     client::{connect::Connection, HttpConnector},
     Body,
@@ -58,26 +58,24 @@ where
 
     fn set_origin<B>(&self, req: Request<B>) -> Result<Request<B>, Error> {
         let (mut parts, body) = req.into_parts();
-        let (scheme, authority) = {
+        let (scheme, authority, base_path) = {
             let scheme = self.base.scheme().unwrap_or(&Scheme::HTTP);
             let authority = self.base.authority().expect("Authority not found");
-            (scheme, authority)
+            let base_path = self.base.path().trim_end_matches('/');
+            (scheme, authority, base_path)
         };
         let path = parts.uri.path_and_query().expect("PathAndQuery not found");
+        let pq: PathAndQuery = format!("{}{}", base_path, path).parse().expect("PathAndQuery invalid");
 
         let uri = Uri::builder()
-            .scheme(scheme.clone())
-            .authority(authority.clone())
-            .path_and_query(path.clone())
-            .build();
-
-        match uri {
-            Ok(u) => {
-                parts.uri = u;
-                Ok(Request::from_parts(parts, body))
-            }
-            Err(e) => Err(Box::new(e)),
-        }
+            .scheme(scheme.as_ref())
+            .authority(authority.as_ref())
+            .path_and_query(pq)
+            .build()
+            .map_err(Box::new)?;
+
+        parts.uri = uri;
+        Ok(Request::from_parts(parts, body))
     }
 }
 
@@ -133,3 +131,50 @@ where
 pub fn build_request() -> http::request::Builder {
     http::Request::builder().header(USER_AGENT_HEADER, USER_AGENT)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_set_origin() {
+        let base = "http://localhost:9001";
+        let client = Client::builder().with_endpoint(base.parse().unwrap()).build().unwrap();
+        let req = build_request()
+            .uri("/2018-06-01/runtime/invocation/next")
+            .body(())
+            .unwrap();
+        let req = client.set_origin(req).unwrap();
+        assert_eq!(
+            "http://localhost:9001/2018-06-01/runtime/invocation/next",
+            &req.uri().to_string()
+        );
+    }
+
+    #[test]
+    fn test_set_origin_with_base_path() {
+        let base = "http://localhost:9001/foo";
+        let client = Client::builder().with_endpoint(base.parse().unwrap()).build().unwrap();
+        let req = build_request()
+            .uri("/2018-06-01/runtime/invocation/next")
+            .body(())
+            .unwrap();
+        let req = client.set_origin(req).unwrap();
+        assert_eq!(
+            "http://localhost:9001/foo/2018-06-01/runtime/invocation/next",
+            &req.uri().to_string()
+        );
+
+        let base = "http://localhost:9001/foo/";
+        let client = Client::builder().with_endpoint(base.parse().unwrap()).build().unwrap();
+        let req = build_request()
+            .uri("/2018-06-01/runtime/invocation/next")
+            .body(())
+            .unwrap();
+        let req = client.set_origin(req).unwrap();
+        assert_eq!(
+            "http://localhost:9001/foo/2018-06-01/runtime/invocation/next",
+            &req.uri().to_string()
+        );
+    }
+}


Test patch: