Skip to content

Commit e4e9bea

Browse files
aidandaly24aidandaly24
committed
feat(gateway): add custom claims validation and TUI wizard for JWT auth
Add custom JWT claims validation support and a full TUI wizard flow for configuring Custom JWT gateway authorization. Schema: - Add ClaimMatchOperator, ClaimMatchValue, InboundTokenClaimValueType, and CustomClaimValidation schemas with strict validation - Add customClaims to CustomJwtAuthorizerConfigSchema and deployed-state - Add --custom-claims CLI flag with JSON parsing and validation TUI Wizard: - Expand JWT config flow with custom claims manager (add/edit/done) - Add claim name, operator, value, and value type sub-steps - Show human-readable claim summary in confirm review - Make client credentials optional (skip with empty Enter) Testing: - Add AddGatewayJwtConfig.test.tsx — full TUI component tests - Add finishJwtConfig.test.ts — unit tests for config assembly - Extend useAddGatewayWizard.test.tsx with JWT + custom claims flows - Add GatewayPrimitive.test.ts for custom claims round-trip - Extend validate.test.ts with custom claims validation cases - Add TUI integration test (add-gateway-jwt.test.ts) Constraint: Stacked on fix/inbound-auth-hardening (aws#598) Confidence: high Scope-risk: moderate
1 parent 1d731e5 commit e4e9bea

File tree

17 files changed

+2339
-159
lines changed

17 files changed

+2339
-159
lines changed

‎integ-tests/tui/add-gateway-jwt.test.ts‎

Lines changed: 415 additions & 0 deletions
Large diffs are not rendered by default.

‎integ-tests/tui/setup.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
/**
2+
* TUI integration test setup.
3+
*
4+
* This file is referenced by vitest.config.ts as a setupFile for the 'tui' project.
5+
* It runs before each test file in integ-tests/tui/.
6+
*/

‎src/cli/commands/add/__tests__/validate.test.ts‎

Lines changed: 91 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -220,40 +220,29 @@ describe('validate', () => {
220220
expect(result.error?.includes('Invalid authorizer type')).toBeTruthy();
221221
});
222222

223-
// AC11: CUSTOM_JWT requires discoveryUrl
224-
it('returns error for CUSTOM_JWT missing discoveryUrl', () => {
225-
const opts = { ...validGatewayOptionsJwt, discoveryUrl: undefined };
226-
const result = validateAddGatewayOptions(opts);
223+
// AC11: CUSTOM_JWT requires discoveryUrl; at least one of allowedAudience/allowedClients/allowedScopes
224+
it('returns error for CUSTOM_JWT missing required fields', () => {
225+
// discoveryUrl is always required
226+
const result = validateAddGatewayOptions({ ...validGatewayOptionsJwt, discoveryUrl: undefined });
227227
expect(result.valid).toBe(false);
228228
expect(result.error).toBe('--discovery-url is required for CUSTOM_JWT authorizer');
229-
});
230229

231-
// AC11b: at least one of audience/clients/scopes required
232-
it('returns error when all of audience, clients, and scopes are missing', () => {
233-
const opts = {
230+
// All three optional fields absent fails
231+
const noneResult = validateAddGatewayOptions({
234232
...validGatewayOptionsJwt,
235233
allowedAudience: undefined,
236234
allowedClients: undefined,
237235
allowedScopes: undefined,
238-
};
239-
const result = validateAddGatewayOptions(opts);
240-
expect(result.valid).toBe(false);
241-
expect(result.error).toContain('At least one of');
242-
});
243-
244-
it('allows CUSTOM_JWT with only allowedScopes', () => {
245-
const opts = {
246-
...validGatewayOptionsJwt,
247-
allowedAudience: undefined,
248-
allowedClients: undefined,
249-
allowedScopes: 'scope1',
250-
};
251-
const result = validateAddGatewayOptions(opts);
252-
expect(result.valid).toBe(true);
236+
});
237+
expect(noneResult.valid).toBe(false);
238+
expect(noneResult.error).toBe(
239+
'At least one of --allowed-audience, --allowed-clients, --allowed-scopes, or --custom-claims must be provided for CUSTOM_JWT authorizer'
240+
);
253241
});
254242

255-
it('allows CUSTOM_JWT with only allowedAudience', () => {
256-
const opts = { ...validGatewayOptionsJwt, allowedClients: undefined, allowedScopes: undefined };
243+
// AC11b: allowedAudience is optional
244+
it('allows CUSTOM_JWT without allowedAudience', () => {
245+
const opts = { ...validGatewayOptionsJwt, allowedAudience: undefined };
257246
const result = validateAddGatewayOptions(opts);
258247
expect(result.valid).toBe(true);
259248
});
@@ -271,21 +260,88 @@ describe('validate', () => {
271260
expect(result.error?.includes('.well-known/openid-configuration')).toBeTruthy();
272261
});
273262

274-
it('returns error for HTTP discoveryUrl (HTTPS required)', () => {
263+
// AC13: At least one of audience/clients/scopes must be non-empty
264+
it('returns error when all of audience, clients, and scopes are empty', () => {
275265
const result = validateAddGatewayOptions({
276266
...validGatewayOptionsJwt,
277-
discoveryUrl: 'http://example.com/.well-known/openid-configuration',
267+
allowedAudience: ' ',
268+
allowedClients: undefined,
269+
allowedScopes: undefined,
278270
});
279271
expect(result.valid).toBe(false);
280-
expect(result.error).toBe('Discovery URL must use HTTPS');
272+
expect(result.error).toBe(
273+
'At least one of --allowed-audience, --allowed-clients, --allowed-scopes, or --custom-claims must be provided for CUSTOM_JWT authorizer'
274+
);
281275
});
282276

283-
it('allows CUSTOM_JWT with only allowedClients', () => {
284-
const opts = { ...validGatewayOptionsJwt, allowedAudience: undefined, allowedScopes: undefined };
285-
const result = validateAddGatewayOptions(opts);
277+
// AC-claims1: --custom-claims with valid JSON passes validation
278+
it('accepts valid --custom-claims JSON', () => {
279+
const result = validateAddGatewayOptions({
280+
...validGatewayOptionsJwt,
281+
customClaims: JSON.stringify([
282+
{
283+
inboundTokenClaimName: 'dept',
284+
inboundTokenClaimValueType: 'STRING',
285+
authorizingClaimMatchValue: {
286+
claimMatchOperator: 'EQUALS',
287+
claimMatchValue: { matchValueString: 'engineering' },
288+
},
289+
},
290+
]),
291+
});
286292
expect(result.valid).toBe(true);
287293
});
288294

295+
// AC-claims2: --custom-claims alone satisfies the "at least one constraint" check
296+
it('allows CUSTOM_JWT with only --custom-claims (no audience/clients/scopes)', () => {
297+
const result = validateAddGatewayOptions({
298+
name: 'test-gw',
299+
authorizerType: 'CUSTOM_JWT',
300+
discoveryUrl: 'https://example.com/.well-known/openid-configuration',
301+
customClaims: JSON.stringify([
302+
{
303+
inboundTokenClaimName: 'role',
304+
inboundTokenClaimValueType: 'STRING_ARRAY',
305+
authorizingClaimMatchValue: {
306+
claimMatchOperator: 'CONTAINS_ANY',
307+
claimMatchValue: { matchValueStringList: ['admin'] },
308+
},
309+
},
310+
]),
311+
});
312+
expect(result.valid).toBe(true);
313+
});
314+
315+
// AC-claims3: --custom-claims with invalid JSON fails
316+
it('returns error for --custom-claims with invalid JSON', () => {
317+
const result = validateAddGatewayOptions({
318+
...validGatewayOptionsJwt,
319+
customClaims: 'not json',
320+
});
321+
expect(result.valid).toBe(false);
322+
expect(result.error).toBe('--custom-claims must be valid JSON');
323+
});
324+
325+
// AC-claims4: --custom-claims with empty array fails
326+
it('returns error for --custom-claims with empty array', () => {
327+
const result = validateAddGatewayOptions({
328+
...validGatewayOptionsJwt,
329+
customClaims: '[]',
330+
});
331+
expect(result.valid).toBe(false);
332+
expect(result.error).toBe('--custom-claims must be a non-empty JSON array');
333+
});
334+
335+
// AC-claims5: --custom-claims with invalid claim structure fails
336+
it('returns error for --custom-claims with invalid claim structure', () => {
337+
const result = validateAddGatewayOptions({
338+
...validGatewayOptionsJwt,
339+
customClaims: JSON.stringify([{ badField: 'value' }]),
340+
});
341+
expect(result.valid).toBe(false);
342+
expect(result.error).toContain('Invalid custom claim at index 0');
343+
});
344+
289345
// AC14: Valid options pass
290346
it('passes for valid options', () => {
291347
expect(validateAddGatewayOptions(validGatewayOptionsNone)).toEqual({ valid: true });
@@ -311,8 +367,8 @@ describe('validate', () => {
311367
expect(result.error).toBe('Both --client-id and --client-secret must be provided together');
312368
});
313369

314-
// AC16: OAuth credentials only valid with CUSTOM_JWT
315-
it('returns error when OAuth credentials used with non-CUSTOM_JWT authorizer', () => {
370+
// AC16: OAuth client credentials only valid with CUSTOM_JWT
371+
it('returns error when OAuth client credentials used with non-CUSTOM_JWT authorizer', () => {
316372
const result = validateAddGatewayOptions({
317373
...validGatewayOptionsNone,
318374
clientId: 'my-client-id',
@@ -322,8 +378,8 @@ describe('validate', () => {
322378
expect(result.error).toBe('OAuth client credentials are only valid with CUSTOM_JWT authorizer');
323379
});
324380

325-
// AC17: valid CUSTOM_JWT with OAuth credentials passes
326-
it('passes for CUSTOM_JWT with OAuth credentials', () => {
381+
// AC17: valid CUSTOM_JWT with OAuth client credentials passes
382+
it('passes for CUSTOM_JWT with OAuth client credentials', () => {
327383
const result = validateAddGatewayOptions({
328384
...validGatewayOptionsJwt,
329385
clientId: 'my-client-id',

‎src/cli/commands/add/types.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ export interface AddGatewayOptions {
3737
allowedAudience?: string;
3838
allowedClients?: string;
3939
allowedScopes?: string;
40+
customClaims?: string;
4041
clientId?: string;
4142
clientSecret?: string;
4243
agents?: string;

‎src/cli/commands/add/validate.ts‎

Lines changed: 24 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import { ConfigIO, findConfigRoot } from '../../../lib';
22
import {
33
AgentNameSchema,
44
BuildTypeSchema,
5+
CustomClaimValidationSchema,
56
GatewayExceptionLevelSchema,
67
GatewayNameSchema,
78
ModelProviderSchema,
@@ -275,16 +276,36 @@ export function validateAddGatewayOptions(options: AddGatewayOptions): Validatio
275276
return { valid: false, error: `Discovery URL must end with ${OIDC_WELL_KNOWN_SUFFIX}` };
276277
}
277278

278-
// allowedAudience, allowedClients, allowedScopes are all optional individually,
279+
// Validate custom claims JSON if provided
280+
if (options.customClaims) {
281+
let parsed: unknown;
282+
try {
283+
parsed = JSON.parse(options.customClaims);
284+
} catch {
285+
return { valid: false, error: '--custom-claims must be valid JSON' };
286+
}
287+
if (!Array.isArray(parsed) || parsed.length === 0) {
288+
return { valid: false, error: '--custom-claims must be a non-empty JSON array' };
289+
}
290+
for (const [i, entry] of parsed.entries()) {
291+
const result = CustomClaimValidationSchema.safeParse(entry);
292+
if (!result.success) {
293+
return { valid: false, error: `Invalid custom claim at index ${i}: ${result.error.issues[0]?.message}` };
294+
}
295+
}
296+
}
297+
298+
// allowedAudience, allowedClients, allowedScopes, customClaims are all optional individually,
279299
// but at least one must be provided
280300
const hasAudience = !!options.allowedAudience?.trim();
281301
const hasClients = !!options.allowedClients?.trim();
282302
const hasScopes = !!options.allowedScopes?.trim();
283-
if (!hasAudience && !hasClients && !hasScopes) {
303+
const hasClaims = !!options.customClaims?.trim();
304+
if (!hasAudience && !hasClients && !hasScopes && !hasClaims) {
284305
return {
285306
valid: false,
286307
error:
287-
'At least one of --allowed-audience, --allowed-clients, or --allowed-scopes must be provided for CUSTOM_JWT authorizer',
308+
'At least one of --allowed-audience, --allowed-clients, --allowed-scopes, or --custom-claims must be provided for CUSTOM_JWT authorizer',
288309
};
289310
}
290311
}

‎src/cli/primitives/GatewayPrimitive.ts‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,11 @@
11
import { findConfigRoot, setEnvVar } from '../../lib';
2-
import type { AgentCoreGateway, AgentCoreGatewayTarget, AgentCoreMcpSpec, GatewayAuthorizerType } from '../../schema';
2+
import type {
3+
AgentCoreGateway,
4+
AgentCoreGatewayTarget,
5+
AgentCoreMcpSpec,
6+
CustomClaimValidation,
7+
GatewayAuthorizerType,
8+
} from '../../schema';
39
import { AgentCoreGatewaySchema, PolicyEngineModeSchema } from '../../schema';
410
import type { AddGatewayOptions as CLIAddGatewayOptions } from '../commands/add/types';
511
import { validateAddGatewayOptions } from '../commands/add/validate';
@@ -23,6 +29,7 @@ export interface AddGatewayOptions {
2329
allowedAudience?: string;
2430
allowedClients?: string;
2531
allowedScopes?: string;
32+
customClaims?: CustomClaimValidation[];
2633
clientId?: string;
2734
clientSecret?: string;
2835
agents?: string;
@@ -157,6 +164,7 @@ export class GatewayPrimitive extends BasePrimitive<AddGatewayOptions, Removable
157164
.option('--allowed-audience <audience>', 'Comma-separated allowed audiences (for CUSTOM_JWT)')
158165
.option('--allowed-clients <clients>', 'Comma-separated allowed client IDs (for CUSTOM_JWT)')
159166
.option('--allowed-scopes <scopes>', 'Comma-separated allowed scopes (for CUSTOM_JWT)')
167+
.option('--custom-claims <json>', 'Custom claim validations as JSON array (for CUSTOM_JWT)')
160168
.option('--client-id <id>', 'OAuth client ID for gateway bearer token')
161169
.option('--client-secret <secret>', 'OAuth client secret')
162170
.option('--agents <agents>', 'Comma-separated agent names')
@@ -183,6 +191,11 @@ export class GatewayPrimitive extends BasePrimitive<AddGatewayOptions, Removable
183191
process.exit(1);
184192
}
185193

194+
// Parse custom claims JSON if provided (already validated)
195+
const parsedCustomClaims = cliOptions.customClaims
196+
? (JSON.parse(cliOptions.customClaims) as CustomClaimValidation[])
197+
: undefined;
198+
186199
const result = await this.add({
187200
name: cliOptions.name!,
188201
description: cliOptions.description,
@@ -191,6 +204,7 @@ export class GatewayPrimitive extends BasePrimitive<AddGatewayOptions, Removable
191204
allowedAudience: cliOptions.allowedAudience,
192205
allowedClients: cliOptions.allowedClients,
193206
allowedScopes: cliOptions.allowedScopes,
207+
customClaims: parsedCustomClaims,
194208
clientId: cliOptions.clientId,
195209
clientSecret: cliOptions.clientSecret,
196210
agents: cliOptions.agents,
@@ -327,6 +341,7 @@ export class GatewayPrimitive extends BasePrimitive<AddGatewayOptions, Removable
327341
...(allowedAudience?.length ? { allowedAudience } : {}),
328342
...(allowedClients?.length ? { allowedClients } : {}),
329343
...(allowedScopes?.length ? { allowedScopes } : {}),
344+
...(options.customClaims?.length ? { customClaims: options.customClaims } : {}),
330345
...(options.clientId ? { clientId: options.clientId } : {}),
331346
...(options.clientSecret ? { clientSecret: options.clientSecret } : {}),
332347
};
@@ -429,6 +444,7 @@ export class GatewayPrimitive extends BasePrimitive<AddGatewayOptions, Removable
429444
...(config.jwtConfig.allowedAudience?.length ? { allowedAudience: config.jwtConfig.allowedAudience } : {}),
430445
...(config.jwtConfig.allowedClients?.length ? { allowedClients: config.jwtConfig.allowedClients } : {}),
431446
...(config.jwtConfig.allowedScopes?.length ? { allowedScopes: config.jwtConfig.allowedScopes } : {}),
447+
...(config.jwtConfig.customClaims?.length ? { customClaims: config.jwtConfig.customClaims } : {}),
432448
},
433449
};
434450
}

‎src/cli/primitives/__tests__/GatewayPrimitive.test.ts‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,78 @@ describe('GatewayPrimitive', () => {
3939
primitive = new GatewayPrimitive();
4040
});
4141

42+
describe('customClaims pipeline', () => {
43+
const SAMPLE_CLAIMS = [
44+
{
45+
inboundTokenClaimName: 'department',
46+
inboundTokenClaimValueType: 'STRING_ARRAY' as const,
47+
authorizingClaimMatchValue: {
48+
claimMatchOperator: 'CONTAINS_ANY' as const,
49+
claimMatchValue: { matchValueStringList: ['engineering', 'sales'] },
50+
},
51+
},
52+
];
53+
54+
it('custom claims from TUI flow are written to authorizerConfiguration', async () => {
55+
await primitive.add({
56+
name: 'jwt-gw',
57+
authorizerType: 'CUSTOM_JWT',
58+
discoveryUrl: 'https://example.com/.well-known/openid-configuration',
59+
allowedAudience: 'aud1',
60+
customClaims: SAMPLE_CLAIMS,
61+
});
62+
63+
const gw = getWrittenGateway();
64+
expect(gw.authorizerConfiguration?.customJwtAuthorizer).toBeDefined();
65+
expect(gw.authorizerConfiguration!.customJwtAuthorizer!.customClaims).toEqual(SAMPLE_CLAIMS);
66+
});
67+
68+
it('custom claims are preserved alongside audience and clients', async () => {
69+
await primitive.add({
70+
name: 'jwt-gw',
71+
authorizerType: 'CUSTOM_JWT',
72+
discoveryUrl: 'https://example.com/.well-known/openid-configuration',
73+
allowedAudience: 'aud1,aud2',
74+
allowedClients: 'client1',
75+
customClaims: SAMPLE_CLAIMS,
76+
});
77+
78+
const gw = getWrittenGateway();
79+
const jwtConfig = gw.authorizerConfiguration!.customJwtAuthorizer!;
80+
expect(jwtConfig.allowedAudience).toEqual(['aud1', 'aud2']);
81+
expect(jwtConfig.allowedClients).toEqual(['client1']);
82+
expect(jwtConfig.customClaims).toEqual(SAMPLE_CLAIMS);
83+
});
84+
85+
it('omits customClaims from authorizerConfiguration when not provided', async () => {
86+
await primitive.add({
87+
name: 'jwt-gw',
88+
authorizerType: 'CUSTOM_JWT',
89+
discoveryUrl: 'https://example.com/.well-known/openid-configuration',
90+
allowedAudience: 'aud1',
91+
});
92+
93+
const gw = getWrittenGateway();
94+
expect(gw.authorizerConfiguration!.customJwtAuthorizer!.customClaims).toBeUndefined();
95+
});
96+
97+
it('custom claims only (no audience/clients/scopes) produces valid config', async () => {
98+
await primitive.add({
99+
name: 'jwt-gw',
100+
authorizerType: 'CUSTOM_JWT',
101+
discoveryUrl: 'https://example.com/.well-known/openid-configuration',
102+
customClaims: SAMPLE_CLAIMS,
103+
});
104+
105+
const gw = getWrittenGateway();
106+
const jwtConfig = gw.authorizerConfiguration!.customJwtAuthorizer!;
107+
expect(jwtConfig.allowedAudience).toBeUndefined();
108+
expect(jwtConfig.allowedClients).toBeUndefined();
109+
expect(jwtConfig.allowedScopes).toBeUndefined();
110+
expect(jwtConfig.customClaims).toEqual(SAMPLE_CLAIMS);
111+
});
112+
});
113+
42114
describe('exceptionLevel', () => {
43115
it('defaults to exceptionLevel NONE', async () => {
44116
await primitive.add({ name: 'test-gw', authorizerType: 'NONE' });

‎src/cli/tui/hooks/useCreateMcp.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ export function useCreateGateway() {
3131
allowedAudience: config.jwtConfig?.allowedAudience?.join(','),
3232
allowedClients: config.jwtConfig?.allowedClients?.join(','),
3333
allowedScopes: config.jwtConfig?.allowedScopes?.join(','),
34+
customClaims: config.jwtConfig?.customClaims,
3435
clientId: config.jwtConfig?.clientId,
3536
clientSecret: config.jwtConfig?.clientSecret,
3637
enableSemanticSearch: config.enableSemanticSearch,

0 commit comments

Comments
 (0)