Skip to content

Commit 974dfa5

Browse files
maxdaymaxday
committed
feat: w3c fields
1 parent 33fa573 commit 974dfa5

4 files changed

Lines changed: 313 additions & 1 deletion

File tree

‎lib/aws_lambda_ric/lambda_context.rb‎

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,11 @@
11
# frozen_string_literal: true
22

33
class LambdaContext
4+
# Allowlist of W3C trace-context fields that may be surfaced through
5+
# LambdaContext#w3c. Any other key carried on clientContext.w3c is ignored,
6+
# and any allowlisted key whose value is not a string is dropped.
7+
W3C_ALLOWED_FIELDS = %w[traceparent tracestate baggage].freeze
8+
49
attr_reader :aws_request_id, :invoked_function_arn, :log_group_name,
510
:log_stream_name, :function_name, :memory_limit_in_mb, :function_version,
611
:identity, :tenant_id, :client_context, :deadline_ms
@@ -17,12 +22,40 @@ def initialize(request)
1722
@function_version = ENV['AWS_LAMBDA_FUNCTION_VERSION']
1823
@identity = JSON.parse(request['Lambda-Runtime-Cognito-Identity']) unless request['Lambda-Runtime-Cognito-Identity'].to_s.empty?
1924
@tenant_id = request['Lambda-Runtime-Aws-Tenant-Id'] unless request['Lambda-Runtime-Aws-Tenant-Id'].to_s.empty?
20-
@client_context = JSON.parse(request['Lambda-Runtime-Client-Context']) unless request['Lambda-Runtime-Client-Context'].to_s.empty?
25+
@w3c_fields = {}
26+
unless request['Lambda-Runtime-Client-Context'].to_s.empty?
27+
client_context = JSON.parse(request['Lambda-Runtime-Client-Context'])
28+
@w3c_fields = self.class.extract_and_strip_w3c(client_context)
29+
@client_context = client_context
30+
end
2131
end
2232

2333
def get_remaining_time_in_millis
2434
now = Process.clock_gettime(Process::CLOCK_MONOTONIC, :millisecond) + @clock_diff
2535
remaining = @deadline_ms - now
2636
remaining.positive? ? remaining : 0
2737
end
38+
39+
# Return the W3C trace-context captured at invoke time, as a fresh copy so
40+
# callers cannot mutate the context's internal state.
41+
def w3c
42+
@w3c_fields.dup
43+
end
44+
45+
# Pop "w3c" out of the parsed client_context hash and return a normalized
46+
# copy of its allowlisted string fields (see W3C_ALLOWED_FIELDS).
47+
def self.extract_and_strip_w3c(client_context)
48+
return {} unless client_context.is_a?(Hash)
49+
return {} unless client_context.key?('w3c')
50+
51+
raw_w3c = client_context.delete('w3c')
52+
return {} unless raw_w3c.is_a?(Hash)
53+
54+
fields = {}
55+
W3C_ALLOWED_FIELDS.each do |key|
56+
value = raw_w3c[key]
57+
fields[key] = value if value.is_a?(String)
58+
end
59+
fields
60+
end
2861
end

