diff --git a/api/strategies/openidStrategy.js b/api/strategies/openidStrategy.js index 3f2f9c22d10..99ad2b0c899 100644 --- a/api/strategies/openidStrategy.js +++ b/api/strategies/openidStrategy.js @@ -13,9 +13,11 @@ const { findOpenIDUser, getOpenIdEmail, getOpenIdIssuer, + createOpenIDUser, getBalanceConfig, selectOpenIdRole, getTokenCacheTtlMs, + applyOpenIDProfile, getAvatarSaveParams, isEmailDomainAllowed, getAvatarFileStrategy, @@ -28,7 +30,13 @@ const { } = require('@librechat/api'); const { getStrategyFunctions } = require('~/server/services/Files/strategies'); const { resizeAvatar } = require('~/server/services/Files/images/avatar'); -const { findUser, createUser, updateUser, findRolesByNames } = require('~/models'); +const { + findUser, + updateUser, + findRolesByNames, + findBalanceByUser, + createUserIfAbsent, +} = require('~/models'); const { getAppConfig } = require('~/server/services/Config'); const getLogStores = require('~/cache/getLogStores'); @@ -580,14 +588,15 @@ async function processOpenIDAuth(tokenset, existingUsersOnly = false) { throw new Error('Email domain not allowed'); } - const result = await findOpenIDUser({ + const lookup = { findUser, email: email, openidId: claims.sub || userinfo.sub, openidIssuer, idOnTheSource: claims.oid || userinfo.oid, strategyName: 'openidStrategy', - }); + }; + const result = await findOpenIDUser(lookup); let user = result.user; const error = result.error; @@ -595,7 +604,7 @@ async function processOpenIDAuth(tokenset, existingUsersOnly = false) { throw new Error(ErrorTypes.AUTH_FAILED); } - const appConfig = user?.tenantId ? await resolveAppConfigForUser(getAppConfig, user) : baseConfig; + let appConfig = user?.tenantId ? await resolveAppConfigForUser(getAppConfig, user) : baseConfig; if (!isEmailDomainAllowed(email, appConfig?.registration?.allowedDomains)) { logger.error( @@ -676,33 +685,28 @@ async function processOpenIDAuth(tokenset, existingUsersOnly = false) { throw new Error('User does not exist'); } - if (!user) { - user = { - provider: 'openid', - openidId: userinfo.sub, - username, - email: email || '', - emailVerified: userinfo.email_verified || false, - name: fullName, - idOnTheSource: userinfo.oid, - openidIssuer, - }; + const profile = { + openidId: userinfo.sub, + openidIssuer, + username, + name: fullName, + email, + emailVerified: userinfo.email_verified || false, + idOnTheSource: userinfo.oid, + }; - const balanceConfig = getBalanceConfig(appConfig); - user = await createUser(user, balanceConfig, true, true); + if (!user) { + ({ user, appConfig } = await createOpenIDUser({ + lookup, + profile, + appConfig, + getAppConfig, + getBalanceConfig, + createUserIfAbsent, + findBalanceByUser, + })); } else { - user.provider = 'openid'; - user.openidId = userinfo.sub; - if (openidIssuer) { - user.openidIssuer = openidIssuer; - } - user.username = username; - user.name = fullName; - user.idOnTheSource = userinfo.oid; - if (email && email !== user.email) { - user.email = email; - user.emailVerified = userinfo.email_verified || false; - } + user = applyOpenIDProfile(user, profile); } const adminRole = process.env.OPENID_ADMIN_ROLE; diff --git a/api/strategies/openidStrategy.spec.js b/api/strategies/openidStrategy.spec.js index 641619d1005..d62eb20200c 100644 --- a/api/strategies/openidStrategy.spec.js +++ b/api/strategies/openidStrategy.spec.js @@ -2,7 +2,7 @@ const undici = require('undici'); const fetch = require('node-fetch'); const jwtDecode = require('jsonwebtoken/decode'); const { ErrorTypes, FileSources } = require('librechat-data-provider'); -const { findUser, createUser, updateUser, findRolesByNames } = require('~/models'); +const { findUser, updateUser, findRolesByNames, createUserIfAbsent } = require('~/models'); const { getOpenIdProxyDispatcher, resolveAppConfigForUser, @@ -97,9 +97,10 @@ jest.mock('@librechat/api', () => { }); jest.mock('~/models', () => ({ findUser: jest.fn(), - createUser: jest.fn(), + createUserIfAbsent: jest.fn(), updateUser: jest.fn(), findRolesByNames: jest.fn(), + findBalanceByUser: jest.fn(), })); jest.mock('@librechat/data-schemas', () => ({ ...jest.requireActual('@librechat/api'), @@ -264,11 +265,11 @@ describe('setupOpenId', () => { permissions: ['admin'], }); - // By default, assume that no user is found, so createUser will be called + // By default, assume that no user is found, so createUserIfAbsent will be called findUser.mockResolvedValue(null); - createUser.mockImplementation(async (userData) => { + createUserIfAbsent.mockImplementation(async (userData) => { // simulate created user with an _id property - return { _id: 'newUserId', ...userData }; + return { ok: true, value: { _id: 'newUserId', ...userData } }; }); updateUser.mockImplementation(async (id, userData) => { return { _id: id, ...userData }; @@ -420,7 +421,7 @@ describe('setupOpenId', () => { expect(getOpenIdIssuer.mock.calls[0][1]).toEqual( expect.objectContaining({ issuer: 'https://fake-issuer.com' }), ); - expect(createUser).toHaveBeenCalledWith( + expect(createUserIfAbsent).toHaveBeenCalledWith( expect.objectContaining({ provider: 'openid', openidId: userinfo.sub, @@ -430,8 +431,6 @@ describe('setupOpenId', () => { name: `${userinfo.given_name} ${userinfo.family_name}`, }), { enabled: false }, - true, - true, ); }); @@ -447,11 +446,9 @@ describe('setupOpenId', () => { // Assert expect(user.username).toBe(expectUsername); - expect(createUser).toHaveBeenCalledWith( + expect(createUserIfAbsent).toHaveBeenCalledWith( expect.objectContaining({ username: expectUsername }), { enabled: false }, - true, - true, ); }); @@ -467,11 +464,9 @@ describe('setupOpenId', () => { // Assert expect(user.username).toBe(expectUsername); - expect(createUser).toHaveBeenCalledWith( + expect(createUserIfAbsent).toHaveBeenCalledWith( expect.objectContaining({ username: expectUsername }), { enabled: false }, - true, - true, ); }); @@ -485,11 +480,9 @@ describe('setupOpenId', () => { // Assert – username should equal the sub (converted as-is) expect(user.username).toBe(userinfo.sub); - expect(createUser).toHaveBeenCalledWith( + expect(createUserIfAbsent).toHaveBeenCalledWith( expect.objectContaining({ username: userinfo.sub }), { enabled: false }, - true, - true, ); }); @@ -575,7 +568,7 @@ describe('setupOpenId', () => { // Assert – verify that the strategy rejects login expect(result.user).toBe(false); expect(result.details.message).toBe(ErrorTypes.AUTH_FAILED); - expect(createUser).not.toHaveBeenCalled(); + expect(createUserIfAbsent).not.toHaveBeenCalled(); expect(updateUser).not.toHaveBeenCalled(); }); @@ -602,10 +595,109 @@ describe('setupOpenId', () => { expect(result.user).toBe(false); expect(result.details.message).toBe(ErrorTypes.AUTH_FAILED); - expect(createUser).not.toHaveBeenCalled(); + expect(createUserIfAbsent).not.toHaveBeenCalled(); expect(updateUser).not.toHaveBeenCalled(); }); + describe('concurrent first login', () => { + const raceCreateWith = (existingUser) => { + let created = false; + const matches = (query) => + Object.entries(query).every(([field, value]) => existingUser[field] === value); + findUser.mockImplementation(async (query) => + created && matches(query) ? existingUser : null, + ); + createUserIfAbsent.mockImplementation(async () => { + created = true; + return { ok: false, error: { code: 'user_exists' } }; + }); + }; + + it('continues as the user the other request created', async () => { + raceCreateWith({ + _id: 'winnerUserId', + provider: 'openid', + openidId: tokenset.claims().sub, + openidIssuer: 'https://fake-issuer.com', + email: tokenset.claims().email, + username: 'testusername', + name: 'First Last', + }); + + const { user } = await validate(tokenset); + + expect(createUserIfAbsent).toHaveBeenCalledTimes(1); + expect(updateUser).toHaveBeenCalledWith( + 'winnerUserId', + expect.objectContaining({ openidId: tokenset.claims().sub }), + ); + expect(user).toEqual( + expect.objectContaining({ _id: 'winnerUserId', email: tokenset.claims().email }), + ); + }); + + it("writes this callback's claims over the ones the other request stored", async () => { + raceCreateWith({ + _id: 'winnerUserId', + provider: 'openid', + openidId: tokenset.claims().sub, + openidIssuer: 'https://fake-issuer.com', + email: 'old@example.com', + emailVerified: false, + username: 'old-username', + name: 'Old Name', + }); + + await validate(tokenset); + + expect(updateUser).toHaveBeenCalledWith( + 'winnerUserId', + expect.objectContaining({ + email: tokenset.claims().email, + emailVerified: true, + username: tokenset.claims().preferred_username, + name: `${tokenset.claims().given_name} ${tokenset.claims().family_name}`, + }), + ); + }); + + it('fails the login while the account it found has no start balance yet', async () => { + const { findBalanceByUser } = require('~/models'); + const { getBalanceConfig } = require('@librechat/api'); + getBalanceConfig.mockReturnValue({ enabled: true, startBalance: 500 }); + findBalanceByUser.mockResolvedValue(null); + raceCreateWith({ + _id: 'winnerUserId', + provider: 'openid', + openidId: tokenset.claims().sub, + openidIssuer: 'https://fake-issuer.com', + email: tokenset.claims().email, + }); + + const result = await validate(tokenset); + + expect(findBalanceByUser).toHaveBeenCalledWith('winnerUserId'); + expect(result.user).toBe(false); + expect(result.details.message).toBe(ErrorTypes.AUTH_FAILED); + expect(updateUser).not.toHaveBeenCalled(); + getBalanceConfig.mockReturnValue({ enabled: false }); + }); + + it('fails the login when another provider took the email in the meantime', async () => { + raceCreateWith({ + _id: 'localUserId', + provider: 'local', + email: tokenset.claims().email, + }); + + const result = await validate(tokenset); + + expect(result.user).toBe(false); + expect(result.details.message).toBe(ErrorTypes.AUTH_FAILED); + expect(updateUser).not.toHaveBeenCalled(); + }); + }); + it('should enforce the required role and reject login if missing', async () => { // Arrange – simulate a token without the required role. jwtDecode.mockReturnValue({ @@ -650,7 +742,7 @@ describe('setupOpenId', () => { // Assert – login succeeds when required role is present after splitting expect(user).toBeTruthy(); - expect(createUser).toHaveBeenCalled(); + expect(createUserIfAbsent).toHaveBeenCalled(); }); it('should allow login when roles claim is a comma-separated string containing the required role', async () => { @@ -664,7 +756,7 @@ describe('setupOpenId', () => { // Assert – login succeeds when required role is present after splitting expect(user).toBeTruthy(); - expect(createUser).toHaveBeenCalled(); + expect(createUserIfAbsent).toHaveBeenCalled(); }); it('should allow login when roles claim is a mixed comma-and-space-separated string containing the required role', async () => { @@ -678,7 +770,7 @@ describe('setupOpenId', () => { // Assert – login succeeds when required role is present after splitting expect(user).toBeTruthy(); - expect(createUser).toHaveBeenCalled(); + expect(createUserIfAbsent).toHaveBeenCalled(); }); it('should reject login when roles claim is a space-separated string that does not contain the required role', async () => { @@ -710,7 +802,7 @@ describe('setupOpenId', () => { expect(user).toBeTruthy(); expect(user.email).toBe(tokenset.claims().email); expect(user.username).toBe(tokenset.claims().preferred_username); - expect(createUser).toHaveBeenCalled(); + expect(createUserIfAbsent).toHaveBeenCalled(); }); it('should allow login when required role is found in userinfo claims', async () => { @@ -2422,11 +2514,9 @@ describe('setupOpenId', () => { const { user } = await validate({ ...tokenset, claims: () => userinfo }); expect(user.email).toBe('user@corp.example.com'); - expect(createUser).toHaveBeenCalledWith( + expect(createUserIfAbsent).toHaveBeenCalledWith( expect.objectContaining({ email: 'user@corp.example.com' }), expect.anything(), - true, - true, ); }); diff --git a/packages/api/src/auth/openid.spec.ts b/packages/api/src/auth/openid.spec.ts index 57e9aedd7a0..5548bba3128 100644 --- a/packages/api/src/auth/openid.spec.ts +++ b/packages/api/src/auth/openid.spec.ts @@ -1,11 +1,17 @@ import mongoose, { Types } from 'mongoose'; import { ErrorTypes } from 'librechat-data-provider'; import { MongoMemoryServer } from 'mongodb-memory-server'; -import { logger, createMethods, createModels } from '@librechat/data-schemas'; -import type { IUser, UserMethods } from '@librechat/data-schemas'; +import { logger, tenantStorage, createMethods, createModels } from '@librechat/data-schemas'; +import type { IUser, AppConfig, UserMethods } from '@librechat/data-schemas'; import type { CommandStartedEvent } from 'mongodb'; import type { FilterQuery } from 'mongoose'; -import { findOpenIDUser, getOpenIdEmail, getOpenIdIssuer, normalizeOpenIdIssuer } from './openid'; +import { + findOpenIDUser, + getOpenIdEmail, + getOpenIdIssuer, + createOpenIDUser, + normalizeOpenIdIssuer, +} from './openid'; import { recordOpenIDUserLookup } from '~/app/metrics'; function newId() { @@ -17,6 +23,7 @@ jest.mock('@librechat/data-schemas', () => ({ logger: { warn: jest.fn(), info: jest.fn(), + error: jest.fn(), }, })); @@ -997,3 +1004,214 @@ describe('getOpenIdEmail', () => { ); }); }); + +describe('createOpenIDUser', () => { + let mongoServer: MongoMemoryServer; + let User: mongoose.Model; + let methods: ReturnType; + + const issuer = 'https://issuer.example.com'; + const email = 'first-login@example.com'; + const openidId = 'first-login-sub'; + const tenantId = 'tenant-a'; + const profile = { + openidId, + openidIssuer: issuer, + email, + username: 'first-login', + name: 'First Login', + emailVerified: true, + }; + const getAppConfig = jest.fn(async (): Promise => appConfigWith()); + + function appConfigWith( + balance?: { enabled: boolean; startBalance: number }, + allowedDomains?: string[], + ): AppConfig { + return { balance, registration: { allowedDomains } } as Partial as AppConfig; + } + + function getLookup() { + return { findUser: methods.findUser, email, openidId, openidIssuer: issuer }; + } + + function firstLogin( + balance?: { enabled: boolean; startBalance: number }, + overrides: Partial[0]> = {}, + ) { + return createOpenIDUser({ + lookup: getLookup(), + profile, + appConfig: appConfigWith(balance), + getAppConfig, + getBalanceConfig: (config: AppConfig) => config.balance ?? null, + createUserIfAbsent: methods.createUserIfAbsent, + findBalanceByUser: (userId) => methods.findBalanceByUser(userId), + ...overrides, + }); + } + + beforeAll(async () => { + mongoServer = await MongoMemoryServer.create(); + await mongoose.connect(mongoServer.getUri(), { autoIndex: false }); + createModels(mongoose); + User = mongoose.models.User as mongoose.Model; + methods = createMethods(mongoose); + }); + + afterAll(async () => { + await mongoose.disconnect(); + await mongoServer.stop(); + }); + + beforeEach(async () => { + jest.clearAllMocks(); + await mongoose.connection.dropDatabase(); + await User.syncIndexes(); + }); + + it('resolves concurrent first logins to the one user that was created', async () => { + const lookups = await Promise.all([findOpenIDUser(getLookup()), findOpenIDUser(getLookup())]); + expect(lookups.map((lookup) => lookup.user)).toEqual([null, null]); + + const logins = await Promise.all([firstLogin(), firstLogin()]); + + expect(await User.countDocuments({ email })).toBe(1); + expect(logins[0].user._id.toString()).toBe(logins[1].user._id.toString()); + expect(logins[1].user).toEqual( + expect.objectContaining({ openidId, email, provider: 'openid' }), + ); + }); + + it("refreshes the recovered user with the losing request's claims", async () => { + await firstLogin(); + const renamed = { + ...profile, + email: 'renamed@example.com', + name: 'Renamed', + emailVerified: false, + }; + + const { user } = await firstLogin(undefined, { + lookup: { ...getLookup(), email: renamed.email }, + profile: renamed, + }); + + expect(await User.countDocuments()).toBe(1); + expect(user).toEqual( + expect.objectContaining({ + openidId, + email: 'renamed@example.com', + name: 'Renamed', + emailVerified: false, + }), + ); + }); + + it('credits the start balance once across concurrent first logins', async () => { + const balance = { enabled: true, startBalance: 500 }; + + const [{ user }] = await Promise.all([firstLogin(balance), firstLogin(balance)]); + + const balances = await mongoose.models.Balance.find({ user: user._id }).lean< + Array<{ tokenCredits: number }> + >(); + expect(balances.map((record) => record.tokenCredits)).toEqual([500]); + }); + + it('fails the recovered login while the winner has not credited its start balance yet', async () => { + await User.create({ email, provider: 'openid', openidId, openidIssuer: issuer }); + + const login = firstLogin({ enabled: true, startBalance: 500 }); + + await expect(login).rejects.toThrow(ErrorTypes.AUTH_FAILED); + expect(await mongoose.models.Balance.countDocuments()).toBe(0); + }); + + it('recovers an account created without a balance when no start balance is configured', async () => { + await User.create({ email, provider: 'openid', openidId, openidIssuer: issuer }); + + const { user } = await firstLogin({ enabled: true, startBalance: 0 }); + + expect(user).toEqual(expect.objectContaining({ openidId, email })); + }); + + it("applies the recovered tenant account's email-domain policy", async () => { + await tenantStorage.run({ tenantId }, () => + User.create({ email, provider: 'openid', openidId, openidIssuer: issuer }), + ); + getAppConfig.mockResolvedValueOnce(appConfigWith(undefined, ['other.example'])); + + const login = tenantStorage.run({ tenantId }, () => firstLogin()); + + await expect(login).rejects.toThrow('Email domain not allowed'); + expect(getAppConfig).toHaveBeenCalledWith(expect.objectContaining({ tenantId })); + }); + + it("continues when only the recovered tenant account's config sets a start balance", async () => { + await tenantStorage.run({ tenantId }, () => + User.create({ email, provider: 'openid', openidId, openidIssuer: issuer }), + ); + getAppConfig.mockResolvedValueOnce(appConfigWith({ enabled: true, startBalance: 500 })); + + const login = await tenantStorage.run({ tenantId }, () => firstLogin()); + + expect(login.user).toEqual(expect.objectContaining({ openidId, tenantId })); + }); + + it('holds the recovered login on the start balance new users are created with', async () => { + await tenantStorage.run({ tenantId }, () => + User.create({ email, provider: 'openid', openidId, openidIssuer: issuer }), + ); + getAppConfig.mockResolvedValueOnce(appConfigWith({ enabled: false, startBalance: 0 })); + + const login = tenantStorage.run({ tenantId }, () => + firstLogin({ enabled: true, startBalance: 500 }), + ); + + await expect(login).rejects.toThrow(ErrorTypes.AUTH_FAILED); + }); + + it("continues under the recovered tenant account's config when its policy admits the email", async () => { + await tenantStorage.run({ tenantId }, () => + User.create({ email, provider: 'openid', openidId, openidIssuer: issuer }), + ); + const tenantConfig = appConfigWith(undefined, ['example.com']); + getAppConfig.mockResolvedValueOnce(tenantConfig); + + const login = await tenantStorage.run({ tenantId }, () => firstLogin()); + + expect(login.appConfig).toBe(tenantConfig); + expect(login.user).toEqual(expect.objectContaining({ openidId, tenantId })); + }); + + it('fails the login when the email was taken by another provider in the meantime', async () => { + await User.create({ email, provider: 'local', username: 'local-user' }); + + await expect(firstLogin()).rejects.toThrow(ErrorTypes.AUTH_FAILED); + expect(await User.countDocuments({ email })).toBe(1); + }); + + it('throws when the lookup cannot account for the existing account', async () => { + await User.create({ email: 'taken@example.com', provider: 'openid', openidId: 'other-sub' }); + + const login = firstLogin(undefined, { + lookup: { findUser: methods.findUser, openidId, openidIssuer: issuer }, + profile: { ...profile, email: 'taken@example.com' }, + }); + + await expect(login).rejects.toThrow('conflicts with an account the lookup cannot resolve'); + }); + + it('rethrows other create failures without a second lookup', async () => { + const findUser = jest.fn(methods.findUser); + + const login = firstLogin(undefined, { + lookup: { ...getLookup(), findUser }, + profile: { ...profile, email: 'not-an-email' }, + }); + + await expect(login).rejects.toMatchObject({ name: 'ValidationError' }); + expect(findUser).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/api/src/auth/openid.ts b/packages/api/src/auth/openid.ts index 4d346a7591f..e121f484948 100644 --- a/packages/api/src/auth/openid.ts +++ b/packages/api/src/auth/openid.ts @@ -1,9 +1,19 @@ import { logger } from '@librechat/data-schemas'; import { ErrorTypes } from 'librechat-data-provider'; -import type { IUser, UserMethods } from '@librechat/data-schemas'; +import type { + IUser, + UserMethods, + UserRecord, + NewUserData, + BalanceConfig, + CreateUserIfAbsentResult, + AppConfig, +} from '@librechat/data-schemas'; import type { FilterQuery } from 'mongoose'; -import { isMetricsConfigured, recordOpenIDUserLookup } from '~/app/metrics'; import type { OpenIDUserLookupResult } from '~/app/metrics'; +import { isMetricsConfigured, recordOpenIDUserLookup } from '~/app/metrics'; +import { resolveAppConfigForUser } from '~/app/resolve'; +import { isEmailDomainAllowed } from './domain'; export type OpenIdEmailClaims = { email?: unknown; @@ -21,6 +31,17 @@ export type OpenIdIssuerSource = { type OpenIdLookupField = 'openidId' | 'idOnTheSource'; type OpenIdUserResolution = { user: IUser | null; error: string | null; migration: boolean }; +/** The account fields an OpenID login writes from the current callback's claims. */ +export type OpenIDProfile = { + openidId: string; + openidIssuer?: string; + username: string; + name: string; + email?: string; + emailVerified: boolean; + idOnTheSource?: string; +}; + const OPENID_DISCOVERY_PATH = '/.well-known/openid-configuration'; const LEGACY_ISSUER_FILTERS: Array['openidIssuer']> = [ { $exists: false }, @@ -297,3 +318,113 @@ export async function findOpenIDUser({ throw error; } } + +/** Returns `user` carrying the current callback's claims, as every OpenID login refreshes them. */ +export function applyOpenIDProfile(user: T, profile: OpenIDProfile): T { + const updated = { + ...user, + provider: 'openid', + openidId: profile.openidId, + username: profile.username, + name: profile.name, + idOnTheSource: profile.idOnTheSource, + }; + if (profile.openidIssuer) { + updated.openidIssuer = profile.openidIssuer; + } + if (profile.email && profile.email !== user.email) { + updated.email = profile.email; + updated.emailVerified = profile.emailVerified; + } + return updated; +} + +/** Whether new users start with a balance record (the condition data-schemas credits it under). */ +function hasStartBalance(balanceConfig?: BalanceConfig | null): boolean { + return Boolean(balanceConfig?.enabled && balanceConfig.startBalance); +} + +/** + * Creates a first-login OpenID user. Concurrent first logins for one identity all miss + * `findOpenIDUser`, and `createUserIfAbsent` reports `user_exists` for every insert but the + * first; a rejected request repeats the lookup, with its provider and issuer checks, and + * continues as the account that won, admitted and refreshed exactly as if its first lookup had + * found it: a tenant account resolves its tenant config and that config's email-domain policy + * applies, and the account carries this callback's claims. Returns the account with the config + * the login continues under (`appConfig` for a user this request created). + * + * It continues only once the winner finished provisioning that account. When the config new + * users are created under (`appConfig`) sets a start balance, the balance must exist: + * `createUserIfAbsent` writes it before the user, but a winner on an earlier release adds it with + * `$inc` after its insert, and login balance sync must not initialize it first; such a login fails + * as it did before recovery existed. A start balance only the account's tenant config sets is + * initialized by login balance sync with an insert-only write that cannot be added on top, so it + * does not hold the login back. A conflict the lookup cannot account for throws. + */ +export async function createOpenIDUser({ + lookup, + profile, + appConfig, + getAppConfig, + getBalanceConfig, + createUserIfAbsent, + findBalanceByUser, +}: { + lookup: Parameters[0]; + profile: OpenIDProfile; + appConfig: AppConfig; + getAppConfig: Parameters[0]; + getBalanceConfig: (appConfig: AppConfig) => BalanceConfig | null | undefined; + createUserIfAbsent: ( + data: NewUserData, + balanceConfig?: BalanceConfig, + ) => Promise; + findBalanceByUser: (userId: string) => Promise; +}): Promise<{ user: UserRecord; appConfig: AppConfig }> { + const created = await createUserIfAbsent( + { + provider: 'openid', + openidId: profile.openidId, + username: profile.username, + email: profile.email || '', + emailVerified: profile.emailVerified, + name: profile.name, + idOnTheSource: profile.idOnTheSource, + openidIssuer: profile.openidIssuer, + }, + getBalanceConfig(appConfig) ?? undefined, + ); + if (created.ok) return { user: created.value, appConfig }; + + const strategyName = lookup.strategyName ?? 'openid'; + const resolution = await findOpenIDUser(lookup); + if (resolution.error) throw new Error(ErrorTypes.AUTH_FAILED); + if (!resolution.user) { + throw new Error( + `[${strategyName}] New user conflicts with an account the lookup cannot resolve`, + ); + } + + const userId = resolution.user._id.toString(); + const accountConfig = resolution.user.tenantId + ? await resolveAppConfigForUser(getAppConfig, resolution.user) + : appConfig; + if (!isEmailDomainAllowed(profile.email ?? '', accountConfig?.registration?.allowedDomains)) { + logger.error( + `[${strategyName}] Authentication blocked - email domain not allowed for the recovered account [Identifier: ${profile.email}]`, + ); + throw new Error('Email domain not allowed'); + } + + if (hasStartBalance(getBalanceConfig(appConfig)) && !(await findBalanceByUser(userId))) { + logger.warn( + `[${strategyName}] Concurrent first login found user ${userId} before its start balance; failing this login`, + ); + throw new Error(ErrorTypes.AUTH_FAILED); + } + + logger.info( + `[${strategyName}] Concurrent first login for user ${userId}; continuing as that user`, + ); + return { user: applyOpenIDProfile(resolution.user, profile), appConfig: accountConfig }; +} diff --git a/packages/data-schemas/src/methods/user.methods.spec.ts b/packages/data-schemas/src/methods/user.methods.spec.ts index acf4bdafd9f..945260b5989 100644 --- a/packages/data-schemas/src/methods/user.methods.spec.ts +++ b/packages/data-schemas/src/methods/user.methods.spec.ts @@ -424,6 +424,127 @@ describe('User Methods - Database Tests', () => { }); }); + describe('createUserIfAbsent', () => { + const newUser = { + name: 'First Login', + email: 'first@example.com', + provider: 'openid', + openidId: 'first-sub', + openidIssuer: 'https://issuer.example.com', + }; + + beforeEach(async () => { + await User.syncIndexes(); + }); + + test('creates the user, credits the start balance, and returns it', async () => { + const result = await methods.createUserIfAbsent(newUser, { + enabled: true, + startBalance: 500, + }); + + expect(result.ok).toBe(true); + const created = result.ok ? result.value : null; + expect(created).toEqual(expect.objectContaining({ email: 'first@example.com' })); + const balance = await Balance.findOne({ user: created?._id }); + expect(balance?.tokenCredits).toBe(500); + }); + + test('reports user_exists for one of two concurrent inserts and credits the balance once', async () => { + const balanceConfig = { enabled: true, startBalance: 500 }; + + const results = await Promise.all([ + methods.createUserIfAbsent(newUser, balanceConfig), + methods.createUserIfAbsent(newUser, balanceConfig), + ]); + + expect(results.filter((result) => result.ok)).toHaveLength(1); + expect(results.find((result) => !result.ok)).toEqual({ + ok: false, + error: { code: 'user_exists' }, + }); + expect(await User.countDocuments()).toBe(1); + const balances = await Balance.find({}).lean(); + expect(balances.map((balance) => balance.tokenCredits)).toEqual([500]); + }); + + test('reports user_exists when another account holds the provider identity', async () => { + await methods.createUserIfAbsent(newUser); + + const result = await methods.createUserIfAbsent({ ...newUser, email: 'renamed@example.com' }); + + expect(result).toEqual({ ok: false, error: { code: 'user_exists' } }); + }); + + test('initializes the start balance before the user exists', async () => { + const balanceWrite = jest.spyOn(Balance, 'findOneAndUpdate'); + const userSave = jest.spyOn(User.prototype, 'save'); + + await methods.createUserIfAbsent(newUser, { enabled: true, startBalance: 500 }); + + expect(balanceWrite.mock.invocationCallOrder[0]).toBeLessThan( + userSave.mock.invocationCallOrder[0], + ); + balanceWrite.mockRestore(); + userSave.mockRestore(); + }); + + test('removes the start balance it initialized when the account already exists', async () => { + const balanceConfig = { enabled: true, startBalance: 500 }; + const created = await methods.createUserIfAbsent(newUser, balanceConfig); + + const result = await methods.createUserIfAbsent(newUser, balanceConfig); + + expect(result).toEqual({ ok: false, error: { code: 'user_exists' } }); + const balances = await Balance.find({}).lean(); + expect(balances.map((balance) => balance.user.toString())).toEqual([ + created.ok ? created.value._id.toString() : '', + ]); + }); + + test('removes the start balance when its write commits but reports a failure', async () => { + const realWrite = Balance.findOneAndUpdate.bind(Balance); + const balanceWrite = jest.spyOn(Balance, 'findOneAndUpdate').mockImplementationOnce( + (...args: Parameters) => + ({ + lean: async () => { + await realWrite(...args).lean(); + throw new Error('connection closed before acknowledgement'); + }, + }) as never, + ); + + await expect( + methods.createUserIfAbsent(newUser, { enabled: true, startBalance: 500 }), + ).rejects.toThrow('connection closed before acknowledgement'); + expect(await Balance.countDocuments()).toBe(0); + expect(await User.countDocuments()).toBe(0); + balanceWrite.mockRestore(); + }); + + test('keeps the start balance when the insert fails without a unique-index rejection', async () => { + const userSave = jest + .spyOn(User.prototype, 'save') + .mockRejectedValueOnce(new Error('connection closed before acknowledgement')); + + await expect( + methods.createUserIfAbsent(newUser, { enabled: true, startBalance: 500 }), + ).rejects.toThrow('connection closed before acknowledgement'); + expect(await Balance.countDocuments()).toBe(1); + userSave.mockRestore(); + }); + + test('throws failures other than an existing account without writing a balance', async () => { + await expect( + methods.createUserIfAbsent( + { ...newUser, email: 'not-an-email' }, + { enabled: true, startBalance: 500 }, + ), + ).rejects.toMatchObject({ name: 'ValidationError' }); + expect(await Balance.countDocuments()).toBe(0); + }); + }); + describe('updateUser', () => { test('should update user fields', async () => { const user = await User.create({ diff --git a/packages/data-schemas/src/methods/user.ts b/packages/data-schemas/src/methods/user.ts index 5237d415ec4..a78e911968a 100644 --- a/packages/data-schemas/src/methods/user.ts +++ b/packages/data-schemas/src/methods/user.ts @@ -5,7 +5,15 @@ import type { RefillIntervalUnit, StatefulCodeEnvironment, } from 'librechat-data-provider'; -import type { IUser, BalanceConfig, CreateUserRequest, UserDeleteResult } from '~/types'; +import type { + IUser, + BalanceConfig, + CreateUserRequest, + UserRecord, + NewUserData, + UserDeleteResult, + CreateUserIfAbsentResult, +} from '~/types'; import type { TwoFactorEnrollmentGuard, TwoFactorEnrollmentUpdate } from '~/types'; import type { CacheStore } from '~/types'; import { evictAuthUserDocs } from '~/utils/eviction'; @@ -64,6 +72,10 @@ export function createUserMethods( disableTTL?: boolean, returnUser?: boolean, ) => Promise>; + createUserIfAbsent: ( + data: NewUserData, + balanceConfig?: BalanceConfig, + ) => Promise; updateUser: ( userId: string, updateData: Partial, @@ -266,6 +278,48 @@ export function createUserMethods( return await User.countDocuments(filter); } + /** + * Initializes a new user's start balance, with auto-refill settings when complete. The write is + * insert-only, so it never replaces a balance that already exists or the activity recorded on it. + */ + async function creditStartBalance( + userId: mongoose.Types.ObjectId, + balanceConfig?: BalanceConfig, + ): Promise { + if (!balanceConfig?.enabled || !balanceConfig?.startBalance) { + return; + } + + const Balance = mongoose.models.Balance; + const initial: { + tokenCredits: number; + autoRefillEnabled?: boolean; + refillIntervalValue?: number; + refillIntervalUnit?: RefillIntervalUnit; + refillAmount?: number; + refillMode?: BalanceRefillMode; + } = { tokenCredits: balanceConfig.startBalance }; + + if ( + balanceConfig.autoRefillEnabled && + balanceConfig.refillIntervalValue != null && + balanceConfig.refillIntervalUnit != null && + balanceConfig.refillAmount != null + ) { + initial.autoRefillEnabled = true; + initial.refillIntervalValue = balanceConfig.refillIntervalValue; + initial.refillIntervalUnit = balanceConfig.refillIntervalUnit; + initial.refillAmount = balanceConfig.refillAmount; + initial.refillMode = balanceConfig.refillMode ?? 'add'; + } + + await Balance.findOneAndUpdate( + { _id: userId }, + { $setOnInsert: { ...initial, user: userId } }, + { upsert: true, new: true }, + ).lean(); + } + /** * Creates a new user, optionally with a TTL of 1 week. */ @@ -276,7 +330,6 @@ export function createUserMethods( returnUser: boolean = false, ): Promise> { const User = mongoose.models.User; - const Balance = mongoose.models.Balance; const userData: Partial = { ...data, @@ -288,48 +341,71 @@ export function createUserMethods( } const user = await User.create(userData); + await creditStartBalance(user._id, balanceConfig); - // If balance is enabled, create or update a balance record for the user - if (balanceConfig?.enabled && balanceConfig?.startBalance) { - const update: { - $inc: { tokenCredits: number }; - $set?: { - autoRefillEnabled: boolean; - refillIntervalValue: number; - refillIntervalUnit: RefillIntervalUnit; - refillAmount: number; - refillMode?: BalanceRefillMode; - }; - } = { - $inc: { tokenCredits: balanceConfig.startBalance }, - }; + if (returnUser) { + return user.toObject() as Partial; + } + return user._id as mongoose.Types.ObjectId; + } - if ( - balanceConfig.autoRefillEnabled && - balanceConfig.refillIntervalValue != null && - balanceConfig.refillIntervalUnit != null && - balanceConfig.refillAmount != null - ) { - update.$set = { - autoRefillEnabled: true, - refillIntervalValue: balanceConfig.refillIntervalValue, - refillIntervalUnit: balanceConfig.refillIntervalUnit, - refillAmount: balanceConfig.refillAmount, - refillMode: balanceConfig.refillMode ?? 'add', - }; - } + /** + * Creates a user without a TTL, or reports `user_exists` when a unique email or provider + * identity index already holds the account, as when concurrent first logins race to insert it. + * The start balance is initialized under the new user's id before the user is inserted, so + * the account is never visible without it and login balance sync never initializes it first. + * A balance write that fails, or an insert a unique index rejects, removes the balance under + * the id that never became a user; any other insert failure keeps it, since an unacknowledged + * insert may still have committed. + */ + async function createUserIfAbsent( + data: NewUserData, + balanceConfig?: BalanceConfig, + ): Promise { + const User = mongoose.models.User as mongoose.Model; + const userData: Partial = { ...data }; + delete userData.expiresAt; - await Balance.findOneAndUpdate( - { _id: user._id }, - { ...update, $setOnInsert: { user: user._id } }, - { upsert: true, new: true }, - ).lean(); + const user = new User(userData); + await user.validate(); + try { + await creditStartBalance(user._id, balanceConfig); + } catch (error) { + await discardStartBalance(user._id, balanceConfig); + throw error; } - if (returnUser) { - return user.toObject() as Partial; + try { + await user.save({ validateBeforeSave: false }); + } catch (error) { + if ((error as { code?: number }).code !== 11000) { + throw error; + } + await discardStartBalance(user._id, balanceConfig); + return { ok: false, error: { code: 'user_exists' } }; + } + + return { ok: true, value: user.toObject() as UserRecord }; + } + + /** Removes, best effort, the start balance initialized for an id that never became a user. */ + async function discardStartBalance( + userId: mongoose.Types.ObjectId, + balanceConfig?: BalanceConfig, + ): Promise { + if (!balanceConfig?.enabled || !balanceConfig?.startBalance) { + return; + } + try { + await mongoose.models.Balance.deleteOne({ _id: userId }); + } catch { + logger.warn( + '[createUserIfAbsent] Could not remove the start balance of a user never created', + { + userId: userId.toString(), + }, + ); } - return user._id as mongoose.Types.ObjectId; } /** @@ -949,6 +1025,7 @@ export function createUserMethods( findOwnerContactUsers, countUsers, createUser, + createUserIfAbsent, updateUser, awaitAuthUserDocEviction, consumeBackupCode, diff --git a/packages/data-schemas/src/types/user.ts b/packages/data-schemas/src/types/user.ts index 138de99ebc7..53f5bc8d231 100644 --- a/packages/data-schemas/src/types/user.ts +++ b/packages/data-schemas/src/types/user.ts @@ -138,6 +138,23 @@ export interface CreateUserRequest extends Partial { email: string; } +/** A user's own fields, without Mongoose document members. */ +type UserFields = Omit; + +/** + * A stored user as plain data with a storage-neutral id: what `.lean()` queries and `.toObject()` + * return. New contracts take and return this rather than the `IUser` document type. + */ +export type UserRecord = UserFields & { _id: { toString(): string } }; + +/** The fields a new user is created with, as plain data. */ +export type NewUserData = Partial & { email: string }; + +/** The created user, or `user_exists` when an account already holds its email or provider identity. */ +export type CreateUserIfAbsentResult = + | { ok: true; value: UserRecord } + | { ok: false; error: { code: 'user_exists' } }; + export interface UpdateUserRequest { name?: string; username?: string;