feat: w3c fields

.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'

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,12 +22,40 @@ 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
     now = Process.clock_gettime(Process::CLOCK_MONOTONIC, :millisecond) + @clock_diff
     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

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 }
+            ]
+        }
+    ]
+}

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

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