feat(core,schemas): implement sso token storage

packages/core/src/libraries/social.ts

@@ -1,10 +1,10 @@
 import { appInsights } from '@logto/app-insights/node';
-import type { GetSession, SocialUserInfo, TokenResponse } from '@logto/connector-kit';
+import type { GetSession, SocialUserInfo } from '@logto/connector-kit';
 import { socialUserInfoGuard } from '@logto/connector-kit';
 import type { EncryptedTokenSet, SecretSocialConnectorRelationPayload, User } from '@logto/schemas';
 import { ConnectorType } from '@logto/schemas';
 import { generateStandardId } from '@logto/shared';
-import { conditional, type Nullable } from '@silverhand/essentials';
+import { trySafe, type Nullable } from '@silverhand/essentials';
 import type { InteractionResults } from 'oidc-provider';
 import { z } from 'zod';
 
@@ -14,11 +14,7 @@
 import assertThat from '#src/utils/assert-that.js';
 import type { LogtoConnector } from '#src/utils/connectors/types.js';
 
-import {
-  deserializeEncryptedSecret,
-  encryptTokens,
-  serializeEncryptedSecret,
-} from '../utils/secret-encryption.js';
+import { deserializeEncryptedSecret, encryptTokenResponse } from '../utils/secret-encryption.js';
 
 const getUserInfoFromInteractionResult = async (
   connectorId: string,
@@ -43,46 +39,6 @@
   return result.socialUserInfo.userInfo;
 };
 
-const encryptTokenResponse = (tokenResponse?: TokenResponse): EncryptedTokenSet | undefined => {
-  if (!tokenResponse?.access_token) {
-    return;
-  }
-
-  try {
-    const {
-      access_token,
-      id_token,
-      refresh_token,
-      scope,
-      token_type: tokenType,
-      expires_in,
-    } = tokenResponse;
-
-    const requestedAt = Math.floor(Date.now() / 1000);
-
-    const expiresAt = expires_in && requestedAt + expires_in;
-
-    const encryptedTokenSet = encryptTokens({
-      access_token,
-      ...conditional(id_token && { id_token }),
-      ...conditional(refresh_token && { refresh_token }),
-    });
-
-    return {
-      encryptedTokenSetBase64: serializeEncryptedSecret(encryptedTokenSet),
-      metadata: {
-        scope,
-        tokenType,
-        expiresAt,
-      },
-    };
-  } catch (error: unknown) {
-    // Token encryption should not break the normal social authentication flow
-    // Return undefined to indicate no token response is available
-    void appInsights.trackException(error);
-  }
-};
-
 export const createSocialLibrary = (queries: Queries, connectorLibrary: ConnectorLibrary) => {
   const { findUserByEmail, findUserByNormalizedPhone } = queries.users;
   const { getLogtoConnectorById } = connectorLibrary;
@@ -162,7 +118,13 @@
         getConnectorSession
       );
 
