diff --git a/.github/workflows/test-on-push-and-pr.yml b/.github/workflows/test-on-push-and-pr.yml index 3171a8f..96c1d81 100644 --- a/.github/workflows/test-on-push-and-pr.yml +++ b/.github/workflows/test-on-push-and-pr.yml @@ -18,6 +18,26 @@ jobs: - name: Run 'pr' target run: make pr + test-runtime-client: + name: test runtime client + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + # The tests in this job drive the compiled extension against a stub RAPID, + # so unlike the unit tests they need the native code built first. + - name: Install aws-lambda-cpp build dependencies + run: | + sudo apt-get update + sudo apt-get install -y g++ make cmake libcurl4-openssl-dev + + - name: Install requirements + run: make init + + - name: Compile the runtime client extension and run its tests + run: make test-runtime-client + integration-test: runs-on: ubuntu-latest strategy: diff --git a/Makefile b/Makefile index 1a6e673..3803b84 100644 --- a/Makefile +++ b/Makefile @@ -8,8 +8,16 @@ init: pip3 install -r requirements/base.txt -r requirements/dev.txt .PHONY: test +# test_runtime_client_headers.py needs the compiled extension, so it runs in its +# own target (test-runtime-client) rather than here. test: check-format - pytest --cov awslambdaric --cov-report term-missing --cov-fail-under 90 tests + pytest --cov awslambdaric --cov-report term-missing --cov-fail-under 90 \ + --ignore tests/test_runtime_client_headers.py tests + +.PHONY: test-runtime-client +test-runtime-client: + python3 setup.py build_ext --inplace + pytest -v tests/test_runtime_client_headers.py .PHONY: test-integ test-integ: diff --git a/RELEASE.CHANGELOG.md b/RELEASE.CHANGELOG.md index b832ba0..9153a5b 100644 --- a/RELEASE.CHANGELOG.md +++ b/RELEASE.CHANGELOG.md @@ -1,3 +1,7 @@ +### July 15, 2026 +`4.0.2` +- Add `Lambda-Runtime-Invocation-Id` header support for cross-wiring protection. The RIC now echoes the invocation ID received from RAPID on `/next` back on `/response` and `/error`, enabling RAPID to detect and reject stale responses from timed-out invocations. + ### June 25, 2026 `4.0.1` - Support building on Alpine Linux 3.17+ (musl) without `libexecinfo-dev` ([#204](https://github.com/aws/aws-lambda-python-runtime-interface-client/pull/204)) diff --git a/awslambdaric/__init__.py b/awslambdaric/__init__.py index 45877d8..96b7844 100644 --- a/awslambdaric/__init__.py +++ b/awslambdaric/__init__.py @@ -2,4 +2,4 @@ Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved. """ -__version__ = "4.0.1" +__version__ = "4.0.2" diff --git a/awslambdaric/bootstrap.py b/awslambdaric/bootstrap.py index f54ca1a..aa3ae65 100644 --- a/awslambdaric/bootstrap.py +++ b/awslambdaric/bootstrap.py @@ -33,7 +33,6 @@ ) AWS_LAMBDA_INITIALIZATION_TYPE = "AWS_LAMBDA_INITIALIZATION_TYPE" INIT_TYPE_SNAP_START = "snap-start" -PREVIEW_RUNTIME_ENVS = {"AWS_Lambda_python3.15"} def _get_handler(handler): @@ -161,6 +160,7 @@ def handle_event_request( epoch_deadline_time_in_ms, tenant_id, log_sink, + invocation_id=None, ): error_result = None try: @@ -205,11 +205,11 @@ def handle_event_request( log_error(error_result, log_sink) lambda_runtime_client.post_invocation_error( - invoke_id, to_json(error_result), to_json(xray_fault) + invoke_id, to_json(error_result), to_json(xray_fault), invocation_id ) else: lambda_runtime_client.post_invocation_result( - invoke_id, result, result_content_type + invoke_id, result, result_content_type, invocation_id ) @@ -478,6 +478,9 @@ def _setup_logging(log_format, log_level, log_sink): logger.addHandler(logger_handler) +PREVIEW_RUNTIME_ENVS = {"AWS_Lambda_python3.15"} + + def _log_preview_runtime_warning(): """Emit a warning if the runtime version is a preview.""" if os.environ.get("LAMBDA_DISABLE_PREVIEW_WARN", ""): @@ -545,4 +548,5 @@ def run(handler, lambda_runtime_client): event_request.deadline_time_in_ms, event_request.tenant_id, log_sink, + event_request.invocation_id, ) diff --git a/awslambdaric/lambda_runtime_client.py b/awslambdaric/lambda_runtime_client.py index 2cc8f3b..0ac0d61 100644 --- a/awslambdaric/lambda_runtime_client.py +++ b/awslambdaric/lambda_runtime_client.py @@ -167,10 +167,15 @@ def wait_next_invocation(self): tenant_id=headers.get("Lambda-Runtime-Aws-Tenant-Id"), content_type=headers.get("Content-Type"), event_body=response_body, + invocation_id=headers.get("Lambda-Runtime-Invocation-Id"), ) def post_invocation_result( - self, invoke_id, result_data, content_type="application/json" + self, + invoke_id, + result_data, + content_type="application/json", + invocation_id=None, ): try: runtime_client.post_invocation_result( @@ -181,17 +186,22 @@ def post_invocation_result( else result_data.encode("utf-8") ), content_type, + invocation_id, ) except Exception as e: self.handle_exception(e) - def post_invocation_error(self, invoke_id, error_response_data, xray_fault): + def post_invocation_error( + self, invoke_id, error_response_data, xray_fault, invocation_id=None + ): try: max_header_size = 1024 * 1024 xray_fault = ( xray_fault if len(xray_fault.encode()) < max_header_size else "" ) - runtime_client.post_error(invoke_id, error_response_data, xray_fault) + runtime_client.post_error( + invoke_id, error_response_data, xray_fault, invocation_id + ) except Exception as e: self.handle_exception(e) diff --git a/awslambdaric/runtime_client.cpp b/awslambdaric/runtime_client.cpp index 7fb2e95..f29e7a6 100644 --- a/awslambdaric/runtime_client.cpp +++ b/awslambdaric/runtime_client.cpp @@ -53,9 +53,10 @@ static PyObject *method_next(PyObject *self) { auto content_type = response.content_type.c_str(); auto cognito_id = response.cognito_identity.c_str(); auto tenant_id = response.tenant_id.c_str(); + auto invocation_id = response.invocation_id.c_str(); PyObject *payload_bytes = PyBytes_FromStringAndSize(payload.c_str(), payload.length()); - PyObject *result = Py_BuildValue("(O,{s:s,s:s,s:s,s:l,s:s,s:s,s:s,s:s})", + PyObject *result = Py_BuildValue("(O,{s:s,s:s,s:s,s:l,s:s,s:s,s:s,s:s,s:s})", payload_bytes, //Py_BuildValue() increments reference counter "Lambda-Runtime-Aws-Request-Id", request_id, "Lambda-Runtime-Trace-Id", NULL_IF_EMPTY(trace_id), @@ -64,7 +65,8 @@ static PyObject *method_next(PyObject *self) { "Lambda-Runtime-Client-Context", NULL_IF_EMPTY(client_context), "Content-Type", NULL_IF_EMPTY(content_type), "Lambda-Runtime-Cognito-Identity", NULL_IF_EMPTY(cognito_id), - "Lambda-Runtime-Aws-Tenant-Id", NULL_IF_EMPTY(tenant_id) + "Lambda-Runtime-Aws-Tenant-Id", NULL_IF_EMPTY(tenant_id), + "Lambda-Runtime-Invocation-Id", NULL_IF_EMPTY(invocation_id) ); Py_XDECREF(payload_bytes); @@ -79,9 +81,9 @@ static PyObject *method_post_invocation_result(PyObject *self, PyObject *args) { PyObject *invocation_response; Py_ssize_t length; - char *request_id, *content_type, *response_as_c_string; + char *request_id, *content_type, *response_as_c_string, *invocation_id = nullptr; - if (!PyArg_ParseTuple(args, "sSs", &request_id, &invocation_response, &content_type)) { + if (!PyArg_ParseTuple(args, "sSsz", &request_id, &invocation_response, &content_type, &invocation_id)) { PyErr_SetString(PyExc_RuntimeError, "Wrong arguments"); return NULL; } @@ -91,7 +93,8 @@ static PyObject *method_post_invocation_result(PyObject *self, PyObject *args) { std::string response_string(response_as_c_string, response_as_c_string + length); auto response = aws::lambda_runtime::invocation_response::success(response_string, content_type); - auto outcome = CLIENT->post_success(request_id, response); + std::string inv_id = invocation_id ? invocation_id : ""; + auto outcome = CLIENT->post_success(request_id, response, inv_id); if (!outcome.is_success()) { PyErr_SetString(PyExc_RuntimeError, "Failed to post invocation response"); return NULL; @@ -107,15 +110,16 @@ static PyObject *method_post_error(PyObject *self, PyObject *args) { return NULL; } - char *request_id, *response_string, *xray_fault; + char *request_id, *response_string, *xray_fault, *invocation_id = nullptr; - if (!PyArg_ParseTuple(args, "sss", &request_id, &response_string, &xray_fault)) { + if (!PyArg_ParseTuple(args, "sssz", &request_id, &response_string, &xray_fault, &invocation_id)) { PyErr_SetString(PyExc_RuntimeError, "Wrong arguments"); return NULL; } auto response = aws::lambda_runtime::invocation_response(response_string, "application/json", false, xray_fault); - auto outcome = CLIENT->post_failure(request_id, response); + std::string inv_id = invocation_id ? invocation_id : ""; + auto outcome = CLIENT->post_failure(request_id, response, inv_id); if (!outcome.is_success()) { PyErr_SetString(PyExc_RuntimeError, "Failed to post invocation error"); return NULL; diff --git a/deps/aws-lambda-cpp-0.2.6.tar.gz b/deps/aws-lambda-cpp-0.2.6.tar.gz index 7c25651..c087a17 100644 Binary files a/deps/aws-lambda-cpp-0.2.6.tar.gz and b/deps/aws-lambda-cpp-0.2.6.tar.gz differ diff --git a/deps/patches/aws-lambda-cpp-add-invocation-id.patch b/deps/patches/aws-lambda-cpp-add-invocation-id.patch new file mode 100644 index 0000000..bde825f --- /dev/null +++ b/deps/patches/aws-lambda-cpp-add-invocation-id.patch @@ -0,0 +1,137 @@ +diff --git a/include/aws/lambda-runtime/runtime.h b/include/aws/lambda-runtime/runtime.h +--- a/include/aws/lambda-runtime/runtime.h ++++ b/include/aws/lambda-runtime/runtime.h +@@ -67,6 +67,11 @@ + std::string tenant_id; + + /** ++ * The unique invocation ID for cross-wiring protection. ++ */ ++ std::string invocation_id; ++ ++ /** + * Function execution deadline counted in milliseconds since the Unix epoch. + */ + std::chrono::time_point deadline; +@@ -184,11 +189,15 @@ + /** + * 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); + + /** +@@ -203,7 +212,8 @@ + 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); + + private: + std::string const m_user_agent_header; +diff --git a/src/runtime.cpp b/src/runtime.cpp +--- a/src/runtime.cpp ++++ b/src/runtime.cpp +@@ -41,6 +41,7 @@ + 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"; + + enum Endpoints { + INIT, +@@ -294,6 +295,10 @@ + req.tenant_id = resp.get_header(TENANT_ID_HEADER); + } + ++ if (resp.has_header(INVOCATION_ID_HEADER)) { ++ req.invocation_id = resp.get_header(INVOCATION_ID_HEADER); ++ } ++ + if (resp.has_header(DEADLINE_MS_HEADER)) { + auto const& deadline_string = resp.get_header(DEADLINE_MS_HEADER); + constexpr int base = 10; +@@ -310,29 +315,40 @@ + return next_outcome(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()); ++ return do_post(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_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, 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()); ++ return do_post(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) ++{ ++ return post_failure(request_id, handler_response, ""); + } + + runtime::post_outcome runtime::post_init_error(runtime_response const& init_error_response) + { + std::string const url = m_endpoints[Endpoints::INIT]; +- return do_post(url, init_error_response.get_content_type(), init_error_response.get_payload(), init_error_response.get_xray_response()); ++ return do_post(url, init_error_response.get_content_type(), init_error_response.get_payload(), 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(m_curl_handle, CURLOPT_URL, url.c_str()); +@@ -351,6 +366,10 @@ + 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()); +@@ -448,13 +467,13 @@ + 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/deps/patches/aws-lambda-cpp-musl-no-execinfo.patch b/deps/patches/aws-lambda-cpp-musl-no-execinfo.patch new file mode 100644 index 0000000..b5cec23 --- /dev/null +++ b/deps/patches/aws-lambda-cpp-musl-no-execinfo.patch @@ -0,0 +1,21 @@ +--- a/CMakeLists.txt ++++ b/CMakeLists.txt +@@ -11,10 +11,17 @@ + add_library(${PROJECT_NAME} + "src/logging.cpp" + "src/runtime.cpp" +- "src/backward.cpp" + "${CMAKE_CURRENT_BINARY_DIR}/version.cpp" + ) + ++include(CheckIncludeFileCXX) ++check_include_file_cxx("execinfo.h" HAVE_EXECINFO_H) ++if (HAVE_EXECINFO_H) ++ target_sources(${PROJECT_NAME} PRIVATE "src/backward.cpp") ++else() ++ message("-- execinfo.h not found, stack traces disabled (musl/Alpine Linux)") ++endif() ++ + set_target_properties(${PROJECT_NAME} PROPERTIES + SOVERSION 0 + VERSION ${PROJECT_VERSION}) diff --git a/scripts/update_deps.sh b/scripts/update_deps.sh index 841d320..f1e71d3 100755 --- a/scripts/update_deps.sh +++ b/scripts/update_deps.sh @@ -29,10 +29,12 @@ wget -c https://github.com/awslabs/aws-lambda-cpp/archive/v$AWS_LAMBDA_CPP_RELEA patch -p1 < ../patches/aws-lambda-cpp-add-xray-response.patch && \ patch -p1 < ../patches/aws-lambda-cpp-posting-init-errors.patch && \ patch -p1 < ../patches/aws-lambda-cpp-make-the-runtime-client-user-agent-overrideable.patch && \ + patch -p1 < ../patches/aws-lambda-cpp-musl-no-execinfo.patch && \ patch -p1 < ../patches/aws-lambda-cpp-make-lto-optional.patch && \ patch -p1 < ../patches/aws-lambda-cpp-add-content-type.patch && \ patch -p1 < ../patches/aws-lambda-cpp-add-tenant-id.patch && \ - patch -p1 < ../patches/aws-lambda-cpp-logging-error.patch + patch -p1 < ../patches/aws-lambda-cpp-logging-error.patch && \ + patch -p1 < ../patches/aws-lambda-cpp-add-invocation-id.patch ) ## Pack again and remove the folder diff --git a/tests/invocation_id_probe.py b/tests/invocation_id_probe.py new file mode 100644 index 0000000..8ddf5fb --- /dev/null +++ b/tests/invocation_id_probe.py @@ -0,0 +1,51 @@ +""" +Copyright 2026 Amazon.com, Inc. or its affiliates. All Rights Reserved. + +Client half of test_runtime_client_headers.py, run as a subprocess. + +It is a separate process rather than a thread because the native +post_invocation_result/post_error entry points hold the GIL for the duration of +the blocking curl call so an in-process stub server thread could never be scheduled to answer the POST. +""" + +import json +import sys + +import runtime_client + +INVOCATION_ID_HEADER = "Lambda-Runtime-Invocation-Id" +REQUEST_ID_HEADER = "Lambda-Runtime-Aws-Request-Id" + + +def main(mode): + runtime_client.initialize_client("test-agent") + + _, headers = runtime_client.next() + invocation_id = headers[INVOCATION_ID_HEADER] + + # Thread the invocation id straight back, exactly as bootstrap does. Passing + # it through untouched is what makes the POST assertions meaningful. + if mode == "response": + runtime_client.post_invocation_result( + headers[REQUEST_ID_HEADER], + b'{"ok": true}', + "application/json", + invocation_id, + ) + elif mode == "error": + runtime_client.post_error( + headers[REQUEST_ID_HEADER], + '{"errorMessage": "boom"}', + "{}", + invocation_id, + ) + else: + raise SystemExit(f"unknown mode: {mode!r}") + + print(json.dumps({"next_invocation_id": invocation_id})) + + +if __name__ == "__main__": + if len(sys.argv) != 2: + raise SystemExit(__doc__) + main(sys.argv[1]) diff --git a/tests/test_bootstrap.py b/tests/test_bootstrap.py index 1eb2bb0..f4d78b6 100644 --- a/tests/test_bootstrap.py +++ b/tests/test_bootstrap.py @@ -95,6 +95,7 @@ def test_handle_event_request_happy_case(self): "invoke_id", '{"input": "event_body", "aws_request_id": "invoke_id"}', "application/json", + None, ) def test_handle_event_request_invalid_client_context(self): @@ -862,7 +863,7 @@ def test_application_json(self): ) self.lambda_runtime.post_invocation_result.assert_called_once_with( - "invoke-id", '{"response": "foo"}', "application/json" + "invoke-id", '{"response": "foo"}', "application/json", None ) def test_binary_request_binary_response(self): @@ -882,7 +883,7 @@ def test_binary_request_binary_response(self): ) self.lambda_runtime.post_invocation_result.assert_called_once_with( - "invoke-id", event_body, "application/unknown" + "invoke-id", event_body, "application/unknown", None ) def test_json_request_binary_response(self): @@ -902,7 +903,7 @@ def test_json_request_binary_response(self): ) self.lambda_runtime.post_invocation_result.assert_called_once_with( - "invoke-id", binary_data, "application/unknown" + "invoke-id", binary_data, "application/unknown", None ) def test_binary_with_application_json(self): @@ -927,12 +928,56 @@ def test_binary_with_application_json(self): invoke_id, error_result, xray_fault, + invocation_id, ), _ = self.lambda_runtime.post_invocation_error.call_args error_dict = json.loads(error_result) self.assertEqual("invoke-id", invoke_id) self.assertEqual("Runtime.UnmarshalError", error_dict["errorType"]) + def test_invocation_id_wired_to_post_invocation_result(self): + bootstrap.handle_event_request( + lambda_runtime_client=self.lambda_runtime, + request_handler=lambda event, ctx: {"ok": True}, + invoke_id="invoke-id", + event_body=b'{"key":"value"}', + content_type="application/json", + client_context_json=None, + cognito_identity_json=None, + invoked_function_arn="invocation-arn", + epoch_deadline_time_in_ms=1415836801003, + tenant_id=None, + log_sink=bootstrap.StandardLogSink(), + invocation_id="inv-uuid-wiring-test", + ) + + self.lambda_runtime.post_invocation_result.assert_called_once() + _, kwargs = self.lambda_runtime.post_invocation_result.call_args + args = self.lambda_runtime.post_invocation_result.call_args[0] + self.assertEqual("inv-uuid-wiring-test", args[3]) + + def test_invocation_id_wired_to_post_invocation_error(self): + bootstrap.handle_event_request( + lambda_runtime_client=self.lambda_runtime, + request_handler=lambda event, ctx: (_ for _ in ()).throw( + ValueError("boom") + ), + invoke_id="invoke-id", + event_body=b'{"key":"value"}', + content_type="application/json", + client_context_json=None, + cognito_identity_json=None, + invoked_function_arn="invocation-arn", + epoch_deadline_time_in_ms=1415836801003, + tenant_id=None, + log_sink=bootstrap.StandardLogSink(), + invocation_id="inv-uuid-error-test", + ) + + self.lambda_runtime.post_invocation_error.assert_called_once() + args = self.lambda_runtime.post_invocation_error.call_args[0] + self.assertEqual("inv-uuid-error-test", args[3]) + class TestLogError(unittest.TestCase): @patch("sys.stdout", new_callable=StringIO) diff --git a/tests/test_lambda_runtime_client.py b/tests/test_lambda_runtime_client.py index c25e581..85cbdb7 100644 --- a/tests/test_lambda_runtime_client.py +++ b/tests/test_lambda_runtime_client.py @@ -200,6 +200,101 @@ def test_wait_next_invocation_with_empty_tenant_id_header( self.assertEqual(event_request.tenant_id, "") self.assertEqual(event_request.event_body, response_body) + @patch("awslambdaric.lambda_runtime_client.runtime_client") + def test_wait_next_invocation_with_invocation_id(self, mock_runtime_client): + response_body = b"{}" + headers = { + **self.get_next_headers, + "Lambda-Runtime-Invocation-Id": "inv-uuid-1234", + } + mock_runtime_client.next.return_value = response_body, headers + runtime_client = LambdaRuntimeClient("localhost:1234") + + event_request = runtime_client.wait_next_invocation() + + self.assertIsNotNone(event_request) + self.assertEqual(event_request.invocation_id, "inv-uuid-1234") + + @patch("awslambdaric.lambda_runtime_client.runtime_client") + def test_wait_next_invocation_without_invocation_id(self, mock_runtime_client): + response_body = b"{}" + headers = { + "Lambda-Runtime-Aws-Request-Id": "RID1234", + "Lambda-Runtime-Trace-Id": "TID1234", + "Lambda-Runtime-Invoked-Function-Arn": "FARN1234", + "Lambda-Runtime-Deadline-Ms": 12, + "Lambda-Runtime-Client-Context": "client_context", + "Lambda-Runtime-Cognito-Identity": "cognito_identity", + "Content-Type": "application/json", + } + mock_runtime_client.next.return_value = response_body, headers + runtime_client = LambdaRuntimeClient("localhost:1234") + + event_request = runtime_client.wait_next_invocation() + + self.assertIsNotNone(event_request) + self.assertIsNone(event_request.invocation_id) + + @patch("awslambdaric.lambda_runtime_client.runtime_client") + def test_post_invocation_result_with_invocation_id(self, mock_runtime_client): + runtime_client = LambdaRuntimeClient("localhost:1234") + response_data = "data" + invoke_id = "1234" + invocation_id = "inv-uuid-5678" + + runtime_client.post_invocation_result( + invoke_id, response_data, invocation_id=invocation_id + ) + + mock_runtime_client.post_invocation_result.assert_called_once_with( + invoke_id, response_data.encode("utf-8"), "application/json", invocation_id + ) + + @patch("awslambdaric.lambda_runtime_client.runtime_client") + def test_post_invocation_error_with_invocation_id(self, mock_runtime_client): + runtime_client = LambdaRuntimeClient("localhost:1234") + error_data = "data" + invoke_id = "1234" + xray_fault = "xray_fault" + invocation_id = "inv-uuid-5678" + + runtime_client.post_invocation_error( + invoke_id, error_data, xray_fault, invocation_id + ) + + mock_runtime_client.post_error.assert_called_once_with( + invoke_id, error_data, xray_fault, invocation_id + ) + + @patch("awslambdaric.lambda_runtime_client.runtime_client") + def test_post_invocation_result_without_invocation_id_passes_none( + self, mock_runtime_client + ): + runtime_client = LambdaRuntimeClient("localhost:1234") + response_data = "data" + invoke_id = "1234" + + runtime_client.post_invocation_result(invoke_id, response_data) + + mock_runtime_client.post_invocation_result.assert_called_once_with( + invoke_id, response_data.encode("utf-8"), "application/json", None + ) + + @patch("awslambdaric.lambda_runtime_client.runtime_client") + def test_post_invocation_error_without_invocation_id_passes_none( + self, mock_runtime_client + ): + runtime_client = LambdaRuntimeClient("localhost:1234") + error_data = "data" + invoke_id = "1234" + xray_fault = "xray_fault" + + runtime_client.post_invocation_error(invoke_id, error_data, xray_fault) + + mock_runtime_client.post_error.assert_called_once_with( + invoke_id, error_data, xray_fault, None + ) + error_result = { "errorMessage": "Dummy message", "errorType": "Runtime.DummyError", @@ -270,7 +365,7 @@ def test_post_invocation_result(self, mock_runtime_client): runtime_client.post_invocation_result(invoke_id, response_data) mock_runtime_client.post_invocation_result.assert_called_once_with( - invoke_id, response_data.encode("utf-8"), "application/json" + invoke_id, response_data.encode("utf-8"), "application/json", None ) @patch("awslambdaric.lambda_runtime_client.runtime_client") @@ -283,7 +378,7 @@ def test_post_invocation_result_binary_data(self, mock_runtime_client): runtime_client.post_invocation_result(invoke_id, response_data, content_type) mock_runtime_client.post_invocation_result.assert_called_once_with( - invoke_id, response_data, content_type + invoke_id, response_data, content_type, None ) @patch("awslambdaric.lambda_runtime_client.runtime_client") @@ -309,7 +404,7 @@ def test_post_invocation_error(self, mock_runtime_client): runtime_client.post_invocation_error(invoke_id, error_data, xray_fault) mock_runtime_client.post_error.assert_called_once_with( - invoke_id, error_data, xray_fault + invoke_id, error_data, xray_fault, None ) @patch("awslambdaric.lambda_runtime_client.runtime_client") @@ -461,7 +556,7 @@ def test_post_invocation_error_with_large_xray_cause(self, mock_runtime_client): runtime_client.post_invocation_error(invoke_id, error_data, large_xray_fault) mock_runtime_client.post_error.assert_called_once_with( - invoke_id, error_data, large_xray_fault + invoke_id, error_data, large_xray_fault, None ) @patch("awslambdaric.lambda_runtime_client.runtime_client") @@ -476,7 +571,7 @@ def test_post_invocation_error_with_too_large_xray_cause(self, mock_runtime_clie ) mock_runtime_client.post_error.assert_called_once_with( - invoke_id, error_data, "" + invoke_id, error_data, "", None ) @patch("http.client.HTTPConnection", autospec=http.client.HTTPConnection) diff --git a/tests/test_runtime_client_headers.py b/tests/test_runtime_client_headers.py new file mode 100644 index 0000000..d853f46 --- /dev/null +++ b/tests/test_runtime_client_headers.py @@ -0,0 +1,141 @@ +""" +Copyright 2026 Amazon.com, Inc. or its affiliates. All Rights Reserved. +""" + +import http.server +import json +import os +import subprocess +import sys +import threading +import unittest + +INVOCATION_ID_HEADER = "Lambda-Runtime-Invocation-Id" +REQUEST_ID = "request-id-1" + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +CLIENT_PROBE = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "invocation_id_probe.py" +) + + +class StubRapidHandler(http.server.BaseHTTPRequestHandler): + """Minimal stand-in for the RAPID /next, /response and /error endpoints. + + Each request's raw header block is appended to the server's `received` list + so that tests can assert on what actually went over the wire rather than on + what the Python layer believed it was sending. + """ + + protocol_version = "HTTP/1.1" + + def log_message(self, format, *args): + pass # keep the test output clean + + def _record(self): + self.server.received.append((self.path, self.headers)) + + def do_GET(self): + self._record() + body = b"{}" + self.send_response(200) + self.send_header("Lambda-Runtime-Aws-Request-Id", REQUEST_ID) + self.send_header("Lambda-Runtime-Invoked-Function-Arn", "arn:aws:lambda:::f") + self.send_header("Lambda-Runtime-Deadline-Ms", "1735689600000") + self.send_header("Content-Type", "application/json") + if self.server.invocation_id is not None: + self.send_header(INVOCATION_ID_HEADER, self.server.invocation_id) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_POST(self): + self._record() + length = int(self.headers.get("Content-Length") or 0) + if length: + self.rfile.read(length) + self.send_response(202) + self.send_header("Content-Length", "0") + self.end_headers() + + +class StubRapid(http.server.ThreadingHTTPServer): + + def __init__(self): + super().__init__(("127.0.0.1", 0), StubRapidHandler) + self.invocation_id = None + self.received = [] + + @property + def endpoint(self): + return "{}:{}".format(*self.server_address) + + +class TestInvocationIdOnTheWire(unittest.TestCase): + """End-to-end coverage of the Lambda-Runtime-Invocation-Id echo. + + These tests drive the compiled extension against a real socket, so they + cover the header-emitting C++ code in aws-lambda-cpp that the mocked unit + tests in test_lambda_runtime_client.py cannot reach. If + aws-lambda-cpp-add-invocation-id.patch is ever dropped from the vendored + tarball, these are the tests that fail. + """ + + def _run_client(self, invocation_id, mode="response"): + """Serve one invocation to a real client and return its POST headers.""" + server = StubRapid() + server.invocation_id = invocation_id + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + self.addCleanup(thread.join, 5) + self.addCleanup(server.shutdown) + self.addCleanup(server.server_close) + + env = dict( + os.environ, + AWS_LAMBDA_RUNTIME_API=server.endpoint, + PYTHONPATH=os.pathsep.join( + [REPO_ROOT] + + ([os.environ["PYTHONPATH"]] if os.environ.get("PYTHONPATH") else []) + ), + ) + completed = subprocess.run( + [sys.executable, CLIENT_PROBE, mode], + env=env, + cwd=REPO_ROOT, + capture_output=True, + text=True, + timeout=30, + ) + self.assertEqual( + 0, + completed.returncode, + f"client failed:\nstdout={completed.stdout}\nstderr={completed.stderr}", + ) + + posted = [h for path, h in server.received if path.endswith("/" + mode)] + self.assertEqual(1, len(posted), f"expected exactly one POST to /{mode}") + return json.loads(completed.stdout), posted[0] + + def test_invocation_id_is_echoed_on_response(self): + seen, posted = self._run_client("inv-uuid-round-trip") + + self.assertEqual("inv-uuid-round-trip", seen["next_invocation_id"]) + # get_all, not get: a duplicated header must be visible, not normalized. + self.assertEqual(["inv-uuid-round-trip"], posted.get_all(INVOCATION_ID_HEADER)) + + def test_invocation_id_is_echoed_on_error(self): + seen, posted = self._run_client("inv-uuid-error-path", mode="error") + + self.assertEqual("inv-uuid-error-path", seen["next_invocation_id"]) + self.assertEqual(["inv-uuid-error-path"], posted.get_all(INVOCATION_ID_HEADER)) + + def test_no_header_is_sent_when_rapid_did_not_send_one(self): + seen, posted = self._run_client(None) + + self.assertIsNone(seen["next_invocation_id"]) + self.assertIsNone(posted.get_all(INVOCATION_ID_HEADER)) + + +if __name__ == "__main__": + unittest.main()