feat(oidc-jwt): support OIDC with JWT tokens

modules/oidc/lambda/callback/index.js

@@ -14,14 +14,23 @@ exports.handler = (event, context, callback) => {
   const query = event.queryStringParameters || {};
   const headers = event.headers || {};
   const cookies = event.cookies || [];
-  const providerKey = query.auth;
+  let providerKey = query.auth;
   const code = query.code;
   const state = query.state;
 
   //console.log('Callback Lambda - Query params:', { providerKey, code, state });
   //console.log('Callback Lambda - Headers:', JSON.stringify(headers, null, 2));
   //console.log('Callback Lambda - Cookies from event.cookies:', cookies);
 
+  // Try to read providerKey from cookie if not in query
+  if (!providerKey && cookies.length > 0) {
+    const authCookie = cookies.find(c => c.startsWith('auth_provider='));
+    if (authCookie) {
+      providerKey = authCookie.split('=')[1];
+      console.log('Callback Lambda - Retrieved providerKey from cookie:', providerKey);
+    }
+  }
+
   if (!providerKey || !config[providerKey]) {
     console.log('Callback Lambda - No matching provider for:', providerKey);
     return callback(null, {
@@ -105,12 +114,17 @@ exports.handler = (event, context, callback) => {
         // Create a session cookie
         const redirectUrl = new URL(provider.redirect_after_login);
         const cookieDomain = redirectUrl.hostname;
-        const sessionValue = json.access_token;
+        const now = Date.now();
+        const sessionPayload = {
+          access_token: json.access_token,
+          exp: now + (provider.session_duration * 1000),
+        };
+        const payload = Buffer.from(JSON.stringify(sessionPayload)).toString('base64');
         const signature = crypto
           .createHmac('sha256', provider.session_secret)
-          .update(sessionValue)
+          .update(payload)
           .digest('hex');
-        const sessionCookie = `session=${sessionValue}.${signature}; Domain=${cookieDomain}; Path=/; HttpOnly; Secure; SameSite=None; Max-Age=${provider.session_duration}`;
+        const sessionCookie = `session=${payload}.${signature}; Domain=${cookieDomain}; Path=/; HttpOnly; Secure; SameSite=None; Max-Age=${provider.session_duration}`;
 
         //console.log('Callback Lambda - Setting session cookie:', sessionCookie);
         return callback(null, {

modules/oidc/lambda/edge_auth/index.js

@@ -25,94 +25,107 @@ exports.handler = (event, context, callback) => {
     const headers = request.headers || {};
     const query = request.querystring || '';
     const params = querystring.parse(query);
-    const providerKey = params.auth;
+    const authParam = params.auth; // just for logging/debug
 
-    console.log('Edge Lambda - Query params:', { providerKey });
+    let providerKey = null;
+    console.log('Edge Lambda - Query params:', { authParam });
 
-    // Extract session cookie
+    // Extract session cookie & providerKey from cookies
     let session = null;
     if (headers.cookie && Array.isArray(headers.cookie)) {
-      //console.log('Edge Lambda - Cookies:', headers.cookie);
       for (const cookie of headers.cookie) {
         const cookieValue = cookie.value || '';
-        //console.log('Edge Lambda - Processing cookie string:', cookieValue);
-        // Split cookies delimited by semicolon or space
         const cookieEntries = cookieValue.split('; ').map(entry => entry.trim());
         for (const entry of cookieEntries) {
           if (entry.startsWith('session=')) {
             session = entry.split('=')[1];
-            //console.log('Edge Lambda - Extracted session cookie:', session);
-            break;
+          } else if (entry.startsWith('auth_provider=')) {
+            providerKey = entry.split('=')[1];
           }
         }
-        if (session) break;
+        if (session && providerKey) break;
       }
     } else {
       console.log('Edge Lambda - No cookie header present or not an array');
     }
 
+    // Determine provider for this request
+    // If providerKey not found in cookie, use query param ?auth=... or fallback to default
+    let newlySelectedProviderKey = null;
+    if (!providerKey) {
+      newlySelectedProviderKey = params.auth || Object.keys(config)[0];
+      providerKey = newlySelectedProviderKey;
+      console.log('Edge Lambda - No auth_provider cookie, using query param or default:', providerKey);
+    } else {
+      console.log('Edge Lambda - Using provider from cookie:', providerKey);
+    }
+
     // Validate session cookie if present
     if (session) {
-      //console.log('Edge Lambda - Found session cookie:', session);
-      // Check session cookie format (value.signature)
-      if (!session.includes('.')) {
-        console.log('Edge Lambda - Invalid session cookie format:', session);
+      if (!providerKey || !config[providerKey]) {
+        console.log('Edge Lambda - Invalid or missing providerKey in config:', providerKey);
       } else {
-        const [value, signature] = session.split('.');
-        //console.log('Edge Lambda - Session cookie parts - Value:', value, 'Signature:', signature);
+        const provider = config[providerKey];
 
-        // Validate for selected provider if providerKey is present
-        if (providerKey && config[providerKey]) {
-          const provider = config[providerKey];
-          //console.log('Edge Lambda - Provider config:', JSON.stringify(provider, null, 2));
+        // JWT-like format: value.signature
+        if (session.includes('.')) {
+          const [value, signature] = session.split('.');
           const expectedSignature = crypto
             .createHmac('sha256', provider.session_secret || '')
             .update(value)
             .digest('hex');
           console.log('Edge Lambda - Validating session for provider:', providerKey);
-          //console.log('Edge Lambda - Expected signature:', expectedSignature, 'Got:', signature);
           if (signature === expectedSignature) {
-            console.log('Edge Lambda - Session validated successfully for provider:', providerKey);
-            return callback(null, request);
+            try {
+              const sessionJson = Buffer.from(value, 'base64').toString('utf-8');
+              const sessionPayload = JSON.parse(sessionJson);
+              const now = Date.now();
+              if (!sessionPayload.exp || now < sessionPayload.exp) {
+                console.log('Edge Lambda - Session validated successfully for provider:', providerKey);
+                return callback(null, request);
+              } else {
+                console.log('Edge Lambda - Session expired');
+              }
+            } catch (err) {
+              console.log('Edge Lambda - Failed to parse session payload:', err.message);
+              console.log('Edge Lambda - Accepting valid signature as opaque token');
+              return callback(null, request);
+            }
           } else {
             console.log('Edge Lambda - Session validation failed for provider:', providerKey);
           }
         } else {
-          // Loop through providers if providerKey is not present
-          console.log('Edge Lambda - No providerKey, trying all providers');
-          for (const key of Object.keys(config)) {
-            const provider = config[key];
-            //console.log('Edge Lambda - Provider config for', key, ':', JSON.stringify(provider, null, 2));
-            const expectedSignature = crypto
-              .createHmac('sha256', provider.session_secret || '')
-              .update(value)
-              .digest('hex');
-            //console.log(`Edge Lambda - Validating session for provider: ${key}, Expected signature: ${expectedSignature}, Got: ${signature}`);
-            if (signature === expectedSignature) {
-              console.log('Edge Lambda - Session validated successfully for provider:', key);
+          // Opaque token: base64-encoded JSON
+          try {
+            const sessionJson = Buffer.from(session, 'base64').toString('utf-8');
+            const sessionPayload = JSON.parse(sessionJson);
+            const now = Date.now();
+            if (!sessionPayload.exp || now < sessionPayload.exp) {
+              console.log('Edge Lambda - Opaque session validated successfully for provider:', providerKey);
               return callback(null, request);
+            } else {
+              console.log('Edge Lambda - Opaque session expired');
             }
+          } catch (err) {
+            console.log('Edge Lambda - Failed to parse opaque session payload:', err.message);
           }
-          console.log('Edge Lambda - Session validation failed for all providers');
         }
       }
     } else {
       console.log('Edge Lambda - No session cookie found');
     }
 
     // Redirect to OIDC if session is invalid or not present
-    //console.log('Edge Lambda - Config:', JSON.stringify(config, null, 2));
-    const defaultProviderKey = providerKey || Object.keys(config)[0];
-    if (!defaultProviderKey || !config[defaultProviderKey]) {
-      console.log('Edge Lambda - No matching provider for:', defaultProviderKey);
+    if (!providerKey || !config[providerKey]) {
+      console.log('Edge Lambda - No matching provider for:', providerKey);
       return callback(null, {
         status: '403',
         statusDescription: 'Forbidden',
         body: 'No matching OIDC provider.',
       });
     }
 
-    const provider = config[defaultProviderKey];
+    const provider = config[providerKey];
     const state = crypto.randomBytes(16).toString('hex');
     const loginUrl = `${provider.auth_url}?` +
       `client_id=${encodeURIComponent(provider.client_id)}` +
@@ -122,6 +135,20 @@ exports.handler = (event, context, callback) => {
     console.log('Edge Lambda - Redirecting to OIDC provider:', loginUrl);
     console.log('Edge Lambda - Setting state cookie:', state);
 
+    // Prepare set-cookie headers
+    const setCookieHeaders = [{
+      key: 'Set-Cookie',
+      value: `state=${state}; Path=/; HttpOnly; Secure; SameSite=None; Max-Age=300`,
+    }];
+
+    if (newlySelectedProviderKey) {
+      const domain = new URL(provider.redirect_after_login).hostname;
+      setCookieHeaders.push({
+        key: 'Set-Cookie',
+        value: `auth_provider=${newlySelectedProviderKey}; Path=/; HttpOnly; Secure; SameSite=None; Max-Age=300`,
+      });
+    }
+
     return callback(null, {
       status: '302',
       statusDescription: 'Found',
@@ -130,10 +157,7 @@ exports.handler = (event, context, callback) => {
           key: 'Location',
           value: loginUrl,
         }],
-        'set-cookie': [{
-          key: 'Set-Cookie',
-          value: `state=${state}; Path=/; HttpOnly; Secure; SameSite=None; Max-Age=300`,
-        }],
+        'set-cookie': setCookieHeaders,
       },
     });
   } catch (error) {
@@ -144,4 +168,4 @@ exports.handler = (event, context, callback) => {
       body: 'Edge Lambda failed to process the request',
     });
   }
-};
+};

modules/oidc/shared.tf

@@ -8,7 +8,7 @@ locals {
       client_secret        = cfg.client_secret
       auth_url             = cfg.auth_url
       token_url            = cfg.token_url
-      redirect_uri         = "https://${var.application_domain}/callback?auth=${cfg.application_name}"
+      redirect_uri         = "https://${var.application_domain}/callback"
       session_secret       = random_string.session_secret[0].result
       redirect_after_login = "https://${var.application_domain}"
       session_duration     = cfg.session_duration