-      const encryptedTokenSet = encryptTokenResponse(tokenResponse);
+      const encryptedTokenSet = trySafe(
+        () => encryptTokenResponse(tokenResponse),
+        (error) => {
+          // If the token response cannot be encrypted, we log the error but continue to return user info.
+          void appInsights.trackException(error);
+        }
+      );
 
       return {
         userInfo,

packages/core/src/libraries/sso-connector.ts

@@ -1,10 +1,13 @@
+import { appInsights } from '@logto/app-insights/node';
 import { type DirectSignInOptions, Prompt, QueryKey, ReservedScope, UserScope } from '@logto/js';
 import {
   ApplicationType,
   type SsoSamlAssertionContent,
   type CreateSsoConnectorIdpInitiatedAuthConfig,
   type SupportedSsoConnector,
   type SsoConnectorIdpInitiatedAuthConfig,
+  type EncryptedTokenSet,
+  type SecretEnterpriseSsoConnectorRelationPayload,
 } from '@logto/schemas';
 import { generateStandardId } from '@logto/shared';
 import { assert, deduplicate, trySafe } from '@silverhand/essentials';
@@ -17,6 +20,8 @@
 import assertThat from '#src/utils/assert-that.js';
 import { type OmitAutoSetFields } from '#src/utils/sql.js';
 
+import { deserializeEncryptedSecret } from '../utils/secret-encryption.js';
+
 export type SsoConnectorLibrary = ReturnType<typeof createSsoConnectorLibrary>;
 
 export const createSsoConnectorLibrary = (queries: Queries) => {
@@ -213,12 +218,41 @@
     return new URL(`${issuer}/auth?${queryParameters.toString()}`);
   };
 
+  const upsertEnterpriseSsoTokenSetSecret = async (
+    userId: string,
+    {
+      encryptedTokenSet,
+      enterpriseSsoConnectorRelationPayload,
+    }: {
+      encryptedTokenSet: EncryptedTokenSet;
+      enterpriseSsoConnectorRelationPayload: SecretEnterpriseSsoConnectorRelationPayload;
+    }
+  ) => {
+    const { encryptedTokenSetBase64, metadata } = encryptedTokenSet;
+
+    try {
+      await queries.secrets.upsertEnterpriseSsoTokenSetSecret(
+        {
+          id: generateStandardId(),
+          userId,
+          ...deserializeEncryptedSecret(encryptedTokenSetBase64),
+          metadata,
+        },
+        enterpriseSsoConnectorRelationPayload
+      );
+    } catch (error: unknown) {
+      // Upsert token set secret should not break the normal social authentication and link flow
+      void appInsights.trackException(error);
+    }
+  };
+
   return {
     getSsoConnectors,
     getAvailableSsoConnectors,
     getSsoConnectorById,
     createSsoConnectorIdpInitiatedAuthConfig,
     createIdpInitiatedSamlSsoSession,
     getIdpInitiatedSamlSsoSignInUrl,
+    upsertEnterpriseSsoTokenSetSecret,
   };
 };

packages/core/src/queries/secret.ts

@@ -2,6 +2,8 @@
   type CreateSecret,
   type CreateSocialTokenSetSecret,
   type Secret,
+  type SecretEnterpriseSsoConnectorRelationPayload,
+  SecretEnterpriseSsoConnectorRelations,
   type SecretKeys,
   Secrets,
   type SecretSocialConnectorRelationPayload,
@@ -17,6 +19,10 @@
 
 const secrets = convertToIdentifiers(Secrets, true);
 const secretSocialConnectorRelations = convertToIdentifiers(SecretSocialConnectorRelations, true);
