Merge branch 'main' into blog/litellm-alternatives

Author: smakosh

apps/api/src/auth/config.ts

@@ -33,19 +33,44 @@ const isHosted = process.env.HOSTED === "true";
 async function isSSOEnforcedForEmail(
 	email: string | null | undefined,
 ): Promise<boolean> {
+	// Keep enforcement tied to the same flag that surfaces the SSO sign-in entry
+	// point: with SSO_ENABLED off there's no way to sign in via SSO, so blocking
+	// password/social/passkey would just lock the domain out entirely.
+	if (process.env.SSO_ENABLED !== "true") {
+		return false;
+	}
+
 	const domain = email?.trim().toLowerCase().split("@")[1];
 	if (!domain) {
 		return false;
 	}
 	const providers = await db.query.ssoProvider.findMany({
 		where: { enforced: { eq: true } },
-		columns: { domain: true },
+		columns: { domain: true, organizationId: true },
+	});
+	const matching = providers.filter(
+		(provider) =>
+			provider.organizationId &&
+			provider.domain
+				.split(",")
+				.map((d) => d.trim().toLowerCase())
+				.includes(domain),
+	);
+	if (!matching.length) {
+		return false;
+	}
+
+	// Only an active, enterprise org can enforce SSO. A soft-deleted org can no
+	// longer be managed to turn enforcement off, so a stale enforced row must not
+	// keep blocking a domain's users forever.
+	const orgs = await db.query.organization.findMany({
+		where: {
+			id: { in: [...new Set(matching.map((p) => p.organizationId as string))] },
+		},
+		columns: { status: true, plan: true },
 	});
-	return providers.some((provider) =>
-		provider.domain
-			.split(",")
-			.map((d) => d.trim().toLowerCase())
-			.includes(domain),
+	return orgs.some(
+		(org) => org.status !== "deleted" && org.plan === "enterprise",
 	);
 }
 
@@ -1060,9 +1085,12 @@ The LLM Gateway Team`.trim();
 						!ctx.path.startsWith("/sso/") &&
 						(await isSSOEnforcedForEmail(dbUser?.email))
 					) {
+						// Delete only the session this blocked attempt just created — not
+						// every session for the user — so a rejected social/passkey login
+						// doesn't also sign them out of valid SSO sessions elsewhere.
 						await db
 							.delete(tables.session)
-							.where(eq(tables.session.userId, userId));
+							.where(eq(tables.session.id, newSession.session.id));
 						return ssoRequiredResponse();
 					}
 

apps/api/src/auth/handler.spec.ts

@@ -0,0 +1,88 @@
+import { describe, expect, test } from "vitest";
+
+import { app } from "@/index.js";
+
+// Better Auth's sso() plugin mounts its own provider-management endpoints under
+// the /auth/* catch-all. They only require a session, so `handler.ts` blocks
+// direct HTTP access to them (see BLOCKED_SSO_PLUGIN_PATHS) — our own
+// enterprise/admin-gated wrapper lives at /sso and calls the plugin via
+// `apiAuth.api.*`, which bypasses this handler entirely.
+const BLOCKED_PATHS = [
+	"/auth/sso/register",
+	"/auth/sso/update-provider",
+	"/auth/sso/delete-provider",
+	"/auth/sso/get-provider",
+	"/auth/sso/providers",
+	"/auth/sso/request-domain-verification",
+	"/auth/sso/verify-domain",
+];
+
+// Paths the login flow needs must NOT be blocked by us; they pass through to
+// Better Auth's handler.
+const ALLOWED_PATHS = [
+	"/auth/sign-in/sso",
+	"/auth/sso/callback",
+	"/auth/sso/saml2/sp/metadata",
+];
+
+const BLOCK_BODY = { error: "not_found", message: "Not found" };
+
+// Our block responds with exactly this 404 JSON, distinct from any 404 Better
+// Auth itself might return, so allowed paths can be told apart from blocked.
+function isOurBlock(status: number, body: unknown): boolean {
+	return (
+		status === 404 &&
+		!!body &&
+		typeof body === "object" &&
+		(body as { error?: string }).error === "not_found" &&
+		(body as { message?: string }).message === "Not found"
+	);
+}
+
+describe("SSO plugin management endpoint blocking", () => {
+	test.each(BLOCKED_PATHS)("blocks POST %s with 404", async (path) => {
+		const response = await app.request(path, {
+			method: "POST",
+			headers: { "Content-Type": "application/json" },
+			body: "{}",
+		});
+
+		expect(response.status).toBe(404);
+		expect(await response.json()).toEqual(BLOCK_BODY);
+	});
+
+	test.each(BLOCKED_PATHS)("blocks GET %s with 404", async (path) => {
+		const response = await app.request(path, { method: "GET" });
+
+		expect(response.status).toBe(404);
+		expect(await response.json()).toEqual(BLOCK_BODY);
+	});
+
+	test("blocks even with a query string", async () => {
+		const response = await app.request(
+			"/auth/sso/get-provider?providerId=acme",
+			{ method: "GET" },
+		);
+
+		expect(response.status).toBe(404);
+		expect(await response.json()).toEqual(BLOCK_BODY);
+	});
+
+	test("does not block an unrelated /auth path", async () => {
+		const response = await app.request("/auth/ok", { method: "GET" });
+		const body = await response.json().catch(() => null);
+
+		expect(isOurBlock(response.status, body)).toBe(false);
+	});
+
+	test.each(ALLOWED_PATHS)(
+		"does not block the login-flow path %s",
+		async (path) => {
+			const response = await app.request(path, { method: "GET" });
+			const body = await response.json().catch(() => null);
+
+			// Reached Better Auth (any status/body) rather than our 404 sentinel.
+			expect(isOurBlock(response.status, body)).toBe(false);
+		},
+	);
+});

apps/api/src/auth/handler.ts

@@ -7,6 +7,25 @@ import type { ServerTypes } from "@/vars.js";
 // Create a Hono app for auth routes
 export const authHandler = new OpenAPIHono<ServerTypes>();
 
+// Better Auth's sso() plugin mounts its own provider-management endpoints under
+// the /auth/* catch-all. They only require a session (any authenticated user),
+// so exposing them would let a user bypass the enterprise/admin-gated /sso
+// wrapper and register a provider — claiming any email domain — with no
+// organization scoping. Block direct HTTP access to them here; our server-side
+// registration calls `apiAuth.api.registerSSOProvider` directly and never goes
+// through this handler, so it is unaffected. The SAML sign-in/callback/metadata
+// endpoints (/auth/sso/saml2/*, /auth/sso/callback, /auth/sign-in/sso) are left
+// reachable because the login flow needs them.
+const BLOCKED_SSO_PLUGIN_PATHS = new Set([
+	"/auth/sso/register",
+	"/auth/sso/update-provider",
+	"/auth/sso/delete-provider",
+	"/auth/sso/get-provider",
+	"/auth/sso/providers",
+	"/auth/sso/request-domain-verification",
+	"/auth/sso/verify-domain",
+]);
+
 authHandler.use("*", async (c, next) => {
 	const session = await apiAuth.api.getSession({ headers: c.req.raw.headers });
 
@@ -22,5 +41,8 @@ authHandler.use("*", async (c, next) => {
 });
 
 authHandler.on(["POST", "GET"], "/auth/*", (c) => {
+	if (BLOCKED_SSO_PLUGIN_PATHS.has(new URL(c.req.url).pathname)) {
+		return c.json({ error: "not_found", message: "Not found" }, 404);
+	}
 	return apiAuth.handler(c.req.raw);
 });

apps/api/src/lib/sso-roles.ts

@@ -0,0 +1,111 @@
+import { db, eq, tables } from "@llmgateway/db";
+
+export type OrgRole = "owner" | "admin" | "developer";
+
+export interface RoleChange {
+	old: OrgRole;
+	new: OrgRole;
+}
+
+const ROLE_RANK: Record<OrgRole, number> = {
+	developer: 1,
+	admin: 2,
+	owner: 3,
+};
+
+// Recompute an org member's role from their SCIM group memberships and the
+// org's group->role mappings. The highest-precedence mapped role wins; the
+// default is `developer`. Owners are never auto-demoted — owner is only ever
+// assigned manually (or by an explicit owner mapping), so an admin who set up
+// SSO can't be locked out by a group that maps to a lower role.
+//
+// Returns the {old,new} role change when it updated the membership, or null
+// when nothing changed, so callers (e.g. SCIM) can audit the transition.
+export async function recomputeUserRole(
+	userId: string,
+	organizationId: string,
+): Promise<RoleChange | null> {
+	const membership = await db.query.userOrganization.findFirst({
+		where: {
+			userId: { eq: userId },
+			organizationId: { eq: organizationId },
+		},
+		columns: { id: true, role: true },
+	});
+	if (!membership) {
+		return null;
+	}
+
+	const groupMemberships = await db.query.scimGroupMember.findMany({
+		where: { userId: { eq: userId } },
+		columns: { scimGroupId: true },
+	});
+	const groupIds = groupMemberships.map((m) => m.scimGroupId);
+
+	let mappedRole: OrgRole = "developer";
+	if (groupIds.length) {
+		const groups = await db.query.scimGroup.findMany({
+			where: {
+				id: { in: groupIds },
+				organizationId: { eq: organizationId },
+			},
+			columns: { displayName: true },
+		});
+		const names = groups.map((g) => g.displayName);
+		if (names.length) {
+			const mappings = await db.query.ssoRoleMapping.findMany({
+				where: {
+					organizationId: { eq: organizationId },
+					groupName: { in: names },
+				},
+				columns: { role: true },
+			});
+			for (const mapping of mappings) {
+				if (ROLE_RANK[mapping.role] > ROLE_RANK[mappedRole]) {
+					mappedRole = mapping.role;
+				}
+			}
+		}
+	}
+
+	if (membership.role === "owner" && ROLE_RANK[mappedRole] < ROLE_RANK.owner) {
+		return null;
+	}
+	if (membership.role !== mappedRole) {
+		await db
+			.update(tables.userOrganization)
+			.set({ role: mappedRole })
+			.where(eq(tables.userOrganization.id, membership.id));
+		return { old: membership.role, new: mappedRole };
+	}
+	return null;
+}
+
+// Recompute the role of every member of the SCIM group(s) with `groupName` in
+// this org. Used when a role mapping is added or removed after the IdP has
+// already pushed the group and its members, so existing members pick up (or
+// lose) the mapped role without waiting for a later SCIM membership event.
+export async function recomputeRoleForGroupName(
+	organizationId: string,
+	groupName: string,
+) {
+	const groups = await db.query.scimGroup.findMany({
+		where: {
+			organizationId: { eq: organizationId },
+			displayName: { eq: groupName },
+		},
+		columns: { id: true },
+	});
+	if (!groups.length) {
+		return;
+	}
+
+	const members = await db.query.scimGroupMember.findMany({
+		where: { scimGroupId: { in: groups.map((g) => g.id) } },
+		columns: { userId: true },
+	});
+	const userIds = [...new Set(members.map((m) => m.userId))];
+	for (const userId of userIds) {
+		await recomputeUserRole(userId, organizationId);
+	}
+}

apps/api/src/routes/keys-api.spec.ts

@@ -417,6 +417,79 @@ describe("keys route", () => {
 		expect(apiKey?.periodUsageDurationUnit).toBe("day");
 	});
 
+	test("POST /keys/api rejects a limit above the member budget", async () => {
+		await db
+			.update(tables.userOrganization)
+			.set({ usageLimit: "10" })
+			.where(eq(tables.userOrganization.id, "test-user-org-id"));
+
+		const res = await app.request("/keys/api", {
+			method: "POST",
+			headers: {
+				"Content-Type": "application/json",
+				Cookie: token,
+			},
+			body: JSON.stringify({
+				description: "Over-budget key",
+				projectId: "test-project-id",
+				usageLimit: "50",
+			}),
+		});
+
+		expect(res.status).toBe(400);
+		const json = await res.json();
+		expect(json.message).toMatch(/organization limit of \$10\.00/);
+
+		const apiKey = await db.query.apiKey.findFirst({
+			where: { description: { eq: "Over-budget key" } },
+		});
+		expect(apiKey).toBeUndefined();
+	});
+
+	test("POST /keys/api allows a limit at or below the member budget", async () => {
+		await db
+			.update(tables.userOrganization)
+			.set({ usageLimit: "10" })
+			.where(eq(tables.userOrganization.id, "test-user-org-id"));
+
+		const res = await app.request("/keys/api", {
+			method: "POST",
+			headers: {
+				"Content-Type": "application/json",
+				Cookie: token,
+			},
+			body: JSON.stringify({
+				description: "Within-budget key",
+				projectId: "test-project-id",
+				usageLimit: "10",
+			}),
+		});
+
+		expect(res.status).toBe(200);
+		const json = await res.json();
+		expect(json.apiKey.usageLimit).toBe("10");
+	});
+
+	test("PATCH /keys/api/limit/{id} rejects a limit above the member budget", async () => {
+		await db
+			.update(tables.userOrganization)
+			.set({ usageLimit: "10" })
+			.where(eq(tables.userOrganization.id, "test-user-org-id"));
+
+		const res = await app.request("/keys/api/limit/test-api-key-id", {
+			method: "PATCH",
+			headers: {
+				"Content-Type": "application/json",
+				Cookie: token,
+			},
+			body: JSON.stringify({ usageLimit: "50" }),
+		});
+
+		expect(res.status).toBe(400);
+		const json = await res.json();
+		expect(json.message).toMatch(/organization limit of \$10\.00/);
+	});
+
 	test("PATCH /keys/api/limit/{id} updates and resets period usage", async () => {
 		await db
 			.update(tables.apiKey)