feat: Add Lambda-Runtime-Invocation-Id support for cross-wiring protection

Parse the invocation ID from /next response headers, store in
invocation_request, and echo it back on /response and /error via
do_post. Header only sent when non-empty (backward compatible).

Author: vip-amzn

include/aws/lambda-runtime/runtime.h

@@ -71,6 +71,11 @@ struct invocation_request {
      */
     std::string tenant_id;
 
+    /**
+     * The unique invocation ID for cross-wiring protection.
+     */
+    std::string invocation_id;
+
     /**
      * The number of milliseconds left before lambda terminates the current execution.
      */
@@ -199,11 +204,21 @@ class runtime {
     /**
      * Tells lambda that the function has succeeded.
      */
+    post_outcome post_success(
+        std::string const& request_id,
+        invocation_response const& handler_response,
+        std::string const& invocation_id);
+
     post_outcome post_success(std::string const& request_id, invocation_response const& handler_response);
 
     /**
      * Tells lambda that the function has failed.
      */
+    post_outcome post_failure(
+        std::string const& request_id,
+        invocation_response const& handler_response,
+        std::string const& invocation_id);
+
     post_outcome post_failure(std::string const& request_id, invocation_response const& handler_response);
 
     /**
@@ -218,7 +233,8 @@ class runtime {
         std::string const& url,
         std::string const& content_type,
         std::string const& payload,
-        std::string const& xray_response);
+        std::string const& xray_response,
+        std::string const& invocation_id);
     std::string const m_user_agent_header;
     std::array<std::string const, 3> const m_endpoints;
 };

src/runtime.cpp

@@ -42,6 +42,7 @@ static constexpr auto COGNITO_IDENTITY_HEADER = "lambda-runtime-cognito-identity
 static constexpr auto DEADLINE_MS_HEADER = "lambda-runtime-deadline-ms";
 static constexpr auto FUNCTION_ARN_HEADER = "lambda-runtime-invoked-function-arn";
 static constexpr auto TENANT_ID_HEADER = "lambda-runtime-aws-tenant-id";
+static constexpr auto INVOCATION_ID_HEADER = "lambda-runtime-invocation-id";
 struct curl_handle_wrapper {
     CURL* handle;
     curl_handle_wrapper() : handle(curl_easy_init()) {}
@@ -318,6 +319,11 @@ runtime::next_outcome runtime::get_next()
         req.tenant_id = std::move(out).get_result();
     }
 
+    out = resp.get_header(INVOCATION_ID_HEADER);
+    if (out.is_success()) {
+        req.invocation_id = std::move(out).get_result();
+    }
+
     out = resp.get_header(DEADLINE_MS_HEADER);
     if (out.is_success()) {
         auto const& deadline_string = std::move(out).get_result();
@@ -335,18 +341,42 @@ runtime::next_outcome runtime::get_next()
     return {req};
 }
 
-runtime::post_outcome runtime::post_success(std::string const& request_id, invocation_response const& handler_response)
+runtime::post_outcome runtime::post_success(
+    std::string const& request_id,
+    invocation_response const& handler_response,
+    std::string const& invocation_id)
 {
     std::string const url = m_endpoints[Endpoints::RESULT] + request_id + "/response";
     return do_post(
-        url, handler_response.get_content_type(), handler_response.get_payload(), handler_response.get_xray_response());
+        url,
+        handler_response.get_content_type(),
+        handler_response.get_payload(),
+        handler_response.get_xray_response(),
+        invocation_id);
 }
 
-runtime::post_outcome runtime::post_failure(std::string const& request_id, invocation_response const& handler_response)
+runtime::post_outcome runtime::post_failure(
+    std::string const& request_id,
+    invocation_response const& handler_response,
+    std::string const& invocation_id)
 {
     std::string const url = m_endpoints[Endpoints::RESULT] + request_id + "/error";
     return do_post(
-        url, handler_response.get_content_type(), handler_response.get_payload(), handler_response.get_xray_response());
+        url,
+        handler_response.get_content_type(),
+        handler_response.get_payload(),
+        handler_response.get_xray_response(),
+        invocation_id);
+}
+
+runtime::post_outcome runtime::post_success(std::string const& request_id, invocation_response const& handler_response)
+{
+    return post_success(request_id, handler_response, "");
+}
+
+runtime::post_outcome runtime::post_failure(std::string const& request_id, invocation_response const& handler_response)
+{
+    return post_failure(request_id, handler_response, "");
 }
 
 runtime::post_outcome runtime::post_init_error(runtime_response const& init_error_response)
@@ -356,14 +386,16 @@ runtime::post_outcome runtime::post_init_error(runtime_response const& init_erro
         url,
         init_error_response.get_content_type(),
         init_error_response.get_payload(),
-        init_error_response.get_xray_response());
+        init_error_response.get_xray_response(),
+        "");
 }
 
 runtime::post_outcome runtime::do_post(
     std::string const& url,
     std::string const& content_type,
     std::string const& payload,
-    std::string const& xray_response)
+    std::string const& xray_response,
+    std::string const& invocation_id)
 {
     set_curl_post_result_options();
     curl_easy_setopt(lambda_runtime::m_curl_handle, CURLOPT_URL, url.c_str());
@@ -382,6 +414,10 @@ runtime::post_outcome runtime::do_post(
     headers = curl_slist_append(headers, "transfer-encoding:");
     headers = curl_slist_append(headers, m_user_agent_header.c_str());
 
+    if (!invocation_id.empty()) {
+        headers = curl_slist_append(headers, (std::string(INVOCATION_ID_HEADER) + ": " + invocation_id).c_str());
+    }
+
     logging::log_debug(
         LOG_TAG, "calculating content length... %s", ("content-length: " + std::to_string(payload.length())).c_str());
     headers = curl_slist_append(headers, ("content-length: " + std::to_string(payload.length())).c_str());
@@ -479,13 +515,13 @@ void run_handler(std::function<invocation_response(invocation_request const&)> c
         logging::log_info(LOG_TAG, "Invoking user handler completed.");
 
         if (res.is_success()) {
-            const auto post_outcome = rt.post_success(req.request_id, res);
+            const auto post_outcome = rt.post_success(req.request_id, res, req.invocation_id);
             if (!handle_post_outcome(post_outcome, req.request_id)) {
                 return; // TODO: implement a better retry strategy
             }
         }
         else {
-            const auto post_outcome = rt.post_failure(req.request_id, res);
+            const auto post_outcome = rt.post_failure(req.request_id, res, req.invocation_id);
             if (!handle_post_outcome(post_outcome, req.request_id)) {
                 return; // TODO: implement a better retry strategy
             }

tests/unit/thread_local_curl_test.cpp

@@ -60,3 +60,20 @@ TEST(ThreadLocalCurl, sequential_requests_on_same_runtime)
         ASSERT_FALSE(outcome.is_success());
     }
 }
+
+TEST(ThreadLocalCurl, post_success_with_invocation_id_does_not_crash)
+{
+    runtime rt("http://127.0.0.1:9001");
+    auto response = invocation_response::success("test payload", "application/json");
+    auto outcome = rt.post_success("test-request-id", response, "test-invocation-uuid-1234");
+    // Connection will fail but the invocation_id header path should not crash
+    ASSERT_FALSE(outcome.is_success());
+}
+
+TEST(ThreadLocalCurl, post_failure_with_invocation_id_does_not_crash)
+{
+    runtime rt("http://127.0.0.1:9001");
+    auto response = invocation_response::failure("error msg", "TestError", "");
+    auto outcome = rt.post_failure("test-request-id", response, "test-invocation-uuid-5678");
+    ASSERT_FALSE(outcome.is_success());
+}

tests/unit/unit_tests.cpp

@@ -333,6 +333,7 @@ TEST(InvocationRequestTest, default_fields_are_empty)
     EXPECT_TRUE(req.cognito_identity.empty());
     EXPECT_TRUE(req.function_arn.empty());
     EXPECT_TRUE(req.tenant_id.empty());
+    EXPECT_TRUE(req.invocation_id.empty());
 }
 
 // --- version tests (no AWS SDK needed) ---
@@ -353,3 +354,11 @@ TEST(VersionTest, version_format)
     }
     EXPECT_EQ(2, dots);
 }
+
+// --- invocation_id cross-wiring protection tests ---
+
+TEST(InvocationRequestTest, invocation_id_default_empty)
+{
+    invocation_request req;
+    EXPECT_TRUE(req.invocation_id.empty());
+}