+const secretEnterpriseSsoConnectorRelations = convertToIdentifiers(
+  SecretEnterpriseSsoConnectorRelations,
+  true
+);
 
 class SecretQueries extends SchemaQueries<SecretKeys, CreateSecret, Secret> {
   constructor(pool: CommonQueryMethods) {
@@ -80,6 +86,39 @@
           and ${secretSocialConnectorRelations.fields.target} = ${target}
     `);
   }
+
+  public async upsertEnterpriseSsoTokenSetSecret(
+    secret: Omit<CreateSocialTokenSetSecret, 'type'>,
+    ssoConnectorRelationPayload: SecretEnterpriseSsoConnectorRelationPayload
+  ) {
+    return this.pool.transaction(async (transaction) => {
+      // Delete any existing secret for the same identity and issuer
+      // This is to ensure only one secret exists for a given identity
+      await transaction.query(sql`
+        delete from ${secrets.table}
+        using ${secretEnterpriseSsoConnectorRelations.table}
+        where ${secrets.fields.id} = ${secretEnterpriseSsoConnectorRelations.fields.secretId}
+          and ${secrets.fields.userId} = ${secret.userId}
+          and ${secretEnterpriseSsoConnectorRelations.fields.issuer} = ${ssoConnectorRelationPayload.issuer}
+      `);
+
+      // Insert new secrets
+      const insertInto = buildInsertIntoWithPool(transaction)(Secrets, { returning: true });
+      const newSecret = await insertInto({
+        ...secret,
+        type: SecretType.FederatedTokenSet,
+      });
+
+      // Insert enterprise SSO connector relation
+      const insertEnterpriseSsoRelations = buildInsertIntoWithPool(transaction)(
+        SecretEnterpriseSsoConnectorRelations
+      );
+      await insertEnterpriseSsoRelations({
+        secretId: newSecret.id,
+        ...ssoConnectorRelationPayload,
+      });
+    });
+  }
 }
 
 export default SecretQueries;

packages/core/src/routes/experience/classes/experience-interaction.ts

@@ -409,6 +409,7 @@
       queries: { users: userQueries, userSsoIdentities: userSsoIdentityQueries },
       libraries: {
         socials: { upsertSocialTokenSetSecret },
+        ssoConnectors: { upsertEnterpriseSsoTokenSetSecret },
       },
     } = this.tenant;
 
@@ -464,6 +465,7 @@
       syncedEnterpriseSsoIdentity,
       jitOrganizationIds,
       socialConnectorTokenSetSecret,
+      enterpriseSsoConnectorTokenSetSecret,
       ...rest
     } = this.profile.data;
     const { mfaSkipped, mfaVerifications } = this.mfa.toUserMfaVerifications();
@@ -516,6 +518,11 @@
       await upsertSocialTokenSetSecret(user.id, socialConnectorTokenSetSecret);
     }
 
+    // Sync enterprise sso token set secret
+    if (enterpriseSsoConnectorTokenSetSecret) {
+      await upsertEnterpriseSsoTokenSetSecret(user.id, enterpriseSsoConnectorTokenSetSecret);
+    }
+
     // Provision organizations for one-time token that carries organization IDs in the context.
     if (jitOrganizationIds) {
       await this.provisionLibrary.provisionJitOrganization({

packages/core/src/routes/experience/classes/helpers.ts

@@ -43,13 +43,16 @@
         verificationRecord.toSyncedProfile(true),
       ]);
 
-      // TODO: Remove this check once enterprise SSO support getTokenSetSecret
-      const socialConnectorTokenSetSecret =
+      const tokenSetSecretProfile =
         verificationRecord.type === VerificationType.Social
-          ? await verificationRecord.getTokenSetSecret()
-          : undefined;
+          ? { socialConnectorTokenSetSecret: await verificationRecord.getTokenSetSecret() }
+          : { enterpriseSsoConnectorTokenSetSecret: verificationRecord.getTokenSetSecret() };
 
-      return { ...identityProfile, ...syncedProfile, socialConnectorTokenSetSecret };
+      return {
+        ...identityProfile,
+        ...syncedProfile,
+        ...tokenSetSecretProfile,
+      };
     }
     default: {
       // Unsupported verification type for user creation, such as MFA verification.
@@ -83,6 +86,7 @@
     | 'avatar'
     | 'name'
     | 'socialConnectorTokenSetSecret'
+    | 'enterpriseSsoConnectorTokenSetSecret'
   >;
 }> => {
   // Check verification record can be used to identify a user using the `identifyUser` method.
@@ -129,6 +133,7 @@
         const syncedProfile = {
           syncedEnterpriseSsoIdentity: enterpriseSsoIdentity,
           ...(await verificationRecord.toSyncedProfile()),
+          enterpriseSsoConnectorTokenSetSecret: verificationRecord.getTokenSetSecret(),
         };
         return { user, syncedProfile };
       } catch (error: unknown) {
@@ -139,6 +144,7 @@
           const syncedProfile = {
             ...verificationRecord.toUserProfile(),
             ...(await verificationRecord.toSyncedProfile()),
+            enterpriseSsoConnectorTokenSetSecret: verificationRecord.getTokenSetSecret(),
           };
           return { user, syncedProfile };
         }

packages/core/src/routes/experience/classes/libraries/provision-library.ts

@@ -59,6 +59,7 @@
       libraries: {
         users: { generateUserId, insertUser },
         socials: { upsertSocialTokenSetSecret },
+        ssoConnectors: { upsertEnterpriseSsoTokenSetSecret },
       },
     } = this.tenantContext;
 
@@ -68,6 +69,7 @@
       syncedEnterpriseSsoIdentity,
       jitOrganizationIds,
       socialConnectorTokenSetSecret,
+      enterpriseSsoConnectorTokenSetSecret,
       ...rest
     } = profile;
 
@@ -96,6 +98,10 @@
       await upsertSocialTokenSetSecret(user.id, socialConnectorTokenSetSecret);
     }
 
+    if (enterpriseSsoConnectorTokenSetSecret) {
+      await upsertEnterpriseSsoTokenSetSecret(user.id, enterpriseSsoConnectorTokenSetSecret);
+    }
+
     await this.provisionNewUserJitOrganizations(user.id, profile);
 
     this.ctx.appendDataHookContext('User.Created', { user });