From 3a639bfeada2aa76298f6e10dec8acf0eb2a5526 Mon Sep 17 00:00:00 2001 From: Vipin Gupta Date: Wed, 22 Jul 2026 16:22:10 +0000 Subject: [PATCH] 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). --- include/aws/lambda-runtime/runtime.h | 18 ++++++++++--- src/runtime.cpp | 39 ++++++++++++++++++++++----- tests/unit/thread_local_curl_test.cpp | 17 ++++++++++++ tests/unit/unit_tests.cpp | 9 +++++++ 4 files changed, 73 insertions(+), 10 deletions(-) diff --git a/include/aws/lambda-runtime/runtime.h b/include/aws/lambda-runtime/runtime.h index ea32c9af..bf119af5 100644 --- a/include/aws/lambda-runtime/runtime.h +++ b/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,12 +204,18 @@ class runtime { /** * Tells lambda that the function has succeeded. */ - post_outcome post_success(std::string const& request_id, invocation_response const& handler_response); + post_outcome post_success( + std::string const& request_id, + invocation_response const& handler_response, + std::string const& invocation_id = ""); /** * Tells lambda that the function has failed. */ - post_outcome post_failure(std::string const& request_id, invocation_response const& handler_response); + post_outcome post_failure( + std::string const& request_id, + invocation_response const& handler_response, + std::string const& invocation_id = ""); /** * Tells lambda that the runtime has failed during initialization. @@ -218,7 +229,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 const m_endpoints; }; diff --git a/src/runtime.cpp b/src/runtime.cpp index 1efab45a..e20c07c8 100644 --- a/src/runtime.cpp +++ b/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,32 @@ 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_init_error(runtime_response const& init_error_response) @@ -363,7 +383,8 @@ 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 +403,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 +504,13 @@ void run_handler(std::function 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 } diff --git a/tests/unit/thread_local_curl_test.cpp b/tests/unit/thread_local_curl_test.cpp index bf0f6e67..3b0fd394 100644 --- a/tests/unit/thread_local_curl_test.cpp +++ b/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()); +} diff --git a/tests/unit/unit_tests.cpp b/tests/unit/unit_tests.cpp index 7483be53..3e469fe4 100644 --- a/tests/unit/unit_tests.cpp +++ b/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()); +}