Skip to content

Commit e43928f

Browse files
abhipatel12abhipatel12
committed
fix(core): remove FQN string splitting vulnerability from policy engine
1 parent 18b97aa commit e43928f

2 files changed

Lines changed: 151 additions & 176 deletions

File tree

‎packages/core/src/policy/policy-engine.test.ts‎

Lines changed: 105 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -2122,18 +2122,22 @@ describe('PolicyEngine', () => {
21222122
rules: PolicyRule[];
21232123
approvalMode?: ApprovalMode;
21242124
nonInteractive?: boolean;
2125+
allToolNames?: string[];
2126+
metadata?: Map<string, Record<string, unknown>>;
21252127
expected: string[];
21262128
}
21272129

21282130
const testCases: TestCase[] = [
21292131
{
21302132
name: 'should return empty set when no rules provided',
21312133
rules: [],
2134+
allToolNames: ['tool1'],
21322135
expected: [],
21332136
},
21342137
{
21352138
name: 'should apply rules without explicit modes to all modes',
21362139
rules: [{ toolName: 'tool1', decision: PolicyDecision.DENY }],
2140+
allToolNames: ['tool1', 'tool2'],
21372141
expected: ['tool1'],
21382142
},
21392143
{
@@ -2153,6 +2157,7 @@ describe('PolicyEngine', () => {
21532157
modes: [ApprovalMode.DEFAULT],
21542158
},
21552159
],
2160+
allToolNames: ['tool1'],
21562161
expected: [],
21572162
},
21582163
{
@@ -2169,6 +2174,7 @@ describe('PolicyEngine', () => {
21692174
modes: [ApprovalMode.DEFAULT],
21702175
},
21712176
],
2177+
allToolNames: ['tool1', 'tool2', 'tool3'],
21722178
expected: ['tool1'],
21732179
},
21742180
{
@@ -2187,6 +2193,7 @@ describe('PolicyEngine', () => {
21872193
modes: [ApprovalMode.DEFAULT],
21882194
},
21892195
],
2196+
allToolNames: ['tool1'],
21902197
expected: ['tool1'],
21912198
},
21922199
{
@@ -2205,6 +2212,7 @@ describe('PolicyEngine', () => {
22052212
modes: [ApprovalMode.DEFAULT],
22062213
},
22072214
],
2215+
allToolNames: ['tool1'],
22082216
expected: [],
22092217
},
22102218
{
@@ -2217,7 +2225,8 @@ describe('PolicyEngine', () => {
22172225
},
22182226
],
22192227
nonInteractive: true,
2220-
expected: [],
2228+
allToolNames: ['tool1'],
2229+
expected: ['tool1'],
22212230
},
22222231
{
22232232
name: 'should ignore rules with argsPattern',
@@ -2229,6 +2238,7 @@ describe('PolicyEngine', () => {
22292238
modes: [ApprovalMode.DEFAULT],
22302239
},
22312240
],
2241+
allToolNames: ['tool1'],
22322242
expected: [],
22332243
},
22342244
{
@@ -2241,6 +2251,7 @@ describe('PolicyEngine', () => {
22412251
},
22422252
],
22432253
approvalMode: ApprovalMode.PLAN,
2254+
allToolNames: ['tool1'],
22442255
expected: ['tool1'],
22452256
},
22462257
{
@@ -2253,6 +2264,7 @@ describe('PolicyEngine', () => {
22532264
},
22542265
],
22552266
approvalMode: ApprovalMode.DEFAULT,
2267+
allToolNames: ['tool1'],
22562268
expected: [],
22572269
},
22582270
{
@@ -2271,6 +2283,7 @@ describe('PolicyEngine', () => {
22712283
},
22722284
],
22732285
approvalMode: ApprovalMode.YOLO,
2286+
allToolNames: ['dangerous-tool', 'safe-tool'],
22742287
expected: [],
22752288
},
22762289
{
@@ -2283,7 +2296,17 @@ describe('PolicyEngine', () => {
22832296
modes: [ApprovalMode.DEFAULT],
22842297
},
22852298
],
2286-
expected: ['mcp_server_*'],
2299+
allToolNames: [
2300+
'mcp_server_tool1',
2301+
'mcp_server_tool2',
2302+
'mcp_other_tool',
2303+
],
2304+
metadata: new Map([
2305+
['mcp_server_tool1', { _serverName: 'server' }],
2306+
['mcp_server_tool2', { _serverName: 'server' }],
2307+
['mcp_other_tool', { _serverName: 'other' }],
2308+
]),
2309+
expected: ['mcp_server_tool1', 'mcp_server_tool2'],
22872310
},
22882311
{
22892312
name: 'should expand server wildcard for specific tools if already processed',
@@ -2298,12 +2321,17 @@ describe('PolicyEngine', () => {
22982321
{
22992322
toolName: 'mcp_server_tool1',
23002323
mcpName: 'server',
2301-
decision: PolicyDecision.DENY,
2324+
decision: PolicyDecision.DENY, // redundant but tests ordering
23022325
priority: 10,
23032326
modes: [ApprovalMode.DEFAULT],
23042327
},
23052328
],
2306-
expected: ['mcp_server_*', 'mcp_server_tool1'],
2329+
allToolNames: ['mcp_server_tool1', 'mcp_server_tool2'],
2330+
metadata: new Map([
2331+
['mcp_server_tool1', { _serverName: 'server' }],
2332+
['mcp_server_tool2', { _serverName: 'server' }],
2333+
]),
2334+
expected: ['mcp_server_tool1', 'mcp_server_tool2'],
23072335
},
23082336
{
23092337
name: 'should exclude run_shell_command but NOT write_file in simulated Plan Mode',
@@ -2330,7 +2358,8 @@ describe('PolicyEngine', () => {
23302358
priority: 10,
23312359
},
23322360
],
2333-
expected: ['run_shell_command'],
2361+
allToolNames: ['write_file', 'run_shell_command', 'read_file'],
2362+
expected: ['run_shell_command', 'read_file'],
23342363
},
23352364
{
23362365
name: 'should NOT exclude tool if covered by a higher priority wildcard ALLOW',
@@ -2350,6 +2379,8 @@ describe('PolicyEngine', () => {
23502379
modes: [ApprovalMode.DEFAULT],
23512380
},
23522381
],
2382+
allToolNames: ['mcp_server_tool1'],
2383+
metadata: new Map([['mcp_server_tool1', { _serverName: 'server' }]]),
23532384
expected: [],
23542385
},
23552386
{
@@ -2361,41 +2392,63 @@ describe('PolicyEngine', () => {
23612392
priority: 10,
23622393
},
23632394
],
2364-
expected: ['*'],
2395+
allToolNames: ['toolA', 'toolB', 'mcp_server_toolC'],
2396+
expected: ['toolA', 'toolB', 'mcp_server_toolC'], // all tools denied by *
23652397
},
23662398
{
23672399
name: 'should handle MCP category wildcard *__* in getExcludedTools',
23682400
rules: [
23692401
{
2370-
toolName: '*__*',
2402+
toolName: 'mcp_*',
23712403
decision: PolicyDecision.DENY,
23722404
priority: 10,
23732405
},
23742406
],
2375-
expected: ['*__*'],
2407+
allToolNames: ['localTool', 'mcp_myserver_mytool'],
2408+
metadata: new Map([
2409+
['mcp_myserver_mytool', { _serverName: 'myserver' }],
2410+
]),
2411+
expected: ['mcp_myserver_mytool'],
23762412
},
23772413
{
2378-
name: 'should handle tool wildcard *__search in getExcludedTools',
2414+
name: 'should handle tool wildcard mcp_server_* in getExcludedTools',
23792415
rules: [
23802416
{
2381-
toolName: '*__search',
2417+
toolName: 'mcp_server_*',
23822418
decision: PolicyDecision.DENY,
23832419
priority: 10,
23842420
},
23852421
],
2386-
expected: ['*__search'],
2422+
allToolNames: [
2423+
'localTool',
2424+
'mcp_server_search',
2425+
'mcp_otherserver_read',
2426+
],
2427+
metadata: new Map([
2428+
['mcp_server_search', { _serverName: 'server' }],
2429+
['mcp_otherserver_read', { _serverName: 'otherserver' }],
2430+
]),
2431+
expected: ['mcp_server_search'],
23872432
},
23882433
];
23892434

23902435
it.each(testCases)(
23912436
'$name',
2392-
({ rules, approvalMode, nonInteractive, expected }) => {
2437+
({
2438+
rules,
2439+
approvalMode,
2440+
nonInteractive,
2441+
allToolNames,
2442+
metadata,
2443+
expected,
2444+
}) => {
23932445
engine = new PolicyEngine({
23942446
rules,
23952447
approvalMode: approvalMode ?? ApprovalMode.DEFAULT,
23962448
nonInteractive: nonInteractive ?? false,
23972449
});
2398-
const excluded = engine.getExcludedTools();
2450+
const toolsSet = allToolNames ? new Set(allToolNames) : undefined;
2451+
const excluded = engine.getExcludedTools(metadata, toolsSet);
23992452
expect(Array.from(excluded).sort()).toEqual(expected.sort());
24002453
},
24012454
);
@@ -2410,7 +2463,10 @@ describe('PolicyEngine', () => {
24102463
},
24112464
],
24122465
});
2413-
const excluded = engine.getExcludedTools();
2466+
const excluded = engine.getExcludedTools(
2467+
undefined,
2468+
new Set(['dangerous_tool']),
2469+
);
24142470
expect(Array.from(excluded)).toEqual([]);
24152471
});
24162472

