Skip to content

Commit a9527a6

Browse files
b4s36t4b4s36t4
committed
Merge branch 'main' of github.com:Portkey-AI/gateway into feat/ft-batch-improvements
2 parents 3f74159 + cfb5238 commit a9527a6

32 files changed

+908
-82
lines changed

‎conf.json‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,8 @@
77
"pillar",
88
"patronus",
99
"pangea",
10-
"promptsecurity"
10+
"promptsecurity",
11+
"panw-prisma-airs"
1112
],
1213
"credentials": {
1314
"portkey": {

‎package-lock.json‎

Lines changed: 11 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎package.json‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
{
22
"name": "@portkey-ai/gateway",
3-
"version": "1.9.18",
3+
"version": "1.9.19",
44
"description": "A fast AI gateway by Portkey",
55
"repository": {
66
"type": "git",
@@ -51,6 +51,7 @@
5151
"async-retry": "^1.3.3",
5252
"avsc": "^5.7.7",
5353
"hono": "^4.6.10",
54+
"jose": "^6.0.11",
5455
"patch-package": "^8.0.0",
5556
"ws": "^8.18.0",
5657
"zod": "^3.22.4"

‎plugins/default/default.test.ts‎

Lines changed: 154 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,7 @@ import { handler as endsWithHandler } from './endsWith';
1313
import { handler as allLowerCaseHandler } from './alllowercase';
1414
import { handler as modelWhitelistHandler } from './modelWhitelist';
1515
import { handler as characterCountHandler } from './characterCount';
16-
17-
import { z } from 'zod';
16+
import { handler as jwtHandler } from './jwt';
1817
import { PluginContext, PluginParameters } from '../types';
1918

2019
describe('Regex Matcher Plugin', () => {
@@ -2315,3 +2314,156 @@ describe('endsWith handler', () => {
23152314
});
23162315
});
23172316
});
2317+
2318+
describe('jwt handler', () => {
2319+
const mockEventType = 'beforeRequestHook';
2320+
2321+
it('should validate a valid JWT token', async () => {
2322+
const context: PluginContext = {
2323+
headers: {
2324+
Authorization:
2325+
'Bearer eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c',
2326+
},
2327+
};
2328+
const parameters: PluginParameters = {
2329+
jwksUri: 'https://www.googleapis.com/oauth2/v3/certs',
2330+
headerKey: 'Authorization',
2331+
};
2332+
2333+
const result = await jwtHandler(context, parameters, mockEventType);
2334+
2335+
expect(result.error).toBe(null);
2336+
expect(result.verdict).toBe(true);
2337+
expect(result.data).toMatchObject({
2338+
verdict: true,
2339+
explanation: 'JWT token validation succeeded',
2340+
payload: expect.any(Object),
2341+
});
2342+
});
2343+
2344+
it('should handle missing authorization header', async () => {
2345+
const context: PluginContext = {
2346+
headers: {},
2347+
};
2348+
const parameters: PluginParameters = {
2349+
jwksUri: 'https://www.googleapis.com/oauth2/v3/certs',
2350+
};
2351+
2352+
const result = await jwtHandler(context, parameters, mockEventType);
2353+
2354+
expect(result.error).not.toBe(null);
2355+
expect(result.verdict).toBe(false);
2356+
expect(result.data).toMatchObject({
2357+
verdict: false,
2358+
explanation: 'JWT validation error: Missing authorization header',
2359+
});
2360+
});
2361+
2362+
it('should handle invalid authorization header format', async () => {
2363+
const context: PluginContext = {
2364+
headers: {
2365+
Authorization: 'InvalidFormat',
2366+
},
2367+
};
2368+
const parameters: PluginParameters = {
2369+
jwksUri: 'https://www.googleapis.com/oauth2/v3/certs',
2370+
};
2371+
2372+
const result = await jwtHandler(context, parameters, mockEventType);
2373+
2374+
expect(result.error).not.toBe(null);
2375+
expect(result.verdict).toBe(false);
2376+
expect(result.data).toMatchObject({
2377+
verdict: false,
2378+
explanation: 'JWT validation error: Invalid authorization header format',
2379+
});
2380+
});
2381+
2382+
it('should handle missing JWKS URI', async () => {
2383+
const context: PluginContext = {
2384+
headers: {
2385+
Authorization: 'Bearer valid.token.here',
2386+
},
2387+
};
2388+
const parameters: PluginParameters = {};
2389+
2390+
const result = await jwtHandler(context, parameters, mockEventType);
2391+
2392+
expect(result.error).not.toBe(null);
2393+
expect(result.verdict).toBe(false);
2394+
expect(result.data).toMatchObject({
2395+
verdict: false,
2396+
explanation: 'JWT validation error: Missing JWKS URI',
2397+
});
2398+
});
2399+
2400+
it('should use custom header key from parameters', async () => {
2401+
const context: PluginContext = {
2402+
headers: {
2403+
'X-Custom-Auth':
2404+
'Bearer eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c',
2405+
},
2406+
};
2407+
const parameters: PluginParameters = {
2408+
jwksUri: 'https://www.googleapis.com/oauth2/v3/certs',
2409+
headerKey: 'X-Custom-Auth',
2410+
};
2411+
2412+
const result = await jwtHandler(context, parameters, mockEventType);
2413+
2414+
expect(result.error).toBe(null);
2415+
expect(result.verdict).toBe(true);
2416+
expect(result.data).toMatchObject({
2417+
verdict: true,
2418+
explanation: 'JWT token validation succeeded',
2419+
payload: expect.any(Object),
2420+
});
2421+
});
2422+
2423+
it('should use environment variables when parameters not provided', async () => {
2424+
process.env.JWT_AUTH_HEADER = 'X-Env-Auth';
2425+
process.env.JWT_JWKS_URI = 'https://www.googleapis.com/oauth2/v3/certs';
2426+
2427+
const context: PluginContext = {
2428+
headers: {
2429+
'X-Env-Auth':
2430+
'Bearer eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c',
2431+
},
2432+
};
2433+
const parameters: PluginParameters = {};
2434+
2435+
const result = await jwtHandler(context, parameters, mockEventType);
2436+
2437+
expect(result.error).toBe(null);
2438+
expect(result.verdict).toBe(true);
2439+
expect(result.data).toMatchObject({
2440+
verdict: true,
2441+
explanation: 'JWT token validation succeeded',
2442+
payload: expect.any(Object),
2443+
});
2444+
2445+
// Clean up
2446+
delete process.env.JWT_AUTH_HEADER;
2447+
delete process.env.JWT_JWKS_URI;
2448+
});
2449+
2450+
it('should handle invalid JWKS URI', async () => {
2451+
const context: PluginContext = {
2452+
headers: {
2453+
Authorization: 'Bearer valid.token.here',
2454+
},
2455+
};
2456+
const parameters: PluginParameters = {
2457+
jwksUri: 'https://invalid-jwks-uri.example.com/keys',
2458+
};
2459+
2460+
const result = await jwtHandler(context, parameters, mockEventType);
2461+
2462+
expect(result.error).not.toBe(null);
2463+
expect(result.verdict).toBe(false);
2464+
expect(result.data).toMatchObject({
2465+
verdict: false,
2466+
explanation: expect.stringContaining('JWT validation error'),
2467+
});
2468+
});
2469+
});

‎plugins/default/jwt.ts‎

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
import { PluginContext, PluginHandler, PluginParameters } from '../types';
2+
import { jwtVerify, importJWK, JWTHeaderParameters, JWK } from 'jose';
3+
4+
interface JWKSCacheOptions {
5+
maxAge: number; // seconds
6+
clockTolerance?: number; // seconds
7+
maxTokenAge?: string; // 1d, 1h, 1m, 1s
8+
}
9+
10+
async function getMatchingKey(token: string, jwks: any) {
11+
const header = JSON.parse(
12+
Buffer.from(token.split('.')[0], 'base64url').toString()
13+
) as JWTHeaderParameters;
14+
if (!header.kid) return null;
15+
const jwk = (jwks.keys || []).find((key: JWK) => key.kid === header.kid);
16+
if (!jwk) return null;
17+
return importJWK(jwk, jwk.alg);
18+
}
19+
20+
async function fetchAndCacheJWKS(
21+
jwksUri: string,
22+
cacheKey: string,
23+
maxAge: number,
24+
putInCacheWithValue?: Function
25+
) {
26+
const res = await fetch(jwksUri);
27+
if (!res.ok) throw new Error(`Failed to fetch JWKS from ${jwksUri}`);
28+
const jwks = await res.json();
29+
await putInCacheWithValue?.(cacheKey, jwks, maxAge);
30+
return jwks;
31+
}
32+
33+
export const handler: PluginHandler = async (
34+
context: PluginContext,
35+
parameters: PluginParameters,
36+
eventType: string,
37+
options
38+
) => {
39+
let error = null;
40+
let verdict = false;
41+
let data = null;
42+
43+
try {
44+
const jwtAuthHeader = parameters.headerKey || 'Authorization';
45+
const jwksUri = parameters.jwksUri;
46+
if (!jwksUri) throw new Error('Missing JWKS URI');
47+
48+
const authHeader = context.request?.headers?.[jwtAuthHeader];
49+
if (!authHeader) {
50+
return {
51+
error,
52+
verdict,
53+
data: {
54+
verdict: false,
55+
explanation: 'Missing authorization header',
56+
},
57+
};
58+
}
59+
60+
const token = authHeader.replace(/^Bearer\s+/i, '').trim();
61+
62+
const cacheKey = `jwks:${jwksUri}`;
63+
const cacheOptions: JWKSCacheOptions = {
64+
maxAge: parameters.cacheMaxAge || 86400, // 24h
65+
clockTolerance: parameters.clockTolerance || 5, // 5s
66+
maxTokenAge: parameters.maxTokenAge || '1d',
67+
};
68+
69+
let jwks = await options?.getFromCacheByKey?.(cacheKey);
70+
if (!jwks) {
71+
jwks = await fetchAndCacheJWKS(
72+
jwksUri,
73+
cacheKey,
74+
cacheOptions.maxAge,
75+
options?.putInCacheWithValue
76+
);
77+
}
78+
79+
let key = null;
80+
try {
81+
key = await getMatchingKey(token, jwks);
82+
if (!key) {
83+
jwks = await fetchAndCacheJWKS(
84+
jwksUri,
85+
cacheKey,
86+
cacheOptions.maxAge,
87+
options?.putInCacheWithValue
88+
);
89+
key = await getMatchingKey(token, jwks);
90+
}
91+
92+
if (!key) {
93+
return {
94+
error,
95+
verdict,
96+
data: {
97+
verdict: false,
98+
explanation: 'No matching key found for kid',
99+
},
100+
};
101+
}
102+
} catch (e: any) {
103+
return {
104+
error: e,
105+
verdict,
106+
data: {
107+
verdict: false,
108+
explanation: `JWT validation error: ${e.message}`,
109+
},
110+
};
111+
}
112+
113+
try {
114+
await jwtVerify(token, key, {
115+
clockTolerance: cacheOptions.clockTolerance,
116+
maxTokenAge: cacheOptions.maxTokenAge,
117+
});
118+
verdict = true;
119+
data = {
120+
verdict,
121+
explanation: 'JWT token validation succeeded',
122+
};
123+
} catch (e: any) {
124+
data = {
125+
verdict: false,
126+
explanation: `JWT validation error: ${e.message}`,
127+
};
128+
}
129+
} catch (e: any) {
130+
error = e;
131+
data = {
132+
verdict: false,
133+
explanation: `JWT validation error: ${e.message}`,
134+
};
135+
}
136+
return { error, verdict, data };
137+
};

0 commit comments

Comments
 (0)