Skip to content

Commit 89da5e2

Browse files
jindrajjindraj
committed
feat(oidc-jwt): support OIDC with JWT tokens
- also remove ?auth=PROVIDER from callback, save the provider into cookie instead
1 parent 44253ad commit 89da5e2

File tree

3 files changed

+88
-50
lines changed

3 files changed

+88
-50
lines changed

‎modules/oidc/lambda/callback/index.js‎

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,14 +14,23 @@ exports.handler = (event, context, callback) => {
1414
const query = event.queryStringParameters || {};
1515
const headers = event.headers || {};
1616
const cookies = event.cookies || [];
17-
const providerKey = query.auth;
17+
let providerKey = query.auth;
1818
const code = query.code;
1919
const state = query.state;
2020

2121
//console.log('Callback Lambda - Query params:', { providerKey, code, state });
2222
//console.log('Callback Lambda - Headers:', JSON.stringify(headers, null, 2));
2323
//console.log('Callback Lambda - Cookies from event.cookies:', cookies);
2424

25+
// Try to read providerKey from cookie if not in query
26+
if (!providerKey && cookies.length > 0) {
27+
const authCookie = cookies.find(c => c.startsWith('auth_provider='));
28+
if (authCookie) {
29+
providerKey = authCookie.split('=')[1];
30+
console.log('Callback Lambda - Retrieved providerKey from cookie:', providerKey);
31+
}
32+
}
33+
2534
if (!providerKey || !config[providerKey]) {
2635
console.log('Callback Lambda - No matching provider for:', providerKey);
2736
return callback(null, {
@@ -105,12 +114,17 @@ exports.handler = (event, context, callback) => {
105114
// Create a session cookie
106115
const redirectUrl = new URL(provider.redirect_after_login);
107116
const cookieDomain = redirectUrl.hostname;
108-
const sessionValue = json.access_token;
117+
const now = Date.now();
118+
const sessionPayload = {
119+
access_token: json.access_token,
120+
exp: now + (provider.session_duration * 1000),
121+
};
122+
const payload = Buffer.from(JSON.stringify(sessionPayload)).toString('base64');
109123
const signature = crypto
110124
.createHmac('sha256', provider.session_secret)
111-
.update(sessionValue)
125+
.update(payload)
112126
.digest('hex');
113-
const sessionCookie = `session=${sessionValue}.${signature}; Domain=${cookieDomain}; Path=/; HttpOnly; Secure; SameSite=None; Max-Age=${provider.session_duration}`;
127+
const sessionCookie = `session=${payload}.${signature}; Domain=${cookieDomain}; Path=/; HttpOnly; Secure; SameSite=None; Max-Age=${provider.session_duration}`;
114128

115129
//console.log('Callback Lambda - Setting session cookie:', sessionCookie);
116130
return callback(null, {

‎modules/oidc/lambda/edge_auth/index.js‎

Lines changed: 69 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -25,94 +25,107 @@ exports.handler = (event, context, callback) => {
2525
const headers = request.headers || {};
2626
const query = request.querystring || '';
2727
const params = querystring.parse(query);
28-
const providerKey = params.auth;
28+
const authParam = params.auth; // just for logging/debug
2929

30-
console.log('Edge Lambda - Query params:', { providerKey });
30+
let providerKey = null;
31+
console.log('Edge Lambda - Query params:', { authParam });
3132

32-
// Extract session cookie
33+
// Extract session cookie & providerKey from cookies
3334
let session = null;
3435
if (headers.cookie && Array.isArray(headers.cookie)) {
35-
//console.log('Edge Lambda - Cookies:', headers.cookie);
3636
for (const cookie of headers.cookie) {
3737
const cookieValue = cookie.value || '';
38-
//console.log('Edge Lambda - Processing cookie string:', cookieValue);
39-
// Split cookies delimited by semicolon or space
4038
const cookieEntries = cookieValue.split('; ').map(entry => entry.trim());
4139
for (const entry of cookieEntries) {
4240
if (entry.startsWith('session=')) {
4341
session = entry.split('=')[1];
44-
//console.log('Edge Lambda - Extracted session cookie:', session);
45-
break;
42+
} else if (entry.startsWith('auth_provider=')) {
43+
providerKey = entry.split('=')[1];
4644
}
4745
}
48-
if (session) break;
46+
if (session && providerKey) break;
4947
}
5048
} else {
5149
console.log('Edge Lambda - No cookie header present or not an array');
5250
}
5351

52+
// Determine provider for this request
53+
// If providerKey not found in cookie, use query param ?auth=... or fallback to default
54+
let newlySelectedProviderKey = null;
55+
if (!providerKey) {
56+
newlySelectedProviderKey = params.auth || Object.keys(config)[0];
57+
providerKey = newlySelectedProviderKey;
58+
console.log('Edge Lambda - No auth_provider cookie, using query param or default:', providerKey);
59+
} else {
60+
console.log('Edge Lambda - Using provider from cookie:', providerKey);
61+
}
62+
5463
// Validate session cookie if present
5564
if (session) {
56-
//console.log('Edge Lambda - Found session cookie:', session);
57-
// Check session cookie format (value.signature)
58-
if (!session.includes('.')) {
59-
console.log('Edge Lambda - Invalid session cookie format:', session);
65+
if (!providerKey || !config[providerKey]) {
66+
console.log('Edge Lambda - Invalid or missing providerKey in config:', providerKey);
6067
} else {
61-
const [value, signature] = session.split('.');
62-
//console.log('Edge Lambda - Session cookie parts - Value:', value, 'Signature:', signature);
68+
const provider = config[providerKey];
6369

64-
// Validate for selected provider if providerKey is present
65-
if (providerKey && config[providerKey]) {
66-
const provider = config[providerKey];
67-
//console.log('Edge Lambda - Provider config:', JSON.stringify(provider, null, 2));
70+
// JWT-like format: value.signature
71+
if (session.includes('.')) {
72+
const [value, signature] = session.split('.');
6873
const expectedSignature = crypto
6974
.createHmac('sha256', provider.session_secret || '')
7075
.update(value)
7176
.digest('hex');
7277
console.log('Edge Lambda - Validating session for provider:', providerKey);
73-
//console.log('Edge Lambda - Expected signature:', expectedSignature, 'Got:', signature);
7478
if (signature === expectedSignature) {
75-
console.log('Edge Lambda - Session validated successfully for provider:', providerKey);
76-
return callback(null, request);
79+
try {
80+
const sessionJson = Buffer.from(value, 'base64').toString('utf-8');
81+
const sessionPayload = JSON.parse(sessionJson);
82+
const now = Date.now();
83+
if (!sessionPayload.exp || now < sessionPayload.exp) {
84+
console.log('Edge Lambda - Session validated successfully for provider:', providerKey);
85+
return callback(null, request);
86+
} else {
87+
console.log('Edge Lambda - Session expired');
88+
}
89+
} catch (err) {
90+
console.log('Edge Lambda - Failed to parse session payload:', err.message);
91+
console.log('Edge Lambda - Accepting valid signature as opaque token');
92+
return callback(null, request);
93+
}
7794
} else {
7895
console.log('Edge Lambda - Session validation failed for provider:', providerKey);
7996
}
8097
} else {
81-
// Loop through providers if providerKey is not present
82-
console.log('Edge Lambda - No providerKey, trying all providers');
83-
for (const key of Object.keys(config)) {
84-
const provider = config[key];
85-
//console.log('Edge Lambda - Provider config for', key, ':', JSON.stringify(provider, null, 2));
86-
const expectedSignature = crypto
87-
.createHmac('sha256', provider.session_secret || '')
88-
.update(value)
89-
.digest('hex');
90-
//console.log(`Edge Lambda - Validating session for provider: ${key}, Expected signature: ${expectedSignature}, Got: ${signature}`);
91-
if (signature === expectedSignature) {
92-
console.log('Edge Lambda - Session validated successfully for provider:', key);
98+
// Opaque token: base64-encoded JSON
99+
try {
100+
const sessionJson = Buffer.from(session, 'base64').toString('utf-8');
101+
const sessionPayload = JSON.parse(sessionJson);
102+
const now = Date.now();
103+
if (!sessionPayload.exp || now < sessionPayload.exp) {
104+
console.log('Edge Lambda - Opaque session validated successfully for provider:', providerKey);
93105
return callback(null, request);
106+
} else {
107+
console.log('Edge Lambda - Opaque session expired');
94108
}
109+
} catch (err) {
110+
console.log('Edge Lambda - Failed to parse opaque session payload:', err.message);
95111
}
96-
console.log('Edge Lambda - Session validation failed for all providers');
97112
}
98113
}
99114
} else {
100115
console.log('Edge Lambda - No session cookie found');
101116
}
102117

103118
// Redirect to OIDC if session is invalid or not present
104-
//console.log('Edge Lambda - Config:', JSON.stringify(config, null, 2));
105-
const defaultProviderKey = providerKey || Object.keys(config)[0];
106-
if (!defaultProviderKey || !config[defaultProviderKey]) {
107-
console.log('Edge Lambda - No matching provider for:', defaultProviderKey);
119+
if (!providerKey || !config[providerKey]) {
120+
console.log('Edge Lambda - No matching provider for:', providerKey);
108121
return callback(null, {
109122
status: '403',
110123
statusDescription: 'Forbidden',
111124
body: 'No matching OIDC provider.',
112125
});
113126
}
114127

115-
const provider = config[defaultProviderKey];
128+
const provider = config[providerKey];
116129
const state = crypto.randomBytes(16).toString('hex');
117130
const loginUrl = `${provider.auth_url}?` +
118131
`client_id=${encodeURIComponent(provider.client_id)}` +
@@ -122,6 +135,20 @@ exports.handler = (event, context, callback) => {
122135
console.log('Edge Lambda - Redirecting to OIDC provider:', loginUrl);
123136
console.log('Edge Lambda - Setting state cookie:', state);
124137

138+
// Prepare set-cookie headers
139+
const setCookieHeaders = [{
140+
key: 'Set-Cookie',
141+
value: `state=${state}; Path=/; HttpOnly; Secure; SameSite=None; Max-Age=300`,
142+
}];
143+
144+
if (newlySelectedProviderKey) {
145+
const domain = new URL(provider.redirect_after_login).hostname;
146+
setCookieHeaders.push({
147+
key: 'Set-Cookie',
148+
value: `auth_provider=${newlySelectedProviderKey}; Path=/; HttpOnly; Secure; SameSite=None; Max-Age=300`,
149+
});
150+
}
151+
125152
return callback(null, {
126153
status: '302',
127154
statusDescription: 'Found',
@@ -130,10 +157,7 @@ exports.handler = (event, context, callback) => {
130157
key: 'Location',
131158
value: loginUrl,
132159
}],
133-
'set-cookie': [{
134-
key: 'Set-Cookie',
135-
value: `state=${state}; Path=/; HttpOnly; Secure; SameSite=None; Max-Age=300`,
136-
}],
160+
'set-cookie': setCookieHeaders,
137161
},
138162
});
139163
} catch (error) {
@@ -144,4 +168,4 @@ exports.handler = (event, context, callback) => {
144168
body: 'Edge Lambda failed to process the request',
145169
});
146170
}
147-
};
171+
};

‎modules/oidc/shared.tf‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ locals {
88
client_secret = cfg.client_secret
99
auth_url = cfg.auth_url
1010
token_url = cfg.token_url
11-
redirect_uri = "https://${var.application_domain}/callback?auth=${cfg.application_name}"
11+
redirect_uri = "https://${var.application_domain}/callback"
1212
session_secret = random_string.session_secret[0].result
1313
redirect_after_login = "https://${var.application_domain}"
1414
session_duration = cfg.session_duration

0 commit comments

Comments
 (0)