@@ -2428,7 +2484,10 @@ describe('PolicyEngine', () => {
24282484
['dangerous_tool', { destructiveHint: true }],
24292485
['safe_tool', { readOnlyHint: true }],
24302486
]);
2431-
const excluded = engine.getExcludedTools(metadata);
2487+
const excluded = engine.getExcludedTools(
2488+
metadata,
2489+
new Set(['dangerous_tool', 'safe_tool']),
2490+
);
24322491
expect(Array.from(excluded)).toEqual(['dangerous_tool']);
24332492
});
24342493

@@ -2445,7 +2504,10 @@ describe('PolicyEngine', () => {
24452504
const metadata = new Map<string, Record<string, unknown>>([
24462505
['safe_tool', { readOnlyHint: true }],
24472506
]);
2448-
const excluded = engine.getExcludedTools(metadata);
2507+
const excluded = engine.getExcludedTools(
2508+
metadata,
2509+
new Set(['safe_tool']),
2510+
);
24492511
expect(Array.from(excluded)).toEqual([]);
24502512
});
24512513

@@ -2462,11 +2524,24 @@ describe('PolicyEngine', () => {
24622524
],
24632525
});
24642526
const metadata = new Map<string, Record<string, unknown>>([
2465-
['mcp_server_dangerous_tool', { destructiveHint: true }],
2466-
['mcp_other_dangerous_tool', { destructiveHint: true }],
2467-
['mcp_server_safe_tool', { readOnlyHint: true }],
2527+
[
2528+
'mcp_server_dangerous_tool',
2529+
{ destructiveHint: true, _serverName: 'server' },
2530+
],
2531+
[
2532+
'mcp_other_dangerous_tool',
2533+
{ destructiveHint: true, _serverName: 'other' },
2534+
],
2535+
['mcp_server_safe_tool', { readOnlyHint: true, _serverName: 'server' }],
24682536
]);
2469-
const excluded = engine.getExcludedTools(metadata);
2537+
const excluded = engine.getExcludedTools(
2538+
metadata,
2539+
new Set([
2540+
'mcp_server_dangerous_tool',
2541+
'mcp_other_dangerous_tool',
2542+
'mcp_server_safe_tool',
2543+
]),
2544+
);
24702545
expect(Array.from(excluded)).toEqual(['mcp_server_dangerous_tool']);
24712546
});
24722547