‎test/dockerized/suites/w3c.json‎

Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,165 @@
1+
{
2+
"tests": [
3+
{
4+
"name": "w3c_is_callable_on_context",
5+
"handler": "w3c.w3c_is_callable",
6+
"request": {},
7+
"assertions": [
8+
{ "response": { "isCallable": true } }
9+
]
10+
},
11+
{
12+
"name": "w3c_returns_empty_when_no_client_context_header",
13+
"handler": "w3c.get_w3c",
14+
"request": {},
15+
"assertions": [
16+
{ "response": {} }
17+
]
18+
},
19+
{
20+
"name": "w3c_returns_empty_when_client_context_has_no_w3c_key",
21+
"handler": "w3c.get_w3c",
22+
"request": {},
23+
"clientContext": {
24+
"custom": { "value": "test" }
25+
},
26+
"assertions": [
27+
{ "response": {} }
28+
]
29+
},
30+
{
31+
"name": "w3c_returns_baggage_only",
32+
"handler": "w3c.get_w3c",
33+
"request": {},
34+
"clientContext": {
35+
"w3c": { "baggage": "userId=alice" }
36+
},
37+
"assertions": [
38+
{ "response": { "baggage": "userId=alice" } }
39+
]
40+
},
41+
{
42+
"name": "w3c_returns_all_three_allowlisted_fields",
43+
"handler": "w3c.get_w3c",
44+
"request": {},
45+
"clientContext": {
46+
"w3c": {
47+
"traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01",
48+
"tracestate": "rojo=00f067aa0ba902b7",
49+
"baggage": "userId=alice"
50+
}
51+
},
52+
"assertions": [
53+
{
54+
"response": {
55+
"traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01",
56+
"tracestate": "rojo=00f067aa0ba902b7",
57+
"baggage": "userId=alice"
58+
}
59+
}
60+
]
61+
},
62+
{
63+
"name": "w3c_allowlist_drops_non_allowlisted_keys",
64+
"handler": "w3c.get_w3c",
65+
"request": {},
66+
"clientContext": {
67+
"w3c": {
68+
"baggage": "keep=me",
69+
"unknownField": "should-not-appear",
70+
"x-custom-trace": "should-not-appear"
71+
}
72+
},
73+
"assertions": [
74+
{ "response": { "baggage": "keep=me" } }
75+
]
76+
},
77+
{
78+
"name": "w3c_drops_allowlisted_fields_with_non_string_values",
79+
"handler": "w3c.get_w3c",
80+
"request": {},
81+
"clientContext": {
82+
"w3c": {
83+
"traceparent": 42,
84+
"tracestate": null,
85+
"baggage": { "nested": "no" }
86+
}
87+
},
88+
"assertions": [
89+
{ "response": {} }
90+
]
91+
},
92+
{
93+
"name": "w3c_treats_non_object_as_empty",
94+
"handler": "w3c.get_w3c",
95+
"request": {},
96+
"clientContext": {
97+
"w3c": "not-an-object"
98+
},
99+
"assertions": [
100+
{ "response": {} }
101+
]
102+
},
103+
{
104+
"name": "w3c_treats_array_as_empty",
105+
"handler": "w3c.get_w3c",
106+
"request": {},
107+
"clientContext": {
108+
"w3c": ["baggage=abc"]
109+
},
110+
"assertions": [
111+
{ "response": {} }
112+
]
113+
},
114+
{
115+
"name": "w3c_strips_source_client_context_w3c_after_construction",
116+
"handler": "w3c.get_w3c_and_source",
117+
"request": {},
118+
"clientContext": {
119+
"custom": { "value": "test" },
120+
"w3c": {
121+
"traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01",
122+
"baggage": "userId=alice"
123+
}
124+
},
125+
"assertions": [
126+
{
127+
"response": {
128+
"w3c": {
129+
"traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01",
130+
"baggage": "userId=alice"
131+
},
132+
"clientContextIsDefined": true,
133+
"clientContextHasW3c": false,
134+
"clientContext": { "custom": { "value": "test" } }
135+
}
136+
}
137+
]
138+
},
139+
{
140+
"name": "client_context_is_echoed_when_no_w3c_key",
141+
"handler": "w3c.echo_client_context",
142+
"request": {},
143+
"clientContext": {
144+
"custom": { "value": "hello" },
145+
"env": { "stage": "beta" }
146+
},
147+
"assertions": [
148+
{
149+
"response": {
150+
"custom": { "value": "hello" },
151+
"env": { "stage": "beta" }
152+
}
153+
}
154+
]
155+
},
156+
{
157+
"name": "client_context_is_null_when_header_absent",
158+
"handler": "w3c.echo_client_context",
159+
"request": {},
160+
"assertions": [
161+
{ "response": null }
162+
]
163+
}
164+
]
165+
}

