diff --git a/.github/workflows/dockerized-test.yml b/.github/workflows/dockerized-test.yml index f5c5baa..3fb9b0b 100644 --- a/.github/workflows/dockerized-test.yml +++ b/.github/workflows/dockerized-test.yml @@ -34,7 +34,7 @@ jobs: run: docker build . -t local/test -f Dockerfile.test --build-arg BASE_IMAGE=public.ecr.aws/lambda/ruby:${{ matrix.ruby_version }} - name: Run tests - uses: aws/containerized-test-runner-for-aws-lambda@511d270614f2c6b1613848db6dcf920a591c3c89 # main + uses: aws/containerized-test-runner-for-aws-lambda@0863dd17b5fc19585250a2405c0f939a77b4f397 # main with: suiteFileArray: '["./test/dockerized/suites/*.json"]' dockerImageName: 'local/test' @@ -44,7 +44,7 @@ jobs: run: docker build . -t local/test-proxy -f Dockerfile.test.proxy - name: Run proxy tests - uses: aws/containerized-test-runner-for-aws-lambda@511d270614f2c6b1613848db6dcf920a591c3c89 # main + uses: aws/containerized-test-runner-for-aws-lambda@0863dd17b5fc19585250a2405c0f939a77b4f397 # main with: suiteFileArray: '["./test/dockerized/suites/proxy/*.json"]' dockerImageName: 'local/test-proxy' diff --git a/lib/aws_lambda_ric/lambda_context.rb b/lib/aws_lambda_ric/lambda_context.rb index 259ba16..e0ca447 100644 --- a/lib/aws_lambda_ric/lambda_context.rb +++ b/lib/aws_lambda_ric/lambda_context.rb @@ -1,6 +1,11 @@ # frozen_string_literal: true class LambdaContext + # Allowlist of W3C trace-context fields that may be surfaced through + # LambdaContext#w3c. Any other key carried on clientContext.w3c is ignored, + # and any allowlisted key whose value is not a string is dropped. + W3C_ALLOWED_FIELDS = %w[traceparent tracestate baggage].freeze + attr_reader :aws_request_id, :invoked_function_arn, :log_group_name, :log_stream_name, :function_name, :memory_limit_in_mb, :function_version, :identity, :tenant_id, :client_context, :deadline_ms @@ -17,7 +22,12 @@ def initialize(request) @function_version = ENV['AWS_LAMBDA_FUNCTION_VERSION'] @identity = JSON.parse(request['Lambda-Runtime-Cognito-Identity']) unless request['Lambda-Runtime-Cognito-Identity'].to_s.empty? @tenant_id = request['Lambda-Runtime-Aws-Tenant-Id'] unless request['Lambda-Runtime-Aws-Tenant-Id'].to_s.empty? - @client_context = JSON.parse(request['Lambda-Runtime-Client-Context']) unless request['Lambda-Runtime-Client-Context'].to_s.empty? + @w3c_fields = {} + unless request['Lambda-Runtime-Client-Context'].to_s.empty? + client_context = JSON.parse(request['Lambda-Runtime-Client-Context']) + @w3c_fields = self.class.extract_and_strip_w3c(client_context) + @client_context = client_context + end end def get_remaining_time_in_millis @@ -25,4 +35,27 @@ def get_remaining_time_in_millis remaining = @deadline_ms - now remaining.positive? ? remaining : 0 end + + # Return the W3C trace-context captured at invoke time, as a fresh copy so + # callers cannot mutate the context's internal state. + def w3c + @w3c_fields.dup + end + + # Pop "w3c" out of the parsed client_context hash and return a normalized + # copy of its allowlisted string fields (see W3C_ALLOWED_FIELDS). + def self.extract_and_strip_w3c(client_context) + return {} unless client_context.is_a?(Hash) + return {} unless client_context.key?('w3c') + + raw_w3c = client_context.delete('w3c') + return {} unless raw_w3c.is_a?(Hash) + + fields = {} + W3C_ALLOWED_FIELDS.each do |key| + value = raw_w3c[key] + fields[key] = value if value.is_a?(String) + end + fields + end end diff --git a/test/dockerized/suites/w3c.json b/test/dockerized/suites/w3c.json new file mode 100644 index 0000000..45f96b8 --- /dev/null +++ b/test/dockerized/suites/w3c.json @@ -0,0 +1,165 @@ +{ + "tests": [ + { + "name": "w3c_is_callable_on_context", + "handler": "w3c.w3c_is_callable", + "request": {}, + "assertions": [ + { "response": { "isCallable": true } } + ] + }, + { + "name": "w3c_returns_empty_when_no_client_context_header", + "handler": "w3c.get_w3c", + "request": {}, + "assertions": [ + { "response": {} } + ] + }, + { + "name": "w3c_returns_empty_when_client_context_has_no_w3c_key", + "handler": "w3c.get_w3c", + "request": {}, + "clientContext": { + "custom": { "value": "test" } + }, + "assertions": [ + { "response": {} } + ] + }, + { + "name": "w3c_returns_baggage_only", + "handler": "w3c.get_w3c", + "request": {}, + "clientContext": { + "w3c": { "baggage": "userId=alice" } + }, + "assertions": [ + { "response": { "baggage": "userId=alice" } } + ] + }, + { + "name": "w3c_returns_all_three_allowlisted_fields", + "handler": "w3c.get_w3c", + "request": {}, + "clientContext": { + "w3c": { + "traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01", + "tracestate": "rojo=00f067aa0ba902b7", + "baggage": "userId=alice" + } + }, + "assertions": [ + { + "response": { + "traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01", + "tracestate": "rojo=00f067aa0ba902b7", + "baggage": "userId=alice" + } + } + ] + }, + { + "name": "w3c_allowlist_drops_non_allowlisted_keys", + "handler": "w3c.get_w3c", + "request": {}, + "clientContext": { + "w3c": { + "baggage": "keep=me", + "unknownField": "should-not-appear", + "x-custom-trace": "should-not-appear" + } + }, + "assertions": [ + { "response": { "baggage": "keep=me" } } + ] + }, + { + "name": "w3c_drops_allowlisted_fields_with_non_string_values", + "handler": "w3c.get_w3c", + "request": {}, + "clientContext": { + "w3c": { + "traceparent": 42, + "tracestate": null, + "baggage": { "nested": "no" } + } + }, + "assertions": [ + { "response": {} } + ] + }, + { + "name": "w3c_treats_non_object_as_empty", + "handler": "w3c.get_w3c", + "request": {}, + "clientContext": { + "w3c": "not-an-object" + }, + "assertions": [ + { "response": {} } + ] + }, + { + "name": "w3c_treats_array_as_empty", + "handler": "w3c.get_w3c", + "request": {}, + "clientContext": { + "w3c": ["baggage=abc"] + }, + "assertions": [ + { "response": {} } + ] + }, + { + "name": "w3c_strips_source_client_context_w3c_after_construction", + "handler": "w3c.get_w3c_and_source", + "request": {}, + "clientContext": { + "custom": { "value": "test" }, + "w3c": { + "traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01", + "baggage": "userId=alice" + } + }, + "assertions": [ + { + "response": { + "w3c": { + "traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01", + "baggage": "userId=alice" + }, + "clientContextIsDefined": true, + "clientContextHasW3c": false, + "clientContext": { "custom": { "value": "test" } } + } + } + ] + }, + { + "name": "client_context_is_echoed_when_no_w3c_key", + "handler": "w3c.echo_client_context", + "request": {}, + "clientContext": { + "custom": { "value": "hello" }, + "env": { "stage": "beta" } + }, + "assertions": [ + { + "response": { + "custom": { "value": "hello" }, + "env": { "stage": "beta" } + } + } + ] + }, + { + "name": "client_context_is_null_when_header_absent", + "handler": "w3c.echo_client_context", + "request": {}, + "assertions": [ + { "response": null } + ] + } + ] +} diff --git a/test/dockerized/tasks/w3c.rb b/test/dockerized/tasks/w3c.rb new file mode 100644 index 0000000..c952e3f --- /dev/null +++ b/test/dockerized/tasks/w3c.rb @@ -0,0 +1,24 @@ +# Copyright 2026 Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +def get_w3c(event:, context:) + context.w3c +end + +def get_w3c_and_source(event:, context:) + client_context = context.client_context + { + w3c: context.w3c, + clientContextIsDefined: !client_context.nil?, + clientContextHasW3c: !client_context.nil? && client_context.key?('w3c'), + clientContext: client_context + } +end + +def echo_client_context(event:, context:) + context.client_context +end + +def w3c_is_callable(event:, context:) + { isCallable: context.respond_to?(:w3c) } +end diff --git a/test/unit/lambda_context_w3c_test.rb b/test/unit/lambda_context_w3c_test.rb new file mode 100644 index 0000000..1ebfd3b --- /dev/null +++ b/test/unit/lambda_context_w3c_test.rb @@ -0,0 +1,90 @@ +# frozen_string_literal: true + +require 'json' +require_relative '../../lib/aws_lambda_ric/lambda_context' +require 'minitest/autorun' +require 'test/unit/assertions' + +include Test::Unit::Assertions + +class LambdaContextW3CTest < Minitest::Test + def build_context(client_context_hash) + request = { + 'Lambda-Runtime-Aws-Request-Id' => 'invoke-id-w3c', + 'Lambda-Runtime-Deadline-Ms' => '0', + 'Lambda-Runtime-Invoked-Function-Arn' => 'arn:test:w3c' + } + unless client_context_hash.nil? + request['Lambda-Runtime-Client-Context'] = JSON.generate(client_context_hash) + end + LambdaContext.new(request) + end + + def test_w3c_returns_empty_when_no_client_context + assert_equal({}, build_context(nil).w3c) + end + + def test_w3c_returns_empty_and_leaves_client_context_untouched_when_no_w3c_key + context = build_context({ 'custom' => { 'value' => 'test' } }) + assert_equal({}, context.w3c) + assert_equal({ 'custom' => { 'value' => 'test' } }, context.client_context) + end + + def test_w3c_returns_baggage_only + context = build_context({ 'w3c' => { 'baggage' => 'userId=alice' } }) + assert_equal({ 'baggage' => 'userId=alice' }, context.w3c) + end + + def test_w3c_returns_every_allowlisted_field + context = build_context( + 'custom' => { 'value' => 'test' }, + 'w3c' => { + 'traceparent' => '00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01', + 'tracestate' => 'rojo=00f067aa0ba902b7', + 'baggage' => 'userId=alice' + } + ) + assert_equal( + { + 'traceparent' => '00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01', + 'tracestate' => 'rojo=00f067aa0ba902b7', + 'baggage' => 'userId=alice' + }, + context.w3c + ) + end + + def test_w3c_is_stripped_from_client_context_but_siblings_are_kept + context = build_context( + 'custom' => { 'value' => 'test' }, + 'w3c' => { 'baggage' => 'userId=alice' } + ) + refute context.client_context.key?('w3c') + assert_equal({ 'custom' => { 'value' => 'test' } }, context.client_context) + end + + def test_w3c_drops_non_string_values_and_non_allowlisted_keys + context = build_context( + 'w3c' => { + 'baggage' => 'keep=me', + 'traceparent' => 42, + 'tracestate' => nil, + 'unknownField' => 'should-not-appear' + } + ) + assert_equal({ 'baggage' => 'keep=me' }, context.w3c) + end + + def test_w3c_treats_non_hash_values_as_empty + assert_equal({}, build_context('w3c' => 'not-an-object').w3c) + assert_equal({}, build_context('w3c' => ['baggage=abc']).w3c) + end + + def test_w3c_returns_a_fresh_copy_each_call + context = build_context({ 'w3c' => { 'baggage' => 'userId=alice' } }) + first = context.w3c + first['baggage'] = 'tampered' + first['injected'] = 'nope' + assert_equal({ 'baggage' => 'userId=alice' }, context.w3c) + end +end