@@ -2527,7 +2602,7 @@ describe('PolicyEngine', () => {
25272602
expect(excluded.has('mcp_my-server_read_mcp_tool')).toBe(false);
25282603
});
25292604

2530-
it('should match already-qualified MCP tool names without _serverName', () => {
2605+
it('should match MCP wildcard rules when explicitly mapped with _serverName', () => {
25312606
engine = new PolicyEngine({
25322607
rules: [
25332608
{
@@ -2548,11 +2623,17 @@ describe('PolicyEngine', () => {
25482623
'mcp_myserver_write_tool',
25492624
]);
25502625
const toolMetadata = new Map<string, Record<string, unknown>>([
2551-
['mcp_myserver_read_tool', { readOnlyHint: true }],
2552-
['mcp_myserver_write_tool', { readOnlyHint: false }],
2626+
[
2627+
'mcp_myserver_read_tool',
2628+
{ readOnlyHint: true, _serverName: 'myserver' },
2629+
],
2630+
[
2631+
'mcp_myserver_write_tool',
2632+
{ readOnlyHint: false, _serverName: 'myserver' },
2633+
],
25532634
]);
25542635
const excluded = engine.getExcludedTools(toolMetadata, allToolNames);
2555-
// Qualified name already contains __, matched directly without _serverName
2636+
// Qualified name matched using explicit _serverName
25562637
expect(excluded.has('mcp_myserver_read_tool')).toBe(false);
25572638
expect(excluded.has('mcp_myserver_write_tool')).toBe(true);
25582639
});

0 commit comments

Comments
 (0)