fix(core): provision user who has valid sign-in session to organizations when using one-time token (#7220)

packages/core/src/libraries/one-time-token.ts

@@ -14,7 +14,7 @@ export const createOneTimeTokenLibrary = (queries: Queries) => {
     return updateOneTimeTokenStatusQuery(token, status);
   };
 
-  const verifyOneTimeToken = async (token: string, email: string) => {
+  const checkOneTimeToken = async (token: string, email: string) => {
     const oneTimeToken = await getOneTimeTokenByToken(token);
 
     assertThat(
@@ -39,8 +39,14 @@ export const createOneTimeTokenLibrary = (queries: Queries) => {
     );
     assertThat(oneTimeToken.status !== OneTimeTokenStatus.Revoked, 'one_time_token.token_revoked');
 
+    return oneTimeToken;
+  };
+
+  const verifyOneTimeToken = async (token: string, email: string) => {
+    await checkOneTimeToken(token, email);
+
     return updateOneTimeTokenStatus(token, OneTimeTokenStatus.Consumed);
   };
 
-  return { updateOneTimeTokenStatus, verifyOneTimeToken };
+  return { checkOneTimeToken, updateOneTimeTokenStatus, verifyOneTimeToken };
 };

packages/core/src/middleware/koa-consent-guard.test.ts

@@ -1,7 +1,8 @@
+import { OneTimeTokenStatus } from '@logto/schemas';
 import { Provider } from 'oidc-provider';
 
 import RequestError from '#src/errors/RequestError/index.js';
-import { MockQueries } from '#src/test-utils/tenant.js';
+import { MockTenant } from '#src/test-utils/tenant.js';
 import { createContextWithRouteParameters } from '#src/utils/test-utils.js';
 
 import koaConsentGuard from './koa-consent-guard.js';
@@ -12,11 +13,36 @@ describe('koaConsentGuard middleware', () => {
   const provider = new Provider('https://logto.test');
   const interactionDetails = jest.spyOn(provider, 'interactionDetails');
 
-  const mockQueries = new MockQueries({
-    users: {
-      findUserById: jest.fn().mockResolvedValue({ primaryEmail: 'foo@example.com' }),
+  const checkOneTimeToken = jest.fn().mockResolvedValue({
+    token: 'token_value',
+    email: 'foo@example.com',
+    status: OneTimeTokenStatus.Active,
+    context: {
+      jitOrganizationIds: ['org_id'],
     },
   });
+  const updateOneTimeTokenStatus = jest.fn().mockResolvedValue({
+    token: 'token_value',
+    status: OneTimeTokenStatus.Consumed,
+    context: {
+      jitOrganizationIds: ['org_id'],
+    },
+  });
+  const provisionOrganizations = jest.fn();
+
+  const mockTenant = new MockTenant(
+    provider,
+    {
+      users: {
+        findUserById: jest.fn().mockResolvedValue({ primaryEmail: 'foo@example.com' }),
+      },
+    },
+    undefined,
+    {
+      oneTimeTokens: { checkOneTimeToken, updateOneTimeTokenStatus },
+      users: { provisionOrganizations },
+    }
+  );
 
   const next = jest.fn();
 
@@ -33,7 +59,7 @@ describe('koaConsentGuard middleware', () => {
     const ctx = createContextWithRouteParameters({
       url: `/consent`,
     });
-    const guard = koaConsentGuard(provider, mockQueries);
+    const guard = koaConsentGuard(provider, mockTenant.libraries, mockTenant.queries);
 
     await expect(guard(ctx, next)).rejects.toThrow(new RequestError({ code: 'session.not_found' }));
   });
@@ -47,10 +73,10 @@ describe('koaConsentGuard middleware', () => {
     const ctx = createContextWithRouteParameters({
       url: `/consent`,
     });
-    const guard = koaConsentGuard(provider, mockQueries);
+    const guard = koaConsentGuard(provider, mockTenant.libraries, mockTenant.queries);
 
     await guard(ctx, next);
-    expect(mockQueries.users.findUserById).not.toHaveBeenCalled();
+    expect(mockTenant.queries.users.findUserById).not.toHaveBeenCalled();
     expect(next).toHaveBeenCalled();
   });
 
@@ -63,7 +89,7 @@ describe('koaConsentGuard middleware', () => {
     const ctx = createContextWithRouteParameters({
       url: `/consent`,
     });
-    const guard = koaConsentGuard(provider, mockQueries);
+    const guard = koaConsentGuard(provider, mockTenant.libraries, mockTenant.queries);
 
     await guard(ctx, jest.fn());
     expect(ctx.redirect).toHaveBeenCalledWith(
@@ -71,6 +97,38 @@ describe('koaConsentGuard middleware', () => {
     );
   });
 
+  it('should provision user to organizations on consent, if a valid one-time token is provided and there are organizations in token context', async () => {
+    interactionDetails.mockResolvedValue({
+      params: { one_time_token: 'token_value', login_hint: 'foo@example.com' },
+      // @ts-expect-error
+      session: { accountId: 'foo' },
+    });
+    checkOneTimeToken.mockResolvedValue({
+      token: 'token_value',
+      email: 'foo@example.com',
+      status: OneTimeTokenStatus.Active,
+      context: {
+        jitOrganizationIds: ['org_id'],
+      },
+    });
+    const ctx = createContextWithRouteParameters({
+      url: `/consent`,
+    });
+    const guard = koaConsentGuard(provider, mockTenant.libraries, mockTenant.queries);
+
+    await guard(ctx, next);
+
+    expect(provisionOrganizations).toHaveBeenCalledWith({
+      userId: 'foo',
+      organizationIds: ['org_id'],
+    });
+    expect(updateOneTimeTokenStatus).toHaveBeenCalledWith(
+      'token_value',
+      OneTimeTokenStatus.Consumed
+    );
+    expect(next).toHaveBeenCalled();
+  });
+
   it('should call next middleware if validations pass', async () => {
     const ctx = createContextWithRouteParameters({
       url: `/consent`,
@@ -80,9 +138,29 @@ describe('koaConsentGuard middleware', () => {
       // @ts-expect-error
       session: { accountId: 'foo' },
     });
-    const guard = koaConsentGuard(provider, mockQueries);
+    const guard = koaConsentGuard(provider, mockTenant.libraries, mockTenant.queries);
 
     await guard(ctx, next);
     expect(next).toHaveBeenCalled();
   });
+
+  it('should navigate to `/one-time-token` route with error message in URL params, if the one-time token is not valid', async () => {
+    interactionDetails.mockResolvedValue({
+      params: { one_time_token: 'token_value', login_hint: 'foo@example.com' },
+      // @ts-expect-error
+      session: { accountId: 'foo' },
+    });
+    checkOneTimeToken.mockImplementationOnce(() => {
+      throw new RequestError('one_time_token.token_expired');
+    });
+    const ctx = createContextWithRouteParameters({
+      url: `/consent`,
+    });
+    const guard = koaConsentGuard(provider, mockTenant.libraries, mockTenant.queries);
+
+    await guard(ctx, next);
+    expect(ctx.redirect).toHaveBeenCalledWith(
+      expect.stringContaining('one-time-token?errorMessage=The token is expired.')
+    );
+  });
 });

packages/core/src/middleware/koa-consent-guard.ts

@@ -1,9 +1,10 @@
-import { experience } from '@logto/schemas';
+import { experience, OneTimeTokenStatus } from '@logto/schemas';
 import { type MiddlewareType } from 'koa';
 import { type IRouterParamContext } from 'koa-router';
 import { type Provider } from 'oidc-provider';
 
 import RequestError from '#src/errors/RequestError/index.js';
+import type Libraries from '#src/tenants/Libraries.js';
 import type Queries from '#src/tenants/Queries.js';
 import assertThat from '#src/utils/assert-that.js';
 
@@ -15,7 +16,11 @@ export default function koaConsentGuard<
   StateT,
   ContextT extends IRouterParamContext,
   ResponseBodyT,
->(provider: Provider, query: Queries): MiddlewareType<StateT, ContextT, ResponseBodyT> {
+>(
+  provider: Provider,
+  libraries: Libraries,
+  queries: Queries
+): MiddlewareType<StateT, ContextT, ResponseBodyT> {
   return async (ctx, next) => {
     const interactionDetails = await provider.interactionDetails(ctx.req, ctx.res);
     const {
@@ -25,8 +30,9 @@ export default function koaConsentGuard<
 
     assertThat(session, new RequestError({ code: 'session.not_found' }));
 
+    // Handle one-time token before auto-consent
     if (token && loginHint && typeof token === 'string' && typeof loginHint === 'string') {
-      const { primaryEmail } = await query.users.findUserById(session.accountId);
+      const { primaryEmail } = await queries.users.findUserById(session.accountId);
 
       assertThat(primaryEmail, 'user.email_not_exist');
 
@@ -35,6 +41,33 @@ export default function koaConsentGuard<
         ctx.redirect(`${experience.routes.switchAccount}?${searchParams.toString()}`);
         return;
       }
+
+      try {
+        await libraries.oneTimeTokens.checkOneTimeToken(token, loginHint);
+        const {
+          context: { jitOrganizationIds },
+        } = await libraries.oneTimeTokens.updateOneTimeTokenStatus(
+          token,
+          OneTimeTokenStatus.Consumed
+        );
+
+        if (jitOrganizationIds) {
+          await libraries.users.provisionOrganizations({
+            userId: session.accountId,
+            organizationIds: jitOrganizationIds,
+          });
+        }
+      } catch (error: unknown) {
+        if (error instanceof RequestError) {
+          // Green light for token in consumed state
+          if (error.code === 'one_time_token.token_consumed') {
+            return next();
+          }
+          ctx.redirect(`${experience.routes.oneTimeToken}?errorMessage=${error.message}`);
+          return;
+        }
+        throw error;
+      }
     }
 
     return next();

packages/core/src/tenants/Tenant.ts

@@ -204,7 +204,10 @@ export default class Tenant implements TenantContext {
         koaSpaSessionGuard(provider, queries),
         mount(
           `/${experience.routes.consent}`,
-          compose([koaConsentGuard(provider, queries), koaAutoConsent(provider, queries)])
+          compose([
+            koaConsentGuard(provider, libraries, queries),
+            koaAutoConsent(provider, queries),
+          ])
         ),
         koaSpaProxy({ mountedApps, queries }),
       ])

packages/experience/src/App.tsx

@@ -69,7 +69,7 @@ const App = () => {
                     <Route element={<AppLayout />}>
                       {isDevFeaturesEnabled && (
                         <>
-                          <Route path="one-time-token" element={<OneTimeToken />} />
+                          <Route path={experience.routes.oneTimeToken} element={<OneTimeToken />} />
                           <Route
                             path={experience.routes.switchAccount}
                             element={<SwitchAccount />}

packages/experience/src/pages/OneTimeToken/index.tsx

@@ -25,7 +25,7 @@ import ErrorPage from '../ErrorPage';
 
 const OneTimeToken = () => {
   const [params] = useSearchParams();
-  const [oneTimeTokenError, setOneTimeTokenError] = useState<RequestErrorBody | boolean>();
+  const [oneTimeTokenError, setOneTimeTokenError] = useState<string | boolean>();
 
   const asyncIdentifyUserAndSubmit = useApi(identifyAndSubmitInteraction);
   const asyncSignInWithVerifiedIdentifier = useApi(signInWithVerifiedIdentifier);
@@ -91,6 +91,13 @@ const OneTimeToken = () => {
     (async () => {
       const token = params.get(ExtraParamsKey.OneTimeToken);
       const email = params.get(ExtraParamsKey.LoginHint);
+      const errorMessage = params.get('errorMessage');
+
+      if (errorMessage) {
+        setOneTimeTokenError(errorMessage);
+        return;
+      }
+
       if (!token || !email) {
         setOneTimeTokenError(true);
         return;
@@ -112,7 +119,7 @@ const OneTimeToken = () => {
       if (error) {
         await handleError(error, {
           global: (error: RequestErrorBody) => {
-            setOneTimeTokenError(error);
+            setOneTimeTokenError(error.message);
           },
         });
         return;
@@ -139,7 +146,7 @@ const OneTimeToken = () => {
         isNavbarHidden
         title="error.invalid_link"
         message="error.invalid_link_description"
-        rawMessage={condString(typeof oneTimeTokenError !== 'boolean' && oneTimeTokenError.message)}
+        rawMessage={condString(typeof oneTimeTokenError !== 'boolean' && oneTimeTokenError)}
       />
     );
   }

packages/experience/src/pages/Register/index.tsx

@@ -1,4 +1,4 @@
-import { AgreeToTermsPolicy, ExtraParamsKey, SignInMode } from '@logto/schemas';
+import { AgreeToTermsPolicy, experience, ExtraParamsKey, SignInMode } from '@logto/schemas';
 import { useCallback, useContext } from 'react';
 import { useTranslation } from 'react-i18next';
 import { Navigate, useNavigate, useSearchParams } from 'react-router-dom';
@@ -44,7 +44,10 @@ const RegisterFooter = () => {
 
   if (params.get(ExtraParamsKey.OneTimeToken)) {
     return (
-      <Navigate replace to={{ pathname: '/one-time-token', search: `?${params.toString()}` }} />
+      <Navigate
+        replace
+        to={{ pathname: `/${experience.routes.oneTimeToken}`, search: `?${params.toString()}` }}
+      />
     );
   }
 

packages/experience/src/pages/SignIn/index.tsx

@@ -1,4 +1,4 @@
-import { AgreeToTermsPolicy, ExtraParamsKey, SignInMode } from '@logto/schemas';
+import { AgreeToTermsPolicy, experience, ExtraParamsKey, SignInMode } from '@logto/schemas';
 import { useCallback, useContext } from 'react';
 import { useTranslation } from 'react-i18next';
 import { Navigate, useNavigate, useSearchParams } from 'react-router-dom';
@@ -107,7 +107,10 @@ const SignIn = () => {
 
   if (params.get(ExtraParamsKey.OneTimeToken)) {
     return (
-      <Navigate replace to={{ pathname: '/one-time-token', search: `?${params.toString()}` }} />
+      <Navigate
+        replace
+        to={{ pathname: `/${experience.routes.oneTimeToken}`, search: `?${params.toString()}` }}
+      />
     );
   }
 

packages/experience/src/pages/SwitchAccount/index.tsx

@@ -1,4 +1,4 @@
-import { type ConsentInfoResponse } from '@logto/schemas';
+import { experience, type ConsentInfoResponse } from '@logto/schemas';
 import { useContext, useEffect, useState } from 'react';
 import { useNavigate, useSearchParams } from 'react-router-dom';
 
@@ -82,7 +82,7 @@ const SwitchAccount = () => {
           i18nProps={{ name: loginHint }}
           onClick={() => {
             navigate(
-              { pathname: '/one-time-token', search: `?${params.toString()}` },
+              { pathname: `/${experience.routes.oneTimeToken}`, search: `?${params.toString()}` },
               { replace: true }
             );
           }}

packages/schemas/src/consts/experience.ts

@@ -7,7 +7,7 @@ const routes = Object.freeze({
   identifierSignIn: 'identifier-sign-in',
   identifierRegister: 'identifier-register',
   switchAccount: 'switch-account',
-  error: 'error',
+  oneTimeToken: 'one-time-token',
 } as const);
 
 export const experience = Object.freeze({