refactor: native trusted providers callback usage

app/client/modules/auth/routes/__tests__/login.test.tsx

@@ -1,5 +1,5 @@
-import { describe, expect, mock, test } from "bun:test";
-import { render, screen } from "@testing-library/react";
+import { afterEach, describe, expect, mock, test } from "bun:test";
+import { cleanup, render, screen } from "@testing-library/react";
 import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
 
 await mock.module("@tanstack/react-router", () => ({
@@ -28,10 +28,25 @@ await mock.module("~/client/lib/auth-client", () => ({
 
 import { LoginPage } from "../login";
 
-const createTestQueryClient = () => new QueryClient({ defaultOptions: { queries: { retry: false } } });
+const createTestQueryClient = () =>
+	new QueryClient({
+		defaultOptions: {
+			queries: {
+				retry: false,
+				gcTime: Infinity,
+			},
+			mutations: {
+				gcTime: Infinity,
+			},
+		},
+	});
 const inviteOnlyMessage =
 	"Access is invite-only. Ask an organization admin to send you an invitation before signing in with SSO.";
 
+afterEach(() => {
+	cleanup();
+});
+
 describe("LoginPage", () => {
 	test("shows an invite-only message when SSO returns INVITE_REQUIRED code", async () => {
 		const queryClient = createTestQueryClient();
@@ -104,7 +119,7 @@ describe("LoginPage", () => {
 			</QueryClientProvider>,
 		);
 
-		// The error message container should be hidden
+		expect(await screen.findByText("Login to your account")).toBeTruthy();
 		expect(screen.queryByText(inviteOnlyMessage)).toBeNull();
 	});
 });

app/server/lib/auth.ts

@@ -22,13 +22,16 @@ import { validateSsoCallbackUrls } from "./auth/middlewares/validate-sso-callbac
 import { validateSsoProviderId } from "./auth/middlewares/validate-sso-provider-id";
 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 { resolveTrustedProvidersForRequest } from "./auth/middlewares/trust-sso-provider-for-linking";
 
 export type AuthMiddlewareContext = MiddlewareContext<MiddlewareOptions, AuthContext<BetterAuthOptions>>;
 
 export const auth = betterAuth({
 	secret: await cryptoUtils.deriveSecret("better-auth"),
-	baseURL: config.baseUrl,
+	baseURL: {
+		allowedHosts: [new URL(config.baseUrl).host, ...config.trustedOrigins.map((origin) => new URL(origin).host)],
+		protocol: "auto",
+	},
 	trustedOrigins: config.trustedOrigins,
 	advanced: {
 		cookiePrefix: "zerobyte",
@@ -92,6 +95,7 @@ export const auth = betterAuth({
 	account: {
 		accountLinking: {
 			enabled: true,
+			trustedProviders: resolveTrustedProvidersForRequest,
 		},
 	},
 	user: {
@@ -133,7 +137,6 @@ export const auth = betterAuth({
 				defaultRole: "member",
 			},
 		}),
-		ssoTrustedProviderLinkingPlugin(),
 		twoFactor({
 			backupCodeOptions: {
 				storeBackupCodes: "encrypted",

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

@@ -1,9 +1,8 @@
 import { beforeEach, describe, expect, test } from "bun:test";
-import type { GenericEndpointContext } from "@better-auth/core";
 import { eq } from "drizzle-orm";
 import { db } from "~/server/db/db";
 import { account, member, organization, ssoProvider, usersTable } from "~/server/db/schema";
-import { isSsoCallbackPath, trustSsoProviderForLinking } from "../trust-sso-provider-for-linking";
+import { resolveTrustedProvidersForRequest } from "../trust-sso-provider-for-linking";
 
 function randomId() {
 	return Bun.randomUUIDv7();
@@ -13,38 +12,8 @@ function randomSlug(prefix: string) {
 	return `${prefix}-${Math.random().toString(36).slice(2, 8)}`;
 }
 
-function createMockContext(options: {
-	path: string;
-	method?: string;
-	params?: Record<string, string>;
-	trustedProviders?: string[];
-	enabled?: boolean;
-}): GenericEndpointContext {
-	const { path, method = "GET", params = {}, trustedProviders = [], enabled = true } = options;
-
-	const accountLinking = {
-		enabled,
-		trustedProviders,
-	};
-
-	const context = {
-		options: {
-			account: {
-				accountLinking,
-			},
-		},
-	};
-
-	return {
-		path,
-		body: {},
-		query: {},
-		headers: new Headers(),
-		request: new Request(`http://test.local${path}`, { method }),
-		params,
-		method,
-		context: context as GenericEndpointContext["context"],
-	} as GenericEndpointContext;
+function createRequest(path: string): Request {
+	return new Request(`http://test.local${path}`);
 }
 
 async function createSsoProviderRecord(
@@ -86,18 +55,7 @@ async function createSsoProviderRecord(
 	return { organizationId, userId };
 }
 
-describe("isSsoCallbackPath", () => {
-	test("detects OIDC callback paths", () => {
-		expect(isSsoCallbackPath("/sso/callback/pocket-id")).toBe(true);
-	});
-
-	test("ignores non-callback paths", () => {
-		expect(isSsoCallbackPath("/sso/register")).toBe(false);
-		expect(isSsoCallbackPath("/sign-in/email")).toBe(false);
-	});
-});
-
-describe("trustSsoProviderForLinking", () => {
+describe("resolveTrustedProvidersForRequest", () => {
 	beforeEach(async () => {
 		await db.delete(member);
 		await db.delete(account);
@@ -106,110 +64,61 @@ describe("trustSsoProviderForLinking", () => {
 		await db.delete(usersTable);
 	});
 
-	test("adds callback provider to trusted providers", async () => {
+	test("returns [] when request is missing", async () => {
+		expect(await resolveTrustedProvidersForRequest()).toEqual([]);
+	});
+
+	test("returns [] for non-callback paths", async () => {
+		expect(await resolveTrustedProvidersForRequest(createRequest("/sign-in/email"))).toEqual([]);
+	});
+
+	test("returns [] for unknown providers", async () => {
+		expect(await resolveTrustedProvidersForRequest(createRequest("/sso/callback/missing-provider"))).toEqual([]);
+	});
+
+	test("returns auto-link-enabled providers from the callback provider organization", async () => {
+		const { organizationId, userId } = await createSsoProviderRecord("pocket-id", true);
+
+		await createSsoProviderRecord("acme-saml", true, { organizationId, userId });
+		await createSsoProviderRecord("acme-disabled", false, { organizationId, userId });
+		await createSsoProviderRecord("other-org-provider", true);
+
+		const trustedProviders = await resolveTrustedProvidersForRequest(createRequest("/sso/callback/pocket-id"));
+
+		expect([...trustedProviders].sort()).toEqual(["acme-saml", "pocket-id"]);
+	});
+
+	test("supports /sso/saml2/callback/:providerId paths", async () => {
+		await createSsoProviderRecord("saml-provider", true);
+
+		expect(await resolveTrustedProvidersForRequest(createRequest("/sso/saml2/callback/saml-provider"))).toEqual([
+			"saml-provider",
+		]);
+	});
+
+	test("supports callback paths nested under /api/auth", async () => {
+		await createSsoProviderRecord("prefixed-provider", true);
+
+		expect(await resolveTrustedProvidersForRequest(createRequest("/api/auth/sso/callback/prefixed-provider"))).toEqual([
+			"prefixed-provider",
+		]);
+	});
+
+	test("supports /sso/saml2/sp/acs/:providerId paths", async () => {
+		await createSsoProviderRecord("saml-acs-provider", true);
+
+		expect(await resolveTrustedProvidersForRequest(createRequest("/sso/saml2/sp/acs/saml-acs-provider"))).toEqual([
+			"saml-acs-provider",
+		]);
+	});
+
+	test("removes providers from the result when auto-linking is disabled", async () => {
 		await createSsoProviderRecord("pocket-id", true);
 
-		const ctx = createMockContext({
-			path: "/sso/callback/pocket-id",
-			params: { providerId: "pocket-id" },
-		});
-
-		await trustSsoProviderForLinking(ctx);
-
-		expect(ctx.context.options.account?.accountLinking?.trustedProviders).toContain("pocket-id");
-	});
-
-	test("does not trust providers with disabled auto-linking", async () => {
-		await createSsoProviderRecord("pocket-id", false);
-
-		const ctx = createMockContext({
-			path: "/sso/callback/pocket-id",
-			params: { providerId: "pocket-id" },
-		});
-
-		await trustSsoProviderForLinking(ctx);
-
-		expect(ctx.context.options.account?.accountLinking?.trustedProviders).toEqual([]);
-	});
-
-	test("replaces stale trusted providers with database state", async () => {
-		await createSsoProviderRecord("pocket-id", true);
-
-		const ctx = createMockContext({
-			path: "/sso/callback/pocket-id",
-			params: { providerId: "pocket-id" },
-			trustedProviders: ["stale-provider"],
-		});
-
-		await trustSsoProviderForLinking(ctx);
-
-		expect(ctx.context.options.account?.accountLinking?.trustedProviders).toEqual(["pocket-id"]);
-	});
-
-	test("does not trust unknown providers", async () => {
-		const ctx = createMockContext({
-			path: "/sso/callback/missing-provider",
-			params: { providerId: "missing-provider" },
-		});
-
-		await trustSsoProviderForLinking(ctx);
-
-		expect(ctx.context.options.account?.accountLinking?.trustedProviders).toEqual([]);
-	});
-
-	test("does not duplicate an existing provider", async () => {
-		await createSsoProviderRecord("pocket-id", true);
-
-		const ctx = createMockContext({
-			path: "/sso/callback/pocket-id",
-			params: { providerId: "pocket-id" },
-			trustedProviders: ["pocket-id"],
-		});
-
-		await trustSsoProviderForLinking(ctx);
-
-		expect(ctx.context.options.account?.accountLinking?.trustedProviders).toEqual(["pocket-id"]);
-	});
-
-	test("removes provider from trusted providers when auto-linking is disabled", async () => {
-		await createSsoProviderRecord("pocket-id", true);
-
-		const ctx = createMockContext({
-			path: "/sso/callback/pocket-id",
-			params: { providerId: "pocket-id" },
-		});
-
-		await trustSsoProviderForLinking(ctx);
-		expect(ctx.context.options.account?.accountLinking?.trustedProviders).toEqual(["pocket-id"]);
+		expect(await resolveTrustedProvidersForRequest(createRequest("/sso/callback/pocket-id"))).toEqual(["pocket-id"]);
 
 		await db.update(ssoProvider).set({ autoLinkMatchingEmails: false }).where(eq(ssoProvider.providerId, "pocket-id"));
 
-		await trustSsoProviderForLinking(ctx);
-		expect(ctx.context.options.account?.accountLinking?.trustedProviders).toEqual([]);
-	});
-
-	test("does nothing when account linking is disabled", async () => {
-		await createSsoProviderRecord("pocket-id", true);
-
-		const ctx = createMockContext({
-			path: "/sso/callback/pocket-id",
-			params: { providerId: "pocket-id" },
-			enabled: false,
-		});
-
-		await trustSsoProviderForLinking(ctx);
-
-		expect(ctx.context.options.account?.accountLinking?.trustedProviders).toEqual([]);
-	});
-
-	test("does nothing when provider id cannot be extracted", async () => {
-		const ctx = createMockContext({
-			path: "/sso/callback/",
-			params: {},
-		});
-
-		await trustSsoProviderForLinking(ctx);
-
-		expect(ctx.context.options.account?.accountLinking?.trustedProviders).toEqual([]);
+		expect(await resolveTrustedProvidersForRequest(createRequest("/sso/callback/pocket-id"))).toEqual([]);
 	});
 });

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

@@ -1,26 +1,14 @@
-import type { GenericEndpointContext } from "@better-auth/core";
 import { db } from "~/server/db/db";
-import { extractProviderIdFromContext } from "../utils/sso-context";
+import { extractProviderIdFromUrl } from "../utils/sso-context";
 
-export function isSsoCallbackPath(path?: string): boolean {
-	if (!path) {
-		return false;
+export async function resolveTrustedProvidersForRequest(request?: Request): Promise<string[]> {
+	if (!request) {
+		return [];
 	}
 
-	return path.startsWith("/sso/callback/");
-}
-
-export async function trustSsoProviderForLinking(ctx: GenericEndpointContext): Promise<void> {
-	const providerId = extractProviderIdFromContext(ctx);
-
+	const providerId = extractProviderIdFromUrl(request.url);
 	if (!providerId) {
-		return;
-	}
-
-	const accountLinking = ctx.context.options.account?.accountLinking;
-
-	if (!accountLinking || accountLinking.enabled === false) {
-		return;
+		return [];
 	}
 
 	const provider = await db.query.ssoProvider.findFirst({
@@ -28,7 +16,7 @@ export async function trustSsoProviderForLinking(ctx: GenericEndpointContext): P
 		where: { providerId },
 	});
 	if (!provider) {
-		return;
+		return [];
 	}
 
 	const autoLinkingProviders = await db.query.ssoProvider.findMany({
@@ -39,5 +27,5 @@ export async function trustSsoProviderForLinking(ctx: GenericEndpointContext): P
 		},
 	});
 
-	accountLinking.trustedProviders = autoLinkingProviders.map((entry) => entry.providerId);
+	return autoLinkingProviders.map((entry) => entry.providerId);
 }

app/server/lib/auth/middlewares/validate-sso-callback-urls.ts

@@ -29,6 +29,7 @@ export const validateSsoCallbackUrls = async (ctx: AuthMiddlewareContext) => {
 			if (value !== undefined && (typeof value !== "string" || !isValidCallbackPath(value))) {
 				throw new APIError("BAD_REQUEST", {
 					message: `Invalid ${field}. Only relative paths like /login are allowed.`,
+					code: `INVALID_${field.toUpperCase()}`,
 				});
 			}
 		}

app/server/modules/auth/__tests__/auth.sso-security.test.ts

@@ -1,9 +1,21 @@
 import { beforeEach, describe, expect, test } from "bun:test";
 import { createApp } from "~/server/app";
+import { config } from "~/server/core/config";
 import { account, invitation, member, organization, ssoProvider, usersTable } from "~/server/db/schema";
 import { db } from "~/server/db/db";
 
 const app = createApp();
+const ssoSignInUrl = new URL("/api/auth/sign-in/sso", config.baseUrl).toString();
+
+function postSsoSignIn(body: Record<string, unknown>) {
+	return app.request(ssoSignInUrl, {
+		method: "POST",
+		headers: {
+			"Content-Type": "application/json",
+		},
+		body: JSON.stringify(body),
+	});
+}
 
 describe("auth SSO sign-in security", () => {
 	beforeEach(async () => {
@@ -16,15 +28,9 @@ describe("auth SSO sign-in security", () => {
 	});
 
 	test("rejects malicious callback URL", async () => {
-		const response = await app.request("/api/auth/sign-in/sso", {
-			method: "POST",
-			headers: {
-				"Content-Type": "application/json",
-			},
-			body: JSON.stringify({
-				providerId: "missing-provider",
-				callbackURL: "https://evil.example",
-			}),
+		const response = await postSsoSignIn({
+			providerId: "missing-provider",
+			callbackURL: "https://evil.example",
 		});
 
 		expect(response.status).toBe(400);
@@ -35,16 +41,10 @@ describe("auth SSO sign-in security", () => {
 	});
 
 	test("rejects malicious error callback URL", async () => {
-		const response = await app.request("/api/auth/sign-in/sso", {
-			method: "POST",
-			headers: {
-				"Content-Type": "application/json",
-			},
-			body: JSON.stringify({
-				providerId: "missing-provider",
-				callbackURL: "/login",
-				errorCallbackURL: "https://evil.example",
-			}),
+		const response = await postSsoSignIn({
+			providerId: "missing-provider",
+			callbackURL: "/login",
+			errorCallbackURL: "https://evil.example",
 		});
 
 		expect(response.status).toBe(400);
@@ -55,16 +55,10 @@ describe("auth SSO sign-in security", () => {
 	});
 
 	test("rejects malicious new user callback URL", async () => {
-		const response = await app.request("/api/auth/sign-in/sso", {
-			method: "POST",
-			headers: {
-				"Content-Type": "application/json",
-			},
-			body: JSON.stringify({
-				providerId: "missing-provider",
-				callbackURL: "/login",
-				newUserCallbackURL: "https://evil.example",
-			}),
+		const response = await postSsoSignIn({
+			providerId: "missing-provider",
+			callbackURL: "/login",
+			newUserCallbackURL: "https://evil.example",
 		});
 
 		expect(response.status).toBe(400);
@@ -75,20 +69,14 @@ describe("auth SSO sign-in security", () => {
 	});
 
 	test("allows relative callback URL to continue normal flow", async () => {
-		const response = await app.request("/api/auth/sign-in/sso", {
-			method: "POST",
-			headers: {
-				"Content-Type": "application/json",
-			},
-			body: JSON.stringify({
-				providerId: "missing-provider",
-				callbackURL: "/login",
-			}),
+		const response = await postSsoSignIn({
+			providerId: "missing-provider",
+			callbackURL: "/login",
 		});
 
 		expect(response.status).toBe(404);
 
 		const body = await response.json();
-		expect(body.code).toBe("NO_PROVIDER_FOUND_FOR_THE_ISSUER");
+		expect(body.message).toBe("No provider found for the issuer");
 	});
 });

playwright/dex-config.yaml

@@ -14,15 +14,19 @@ staticPasswords:
   - email: "admin@example.com"
     hash: "$2a$10$2b2cU8CPhOTaGrs1HRQuAueS7JTT5ZHsHSzYiFPm1leZck7Mc8T4W"
     username: "admin"
+    userID: "001"
   - email: "user@example.com"
     hash: "$2a$10$2b2cU8CPhOTaGrs1HRQuAueS7JTT5ZHsHSzYiFPm1leZck7Mc8T4W"
     username: "user"
+    userID: "002"
   - email: "test@example.com"
     hash: "$2a$10$2b2cU8CPhOTaGrs1HRQuAueS7JTT5ZHsHSzYiFPm1leZck7Mc8T4W"
     username: "test"
+    userID: "003"
   - email: "test@test.com"
     hash: "$2a$10$2b2cU8CPhOTaGrs1HRQuAueS7JTT5ZHsHSzYiFPm1leZck7Mc8T4W"
     username: "test-local"
+    userID: "004"
 
 staticClients:
   - id: zerobyte-test