refactor: user existing check

app/server/lib/auth.ts

@@ -19,6 +19,7 @@ import { convertLegacyUserOnFirstLogin } from "./auth/middlewares/convert-legacy
 import { createUserDefaultOrg } from "./auth/helpers/create-default-org";
 import { isSsoCallbackRequest, requireSsoInvitation } from "./auth/middlewares/require-sso-invitation";
 import { ssoTrustedProviderLinkingPlugin } from "./auth/plugins/sso-trusted-provider-linking";
+import { logger } from "../utils/logger";
 
 export type AuthMiddlewareContext = MiddlewareContext<MiddlewareOptions, AuthContext<BetterAuthOptions>>;
 
@@ -52,6 +53,7 @@ export const auth = betterAuth({
 			create: {
 				before: async (user, ctx) => {
 					if (isSsoCallbackRequest(ctx)) {
+						logger.debug("SSO callback detected, checking for pending invitations for user", user.email);
 						await requireSsoInvitation(user.email, ctx);
 						user.hasDownloadedResticPassword = true;
 					}

app/server/lib/auth/helpers/create-default-org.ts

@@ -1,18 +1,17 @@
-import { and, eq, gt, inArray } from "drizzle-orm";
+import { eq } from "drizzle-orm";
 import { UnauthorizedError } from "http-errors-enhanced";
 import type { GenericEndpointContext } from "@better-auth/core";
 import { db } from "~/server/db/db";
-import { invitation, member, organization, ssoProvider } from "~/server/db/schema";
+import { invitation, member, organization, type User } from "~/server/db/schema";
 import { cryptoUtils } from "~/server/utils/crypto";
 import { APIError } from "better-auth";
 import { extractProviderIdFromContext, normalizeEmail } from "../utils/sso-context";
+import { logger } from "~/server/utils/logger";
 
 async function findMembershipWithOrganization(userId: string, organizationId?: string) {
 	if (organizationId) {
 		return db.query.member.findFirst({
-			where: {
-				AND: [{ userId }, { organizationId }],
-			},
+			where: { AND: [{ userId }, { organizationId }] },
 			with: { organization: true },
 		});
 	}
@@ -23,95 +22,52 @@ async function findMembershipWithOrganization(userId: string, organizationId?: s
 	});
 }
 
-async function findUserById(userId: string) {
-	return db.query.usersTable.findFirst({
-		where: { id: userId },
-	});
-}
-
 function buildOrgSlug(email: string) {
 	const [emailPrefix] = email.split("@");
 	return `${emailPrefix}-${Math.random().toString(36).slice(-4)}`;
 }
 
-async function getProviderOrganizationIdsFromContext(ctx: GenericEndpointContext | null) {
+async function getProviderFromContext(ctx: GenericEndpointContext | null) {
 	if (!ctx) {
-		return [];
+		return null;
 	}
 
 	const providerId = extractProviderIdFromContext(ctx);
 	if (!providerId) {
-		return [];
-	}
-
-	const provider = await db
-		.select({ organizationId: ssoProvider.organizationId })
-		.from(ssoProvider)
-		.where(eq(ssoProvider.providerId, providerId))
-		.limit(1);
-
-	return provider.map((providerRow) => providerRow.organizationId);
-}
-
-type UserRecord = NonNullable<Awaited<ReturnType<typeof findUserById>>>;
-
-async function tryCreateInvitedMembership(userId: string, email: string, ctx: GenericEndpointContext | null) {
-	let providerOrganizationIds = await getProviderOrganizationIdsFromContext(ctx);
-
-	if (providerOrganizationIds.length === 0) {
-		const linkedAccounts = await db.query.account.findMany({
-			where: { userId },
-		});
-
-		const linkedSsoProviderIds = linkedAccounts
-			.map((linkedAccount) => linkedAccount.providerId)
-			.filter((providerId) => providerId !== "credential");
-
-		if (linkedSsoProviderIds.length === 0) {
-			return null;
-		}
-
-		const linkedSsoProviders = await db
-			.select({ organizationId: ssoProvider.organizationId })
-			.from(ssoProvider)
-			.where(inArray(ssoProvider.providerId, linkedSsoProviderIds));
-
-		if (linkedSsoProviders.length === 0) {
-			return null;
-		}
-
-		providerOrganizationIds = [...new Set(linkedSsoProviders.map((provider) => provider.organizationId))];
-	}
-
-	if (providerOrganizationIds.length === 0) {
 		return null;
 	}
 
+	const provider = await db.query.ssoProvider.findFirst({ where: { providerId } });
+	return provider;
+}
+
+async function tryCreateInvitedMembership(userId: string, email: string, ctx: GenericEndpointContext | null) {
+	logger.debug("Checking for pending invitations for user", userId, email);
+
+	const ssoProvider = await getProviderFromContext(ctx);
+	if (!ssoProvider) {
+		logger.debug("No SSO provider found in context, skipping invitation check");
+		return null;
+	}
+	logger.debug("SSO provider found in context, checking for linked accounts", ssoProvider.providerId);
+
 	const now = new Date();
-	const pendingInvitations = await db
-		.select({
-			id: invitation.id,
-			email: invitation.email,
-			role: invitation.role,
-			organizationId: invitation.organizationId,
-		})
-		.from(invitation)
-		.where(
-			and(
-				eq(invitation.status, "pending"),
-				inArray(invitation.organizationId, providerOrganizationIds),
-				gt(invitation.expiresAt, now),
-			),
-		);
 
-	const matchingInvitation = pendingInvitations.find(
-		(invitationCandidate) => normalizeEmail(invitationCandidate.email) === email,
-	);
+	const pendingInvitation = await db.query.invitation.findFirst({
+		where: {
+			AND: [
+				{ status: "pending" },
+				{ organizationId: ssoProvider.organizationId },
+				{ expiresAt: { gt: now } },
+				{ email: normalizeEmail(email) },
+			],
+		},
+		columns: { id: true, email: true, role: true, organizationId: true },
+	});
 
-	if (!matchingInvitation) {
-		throw new APIError("FORBIDDEN", {
-			message: "SSO sign-in is invite-only for this organization",
-		});
+	if (!pendingInvitation) {
+		logger.debug("No pending invitation found for user", { email });
+		throw new APIError("FORBIDDEN", { message: "SSO sign-in is invite-only for this organization" });
 	}
 
 	db.transaction((tx) => {
@@ -119,16 +75,20 @@ async function tryCreateInvitedMembership(userId: string, email: string, ctx: Ge
 			.values({
 				id: Bun.randomUUIDv7(),
 				userId,
-				role: matchingInvitation.role as "member",
-				organizationId: matchingInvitation.organizationId,
+				role: pendingInvitation.role as "member",
+				organizationId: pendingInvitation.organizationId,
 				createdAt: new Date(),
 			})
 			.run();
 
-		tx.update(invitation).set({ status: "accepted" }).where(eq(invitation.id, matchingInvitation.id)).run();
+		tx.update(invitation).set({ status: "accepted" }).where(eq(invitation.id, pendingInvitation.id)).run();
 	});
 
-	const invitedMembership = await findMembershipWithOrganization(userId, matchingInvitation.organizationId);
+	const invitedMembership = await findMembershipWithOrganization(userId, pendingInvitation.organizationId);
+	logger.debug("Created organization membership from invitation", {
+		userId,
+		organizationId: pendingInvitation.organizationId,
+	});
 
 	if (!invitedMembership) {
 		throw new Error("Failed to create invited organization membership");
@@ -137,7 +97,8 @@ async function tryCreateInvitedMembership(userId: string, email: string, ctx: Ge
 	return invitedMembership;
 }
 
-async function createDefaultOrganizationMembership(user: UserRecord) {
+async function createDefaultOrganizationMembership(user: User) {
+	logger.debug("Creating default organization for user", { userId: user.id });
 	const resticPassword = cryptoUtils.generateResticPassword();
 	const metadata = { resticPassword: await cryptoUtils.sealSecret(resticPassword) };
 
@@ -168,16 +129,17 @@ async function createDefaultOrganizationMembership(user: UserRecord) {
 }
 
 export async function createUserDefaultOrg(userId: string, ctx: GenericEndpointContext | null) {
-	const existingMembership = await findMembershipWithOrganization(userId);
-	if (existingMembership) {
-		return existingMembership;
-	}
-
-	const user = await findUserById(userId);
+	const user = await db.query.usersTable.findFirst({ where: { id: userId } });
 	if (!user) {
 		throw new UnauthorizedError("User not found");
 	}
 
+	const existingMembership = await findMembershipWithOrganization(user.id);
+	if (existingMembership) {
+		logger.debug("User already has an organization membership, skipping default org creation", { userId });
+		return existingMembership;
+	}
+
 	const invitedMembership = await tryCreateInvitedMembership(userId, normalizeEmail(user.email), ctx);
 	if (invitedMembership) {
 		return invitedMembership;

app/server/lib/auth/middlewares/require-sso-invitation.ts

@@ -1,8 +1,6 @@
-import { and, eq, gt } from "drizzle-orm";
 import { APIError } from "better-auth/api";
 import type { GenericEndpointContext } from "@better-auth/core";
 import { db } from "~/server/db/db";
-import { invitation, ssoProvider } from "~/server/db/schema";
 import { logger } from "~/server/utils/logger";
 import { extractProviderIdFromContext, normalizeEmail } from "../utils/sso-context";
 
@@ -16,53 +14,35 @@ export function isSsoCallbackRequest(ctx: GenericEndpointContext | null) {
 
 export const requireSsoInvitation = async (userEmail: string, ctx: GenericEndpointContext | null) => {
 	if (!ctx) {
-		throw new APIError("BAD_REQUEST", {
-			message: "Missing SSO context",
-		});
+		throw new APIError("BAD_REQUEST", { message: "Missing SSO context" });
 	}
 
 	const providerId = extractProviderIdFromContext(ctx);
 	if (!providerId) {
-		throw new APIError("BAD_REQUEST", {
-			message: "Missing providerId in context",
-		});
+		throw new APIError("BAD_REQUEST", { message: "Missing providerId in context" });
 	}
 
-	const provider = await db
-		.select({ organizationId: ssoProvider.organizationId })
-		.from(ssoProvider)
-		.where(eq(ssoProvider.providerId, providerId))
-		.limit(1);
-
-	if (provider.length === 0) {
-		throw new APIError("NOT_FOUND", {
-			message: "SSO provider not found",
-		});
+	const provider = await db.query.ssoProvider.findFirst({ where: { providerId } });
+	if (!provider) {
+		throw new APIError("NOT_FOUND", { message: "SSO provider not found" });
 	}
 
 	const normalizedEmail = normalizeEmail(userEmail);
-	const now = new Date();
+	logger.debug("Checking for pending invitations", normalizedEmail, provider.organizationId);
 
-	logger.debug(
-		"Checking for pending invitations for email %s in organization %s",
-		normalizedEmail,
-		provider[0].organizationId,
-	);
+	const pendingInvitation = await db.query.invitation.findFirst({
+		where: {
+			AND: [
+				{ organizationId: provider.organizationId },
+				{ status: "pending" },
+				{ expiresAt: { gt: new Date() } },
+				{ email: normalizedEmail },
+			],
+		},
+		columns: { id: true, email: true },
+	});
 
-	const pendingInvitations = await db
-		.select({ id: invitation.id, email: invitation.email })
-		.from(invitation)
-		.where(
-			and(
-				eq(invitation.organizationId, provider[0].organizationId),
-				eq(invitation.status, "pending"),
-				gt(invitation.expiresAt, now),
-			),
-		);
-
-	const pendingInvitation = pendingInvitations.find(
-		(invitationCandidate) => normalizeEmail(invitationCandidate.email) === normalizedEmail,
-	);
+	logger.debug("Pending invitation result", { pendingInvitation });
 
 	if (!pendingInvitation) {
 		throw new APIError("FORBIDDEN", {

app/server/lib/auth/middlewares/trust-sso-provider-for-linking.ts

@@ -23,8 +23,8 @@ export async function trustSsoProviderForLinking(ctx: GenericEndpointContext): P
 		return;
 	}
 
-	const provider = await db.query.ssoProvider.findFirst({ where: { providerId } });
-	if (!provider || !provider.autoLinkMatchingEmails) {
+	const provider = await db.query.ssoProvider.findFirst({ where: { providerId, autoLinkMatchingEmails: true } });
+	if (!provider) {
 		return;
 	}