‎test/dockerized/tasks/w3c.rb‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
# Copyright 2026 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
def get_w3c(event:, context:)
5+
context.w3c
6+
end
7+
8+
def get_w3c_and_source(event:, context:)
9+
client_context = context.client_context
10+
{
11+
w3c: context.w3c,
12+
clientContextIsDefined: !client_context.nil?,
13+
clientContextHasW3c: !client_context.nil? && client_context.key?('w3c'),
14+
clientContext: client_context
15+
}
16+
end
17+
18+
def echo_client_context(event:, context:)
19+
context.client_context
20+
end
21+
22+
def w3c_is_callable(event:, context:)
23+
{ isCallable: context.respond_to?(:w3c) }
24+
end

‎test/unit/lambda_context_w3c_test.rb‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
# frozen_string_literal: true
2+
3+
require 'json'
4+
require_relative '../../lib/aws_lambda_ric/lambda_context'
5+
require 'minitest/autorun'
6+
require 'test/unit/assertions'
7+
8+
include Test::Unit::Assertions
9+
10+
class LambdaContextW3CTest < Minitest::Test
11+
def build_context(client_context_hash)
12+
request = {
13+
'Lambda-Runtime-Aws-Request-Id' => 'invoke-id-w3c',
14+
'Lambda-Runtime-Deadline-Ms' => '0',
15+
'Lambda-Runtime-Invoked-Function-Arn' => 'arn:test:w3c'
16+
}
17+
unless client_context_hash.nil?
18+
request['Lambda-Runtime-Client-Context'] = JSON.generate(client_context_hash)
19+
end
20+
LambdaContext.new(request)
21+
end
22+
23+
def test_w3c_returns_empty_when_no_client_context
24+
assert_equal({}, build_context(nil).w3c)
25+
end
26+
27+
def test_w3c_returns_empty_and_leaves_client_context_untouched_when_no_w3c_key
28+
context = build_context({ 'custom' => { 'value' => 'test' } })
29+
assert_equal({}, context.w3c)
30+
assert_equal({ 'custom' => { 'value' => 'test' } }, context.client_context)
31+
end
32+
33+
def test_w3c_returns_baggage_only
34+
context = build_context({ 'w3c' => { 'baggage' => 'userId=alice' } })
35+
assert_equal({ 'baggage' => 'userId=alice' }, context.w3c)
36+
end
37+
38+
def test_w3c_returns_every_allowlisted_field
39+
context = build_context(
40+
'custom' => { 'value' => 'test' },
41+
'w3c' => {
42+
'traceparent' => '00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01',
43+
'tracestate' => 'rojo=00f067aa0ba902b7',
44+
'baggage' => 'userId=alice'
45+
}
46+
)
47+
assert_equal(
48+
{
49+
'traceparent' => '00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01',
50+
'tracestate' => 'rojo=00f067aa0ba902b7',
51+
'baggage' => 'userId=alice'
52+
},
53+
context.w3c
54+
)
55+
end
56+
57+
def test_w3c_is_stripped_from_client_context_but_siblings_are_kept
58+
context = build_context(
59+
'custom' => { 'value' => 'test' },
60+
'w3c' => { 'baggage' => 'userId=alice' }
61+
)
62+
refute context.client_context.key?('w3c')
63+
assert_equal({ 'custom' => { 'value' => 'test' } }, context.client_context)
64+
end
65+
66+
def test_w3c_drops_non_string_values_and_non_allowlisted_keys
67+
context = build_context(
68+
'w3c' => {
69+
'baggage' => 'keep=me',
70+
'traceparent' => 42,
71+
'tracestate' => nil,
72+
'unknownField' => 'should-not-appear'
73+
}
74+
)
75+
assert_equal({ 'baggage' => 'keep=me' }, context.w3c)
76+
end
77+
78+
def test_w3c_treats_non_hash_values_as_empty
79+
assert_equal({}, build_context('w3c' => 'not-an-object').w3c)
80+
assert_equal({}, build_context('w3c' => ['baggage=abc']).w3c)
81+
end
82+
83+
def test_w3c_returns_a_fresh_copy_each_call
84+
context = build_context({ 'w3c' => { 'baggage' => 'userId=alice' } })
85+
first = context.w3c
86+
first['baggage'] = 'tampered'
87+
first['injected'] = 'nope'
88+
assert_equal({ 'baggage' => 'userId=alice' }, context.w3c)
89+
end
90+
end

0 commit comments

Comments
 (0)