Skip to content

Commit ec45db0

Browse files
Add tests for guard.WasmGuard.LabelResponse
Cover previously untested branches in WasmGuard.LabelResponse (wasm.go:610): - Context state extraction with tool_args (lines 618-624) - Non-map context state ignored gracefully (line 619 false branch) - Non-nil capabilities added to WASM input (lines 625-627) - labeled_paths JSON response path via parsePathLabeledResponse (lines 647-649) - items JSON response path via parseCollectionLabeledData (lines 652-654) - Final no-labeling return nil,nil (line 657) Three new WASM fixtures written using i32.store instructions write hardcoded JSON payloads to the output buffer at runtime: - labelResponseWritesEmptyObjectWasm → writes '{}' - labelResponseWritesItemsWasm → writes '{"items":[{"data":"x"}]}' - labelResponseWritesLabeledPathsWasm → writes labeled_paths/items_path response Coverage for WasmGuard.LabelResponse: 59.1% → 100.0% Guard package overall: 92.5% → 93.6% Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent 2af6d3e commit ec45db0

1 file changed

Lines changed: 185 additions & 0 deletions

File tree

Lines changed: 185 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,185 @@
1+
package guard
2+
3+
import (
4+
"context"
5+
"testing"
6+
7+
"github.com/stretchr/testify/assert"
8+
"github.com/stretchr/testify/require"
9+
10+
"github.com/github/gh-aw-mcpg/internal/difc"
11+
)
12+
13+
// labelResponseWritesEmptyObjectWasm exports "label_response" and "memory".
14+
// It writes the following JSON to the output buffer and returns its length:
15+
//
16+
// {}
17+
//
18+
// Used to exercise the 'no labeled_paths, no items → return nil, nil' path in LabelResponse.
19+
// The function writes 4 bytes (little-endian i32) and returns 2 so only the '{' '}' bytes are read.
20+
var labelResponseWritesEmptyObjectWasm = []byte{
21+
0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00, 0x01, 0x09, 0x01, 0x60, 0x04, 0x7f, 0x7f, 0x7f,
22+
0x7f, 0x01, 0x7f, 0x03, 0x02, 0x01, 0x00, 0x05, 0x03, 0x01, 0x00, 0x01, 0x07, 0x1b, 0x02, 0x0e,
23+
0x6c, 0x61, 0x62, 0x65, 0x6c, 0x5f, 0x72, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x00, 0x00,
24+
0x06, 0x6d, 0x65, 0x6d, 0x6f, 0x72, 0x79, 0x02, 0x00, 0x0a, 0x0f, 0x01, 0x0d, 0x00, 0x20, 0x02,
25+
0x41, 0xfb, 0xfa, 0x01, 0x36, 0x00, 0x00, 0x41, 0x02, 0x0b,
26+
}
27+
28+
// labelResponseWritesItemsWasm exports "label_response" and "memory".
29+
// It writes the following JSON to the output buffer and returns its length:
30+
//
31+
// {"items":[{"data":"x"}]}
32+
//
33+
// Used to exercise the 'items' collection labeling path in LabelResponse.
34+
var labelResponseWritesItemsWasm = []byte{
35+
0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00, 0x01, 0x09, 0x01, 0x60, 0x04, 0x7f, 0x7f, 0x7f,
36+
0x7f, 0x01, 0x7f, 0x03, 0x02, 0x01, 0x00, 0x05, 0x03, 0x01, 0x00, 0x01, 0x07, 0x1b, 0x02, 0x0e,
37+
0x6c, 0x61, 0x62, 0x65, 0x6c, 0x5f, 0x72, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x00, 0x00,
38+
0x06, 0x6d, 0x65, 0x6d, 0x6f, 0x72, 0x79, 0x02, 0x00, 0x0a, 0x48, 0x01, 0x46, 0x00, 0x20, 0x02,
39+
0x41, 0xfb, 0xc4, 0xa4, 0xa3, 0x07, 0x36, 0x00, 0x00, 0x20, 0x02, 0x41, 0xe5, 0xda, 0xcd, 0x93,
40+
0x02, 0x36, 0x00, 0x04, 0x20, 0x02, 0x41, 0xba, 0xb6, 0xed, 0x93, 0x02, 0x36, 0x00, 0x08, 0x20,
41+
0x02, 0x41, 0xe4, 0xc2, 0xd1, 0x8b, 0x06, 0x36, 0x00, 0x0c, 0x20, 0x02, 0x41, 0xa2, 0xf4, 0x88,
42+
0xc1, 0x07, 0x36, 0x00, 0x10, 0x20, 0x02, 0x41, 0xa2, 0xfa, 0xf5, 0xea, 0x07, 0x36, 0x00, 0x14,
43+
0x41, 0x18, 0x0b,
44+
}
45+
46+
// labelResponseWritesLabeledPathsWasm exports "label_response" and "memory".
47+
// It writes the following JSON to the output buffer and returns its length:
48+
//
49+
// {"labeled_paths":[],"items_path":"/items"}
50+
//
51+
// Used to exercise the 'labeled_paths' response path in LabelResponse.
52+
// When called with a result that does not contain '/items', parsePathLabeledResponse
53+
// returns an error that is propagated from LabelResponse.
54+
var labelResponseWritesLabeledPathsWasm = []byte{
55+
0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00, 0x01, 0x09, 0x01, 0x60, 0x04, 0x7f, 0x7f, 0x7f,
56+
0x7f, 0x01, 0x7f, 0x03, 0x02, 0x01, 0x00, 0x05, 0x03, 0x01, 0x00, 0x01, 0x07, 0x1b, 0x02, 0x0e,
57+
0x6c, 0x61, 0x62, 0x65, 0x6c, 0x5f, 0x72, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x00, 0x00,
58+
0x06, 0x6d, 0x65, 0x6d, 0x6f, 0x72, 0x79, 0x02, 0x00, 0x0a, 0x7d, 0x01, 0x7b, 0x00, 0x20, 0x02,
59+
0x41, 0xfb, 0xc4, 0xb0, 0x8b, 0x06, 0x36, 0x00, 0x00, 0x20, 0x02, 0x41, 0xe2, 0xca, 0xb1, 0xab,
60+
0x06, 0x36, 0x00, 0x04, 0x20, 0x02, 0x41, 0xe4, 0xbe, 0xc1, 0x8b, 0x06, 0x36, 0x00, 0x08, 0x20,
61+
0x02, 0x41, 0xf4, 0xd0, 0xcd, 0x93, 0x02, 0x36, 0x00, 0x0c, 0x20, 0x02, 0x41, 0xba, 0xb6, 0xf5,
62+
0xe2, 0x02, 0x36, 0x00, 0x10, 0x20, 0x02, 0x41, 0xa2, 0xd2, 0xd1, 0xab, 0x06, 0x36, 0x00, 0x14,
63+
0x20, 0x02, 0x41, 0xed, 0xe6, 0xfd, 0x82, 0x07, 0x36, 0x00, 0x18, 0x20, 0x02, 0x41, 0xe1, 0xe8,
64+
0xa1, 0x93, 0x02, 0x36, 0x00, 0x1c, 0x20, 0x02, 0x41, 0xba, 0xc4, 0xbc, 0xc9, 0x06, 0x36, 0x00,
65+
0x20, 0x20, 0x02, 0x41, 0xf4, 0xca, 0xb5, 0x9b, 0x07, 0x36, 0x00, 0x24, 0x20, 0x02, 0x41, 0xa2,
66+
0xfa, 0x01, 0x36, 0x00, 0x28, 0x41, 0x2a, 0x0b,
67+
}
68+
69+
// --- TestLabelResponse: context state and capabilities input-preparation paths ---
70+
71+
// TestLabelResponse_ContextStateWithToolArgs verifies that when the context contains
72+
// a request state map with a "tool_args" key, LabelResponse extracts it into the WASM
73+
// input and still returns (nil, nil) when the WASM guard returns an empty result.
74+
func TestLabelResponse_ContextStateWithToolArgs(t *testing.T) {
75+
g, cleanup := setupRawWasmModule(t, labelResponseReturnsZeroWasm, "lresp-state-tool-args")
76+
defer cleanup()
77+
78+
ctx := SetRequestStateInContext(context.Background(), map[string]any{
79+
"tool_args": map[string]any{"param": "value"},
80+
})
81+
82+
result, err := g.LabelResponse(ctx, "my_tool", nil, nil, nil)
83+
require.NoError(t, err)
84+
assert.Nil(t, result, "WASM returning empty result should yield nil even with state in context")
85+
}
86+
87+
// TestLabelResponse_ContextStateNonMapIsIgnored verifies that when the context contains
88+
// a request state that is not a map[string]any, LabelResponse ignores the tool_args
89+
// extraction and proceeds normally.
90+
func TestLabelResponse_ContextStateNonMapIsIgnored(t *testing.T) {
91+
g, cleanup := setupRawWasmModule(t, labelResponseReturnsZeroWasm, "lresp-state-non-map")
92+
defer cleanup()
93+
94+
// Pass a non-map state (e.g., a plain string). The type assertion in LabelResponse
95+
// will fail silently and tool_args will not be added to the WASM input.
96+
ctx := SetRequestStateInContext(context.Background(), "not-a-map")
97+
98+
result, err := g.LabelResponse(ctx, "my_tool", nil, nil, nil)
99+
require.NoError(t, err)
100+
assert.Nil(t, result)
101+
}
102+
103+
// TestLabelResponse_WithCapabilities verifies that when a non-nil *difc.Capabilities is
104+
// provided, LabelResponse adds it to the WASM input and still returns (nil, nil) when
105+
// the WASM guard returns an empty result.
106+
func TestLabelResponse_WithCapabilities(t *testing.T) {
107+
g, cleanup := setupRawWasmModule(t, labelResponseReturnsZeroWasm, "lresp-caps")
108+
defer cleanup()
109+
110+
caps := difc.NewCapabilities()
111+
caps.Add(difc.Tag("public"))
112+
113+
result, err := g.LabelResponse(context.Background(), "my_tool", nil, nil, caps)
114+
require.NoError(t, err)
115+
assert.Nil(t, result, "WASM returning empty result should yield nil even with capabilities set")
116+
}
117+
118+
// --- TestLabelResponse: JSON response parsing paths ---
119+
120+
// TestLabelResponse_NoLabelingResponse verifies that when the WASM guard returns a
121+
// valid JSON object that contains neither "labeled_paths" nor "items", LabelResponse
122+
// falls through to the final "no fine-grained labeling" return and returns (nil, nil).
123+
func TestLabelResponse_NoLabelingResponse(t *testing.T) {
124+
// labelResponseWritesEmptyObjectWasm writes "{}" which has no labeling keys.
125+
g, cleanup := setupRawWasmModule(t, labelResponseWritesEmptyObjectWasm, "lresp-no-labeling")
126+
defer cleanup()
127+
128+
result, err := g.LabelResponse(context.Background(), "my_tool", nil, nil, nil)
129+
require.NoError(t, err)
130+
assert.Nil(t, result, "response with no labeling keys should return nil LabeledData")
131+
}
132+
133+
// TestLabelResponse_ItemsResponse verifies that when the WASM guard returns a JSON
134+
// object containing a non-empty "items" array, LabelResponse parses it via
135+
// parseCollectionLabeledData and returns a populated *difc.CollectionLabeledData.
136+
func TestLabelResponse_ItemsResponse(t *testing.T) {
137+
// labelResponseWritesItemsWasm writes: {"items":[{"data":"x"}]}
138+
g, cleanup := setupRawWasmModule(t, labelResponseWritesItemsWasm, "lresp-items")
139+
defer cleanup()
140+
141+
result, err := g.LabelResponse(context.Background(), "my_tool", nil, nil, nil)
142+
require.NoError(t, err)
143+
require.NotNil(t, result, "items response should produce non-nil LabeledData")
144+
145+
collection, ok := result.(*difc.CollectionLabeledData)
146+
require.True(t, ok, "items response should return *CollectionLabeledData")
147+
require.Len(t, collection.Items, 1, "expected exactly one labeled item")
148+
assert.Equal(t, "x", collection.Items[0].Data, "item data should match the WASM-written value")
149+
}
150+
151+
// TestLabelResponse_LabeledPathsResponse verifies that when the WASM guard returns a
152+
// JSON object containing a "labeled_paths" key, LabelResponse delegates to
153+
// parsePathLabeledResponse. When items_path cannot be resolved against the supplied
154+
// original result, parsePathLabeledResponse returns an error that is propagated.
155+
func TestLabelResponse_LabeledPathsResponse(t *testing.T) {
156+
// labelResponseWritesLabeledPathsWasm writes:
157+
// {"labeled_paths":[],"items_path":"/items"}
158+
// The original result below does not contain "/items", so parsePathLabeledResponse
159+
// will return an error.
160+
g, cleanup := setupRawWasmModule(t, labelResponseWritesLabeledPathsWasm, "lresp-labeled-paths")
161+
defer cleanup()
162+
163+
originalResult := map[string]any{"other": "value"}
164+
result, err := g.LabelResponse(context.Background(), "my_tool", originalResult, nil, nil)
165+
166+
require.Error(t, err, "labeled_paths response with unresolvable items_path should return an error")
167+
assert.Nil(t, result)
168+
assert.ErrorContains(t, err, "apply path labels")
169+
}
170+
171+
// TestLabelResponse_ContextStateAndCapabilities verifies that both context state and
172+
// capabilities are handled together, and that the WASM call completes successfully.
173+
func TestLabelResponse_ContextStateAndCapabilities(t *testing.T) {
174+
g, cleanup := setupRawWasmModule(t, labelResponseReturnsZeroWasm, "lresp-state-and-caps")
175+
defer cleanup()
176+
177+
ctx := SetRequestStateInContext(context.Background(), map[string]any{
178+
"tool_args": map[string]any{"arg1": "val1"},
179+
})
180+
caps := difc.NewCapabilities()
181+
182+
result, err := g.LabelResponse(ctx, "my_tool", "some-result", nil, caps)
183+
require.NoError(t, err)
184+
assert.Nil(t, result)
185+
}

0 commit comments

Comments
 (0)