feat(auth): support IdP-initiated SAML and harden the SSO completion flow

Follow-up to the SP-initiated SAML fix (#405), three changes:

- IdP-initiated SSO: an Okta dashboard-tile launch posts an assertion with no
  RelayState. Rather than rejecting it, the ACS now picks the enabled SAML
  config whose certificate validates the assertion's signature (the stored
  issuer is the SP entity id, not a reliable IdP identifier) and defaults the
  post-login destination to /.

- Enforce 2FA for SAML logins: the completion route now reads force2FAAllLogins
  and stamps the same twoFactorRequired / twoFactorSetupRequired claims the
  NextAuth jwt callback sets at sign-in, so SAML sessions hit the existing
  middleware 2FA gate instead of bypassing it.

- Single-use completion token: the token handed to /api/auth/saml/complete
  rides in a redirect URL, so it now registers a single-use jti in Valkey
  (deleted on consume) and its lifetime is cut to two minutes — a replayed link
  is rejected.

Author: therealbrad

testplanit/app/api/auth/callback/saml/route.test.ts

@@ -13,7 +13,7 @@ vi.mock("~/lib/valkey", () => ({ default: null })); // RelayState via signed tok
 
 vi.mock("~/server/db", () => ({
   db: {
-    samlConfiguration: { findUnique: vi.fn() },
+    samlConfiguration: { findUnique: vi.fn(), findMany: vi.fn() },
     user: { findUnique: vi.fn(), update: vi.fn(), create: vi.fn() },
     account: { upsert: vi.fn() },
     roles: { findFirst: vi.fn() },
@@ -59,9 +59,10 @@ function relayFor(providerId: string, callbackUrl = "/dash") {
 describe("POST /api/auth/callback/saml — ACS validator", () => {
   beforeEach(() => {
     vi.clearAllMocks();
+    (db.samlConfiguration.findMany as any).mockResolvedValue([]);
   });
 
-  it("returns 400 when RelayState is missing or invalid (Bug 4 regression)", async () => {
+  it("returns 400 when there is no RelayState and no enabled config validates it", async () => {
     const res = await POST(makeReq(null));
     expect(res.status).toBe(400);
     const body = await res.json();
@@ -115,4 +116,46 @@ describe("POST /api/auth/callback/saml — ACS validator", () => {
     const res = await POST(makeReq(relayFor("nope")));
     expect(res.status).toBe(404);
   });
+
+  it("supports IdP-initiated login (no RelayState) via the config that validates", async () => {
+    (db.samlConfiguration.findMany as any).mockResolvedValue([
+      {
+        id: "cfg",
+        entryPoint: "e",
+        cert: "c",
+        issuer: "i",
+        attributeMapping: {},
+        autoProvisionUsers: false,
+        provider: { name: "okta", enabled: true },
+      },
+    ]);
+    validateSAMLResponse.mockResolvedValue({
+      email: "carol@example.com",
+      nameID: "carol",
+    });
+    (db.user.findUnique as any).mockResolvedValue({
+      id: "user_8",
+      email: "carol@example.com",
+      name: "Carol",
+      authMethod: "SSO",
+      externalId: "carol",
+    });
+    (db.account.upsert as any).mockResolvedValue({});
+
+    const res = await POST(makeReq(null)); // no RelayState = IdP-initiated
+
+    // Picked the config by validating the assertion (findMany over enabled
+    // providers), not by a RelayState-carried id.
+    expect(db.samlConfiguration.findMany).toHaveBeenCalled();
+    expect(db.samlConfiguration.findUnique).not.toHaveBeenCalled();
+
+    const location = res.headers.get("location")!;
+    expect(
+      location.startsWith(
+        "https://app.example.com/api/auth/saml/complete?token="
+      )
+    ).toBe(true);
+    // Post-login destination defaults to /.
+    expect(location).toContain("callbackUrl=%2F");
+  });
 });

testplanit/app/api/auth/callback/saml/route.ts

@@ -14,6 +14,43 @@ import { isEmailDomainAllowed } from "~/lib/utils/email-domain-validation";
 import { db } from "~/server/db";
 import { createSAMLClient, validateSAMLResponse } from "~/server/saml-provider";
 
+/**
+ * Resolve an IdP-initiated SAMLResponse (no RelayState — e.g. an Okta dashboard
+ * tile) to the enabled SAML config whose IdP signed it. The stored issuer field
+ * is the SP entity id, not a reliable IdP identifier, so we let cryptographic
+ * validation pick the config: the assertion only validates against the cert of
+ * the IdP that actually issued it. Returns null if none validate it.
+ */
+async function resolveIdpInitiatedConfig(samlResponse: FormDataEntryValue) {
+  const configs = await db.samlConfiguration.findMany({
+    where: { provider: { enabled: true } },
+    include: { provider: true },
+  });
+
+  for (const samlConfig of configs) {
+    try {
+      const samlClient = await createSAMLClient({
+        name: samlConfig.provider.name,
+        entryPoint: samlConfig.entryPoint,
+        cert: samlConfig.cert,
+        issuer: samlConfig.issuer,
+      });
+      const profile = await validateSAMLResponse(samlClient, {
+        SAMLResponse: samlResponse,
+      });
+      return { samlConfig, profile };
+    } catch {
+      // Not this IdP (or the assertion is invalid) — try the next config.
+    }
+  }
+
+  return null;
+}
+
+type ResolvedSaml = NonNullable<
+  Awaited<ReturnType<typeof resolveIdpInitiatedConfig>>
+>;
+
 /**
  * SAML Assertion Consumer Service (ACS).
  *
@@ -54,48 +91,57 @@ export async function POST(request: NextRequest) {
       );
     }
 
-    // Recover the provider and destination from RelayState (the IdP echoes it
-    // back here; same-site cookies are not sent on this cross-site POST). The
-    // token is single-use and short-lived, which is what guards this endpoint.
+    // Recover the provider and destination from RelayState. The IdP echoes it
+    // back on this cross-site POST (where same-site cookies are not sent); the
+    // token is single-use and short-lived, which guards SP-initiated logins.
     const relay = await consumeSamlRelayState(
       typeof relayState === "string" ? relayState : null
     );
 
-    if (!relay) {
-      return NextResponse.json(
-        { error: "Invalid or expired SAML request" },
-        { status: 400 }
-      );
-    }
-
-    const providerId = relay.providerId;
-    const callbackUrl = sanitizeCallbackUrl(relay.callbackUrl);
+    let samlConfig: ResolvedSaml["samlConfig"];
+    let profile: ResolvedSaml["profile"];
+    let callbackUrl = "/";
 
-    // Fetch SAML configuration. RelayState carries the SsoProvider id, which is
-    // the unique foreign key on SamlConfiguration (not its own id).
-    const samlConfig = await db.samlConfiguration.findUnique({
-      where: { providerId },
-      include: { provider: true },
-    });
+    if (relay) {
+      // SP-initiated: RelayState carries the SsoProvider id, which is the unique
+      // foreign key on SamlConfiguration (not its own id).
+      callbackUrl = sanitizeCallbackUrl(relay.callbackUrl);
+      const found = await db.samlConfiguration.findUnique({
+        where: { providerId: relay.providerId },
+        include: { provider: true },
+      });
 
-    if (!samlConfig || !samlConfig.provider.enabled) {
-      return NextResponse.json(
-        { error: "SAML provider not found or disabled" },
-        { status: 404 }
-      );
-    }
+      if (!found || !found.provider.enabled) {
+        return NextResponse.json(
+          { error: "SAML provider not found or disabled" },
+          { status: 404 }
+        );
+      }
 
-    // Create SAML client and validate response
-    const samlClient = await createSAMLClient({
-      name: samlConfig.provider.name,
-      entryPoint: samlConfig.entryPoint,
-      cert: samlConfig.cert,
-      issuer: samlConfig.issuer,
-    });
+      const samlClient = await createSAMLClient({
+        name: found.provider.name,
+        entryPoint: found.entryPoint,
+        cert: found.cert,
+        issuer: found.issuer,
+      });
 
-    const profile = await validateSAMLResponse(samlClient, {
-      SAMLResponse: samlResponse,
-    });
+      samlConfig = found;
+      profile = await validateSAMLResponse(samlClient, {
+        SAMLResponse: samlResponse,
+      });
+    } else {
+      // IdP-initiated (no RelayState — e.g. an Okta dashboard tile): pick the
+      // enabled config whose cert validates this assertion's signature.
+      const resolved = await resolveIdpInitiatedConfig(samlResponse);
+      if (!resolved) {
+        return NextResponse.json(
+          { error: "Invalid or expired SAML request" },
+          { status: 400 }
+        );
+      }
+      samlConfig = resolved.samlConfig;
+      profile = resolved.profile;
+    }
 
     // Validate SAML response timestamps if available
     if (profile.notBefore || profile.notOnOrAfter) {
@@ -269,8 +315,8 @@ export async function POST(request: NextRequest) {
       },
     });
 
-    // Create a temporary session token for secure user info transfer
-    const tempToken = createTempSessionToken({
+    // Create a one-time session token for secure user info transfer.
+    const tempToken = await createTempSessionToken({
       userId: user.id,
       provider: `saml-${samlConfig.provider.name}`,
       email: user.email,

testplanit/app/api/auth/saml/complete/route.test.ts

@@ -21,7 +21,10 @@ vi.mock("next/headers", () => ({
 }));
 
 vi.mock("~/server/db", () => ({
-  db: { user: { findUnique: vi.fn() } },
+  db: {
+    user: { findUnique: vi.fn() },
+    registrationSettings: { findFirst: vi.fn() },
+  },
 }));
 
 import { db } from "~/server/db";
@@ -35,6 +38,7 @@ const user = {
   isApi: false,
   passwordChangedAt: null,
   mustChangePassword: false,
+  twoFactorEnabled: false,
 };
 
 function makeReq(token: string, callbackUrl = "/dashboard") {
@@ -58,6 +62,7 @@ describe("GET /api/auth/saml/complete — post-Okta session handoff", () => {
     cookieJar.clear();
     vi.clearAllMocks();
     (db.user.findUnique as any).mockResolvedValue(user);
+    (db.registrationSettings.findFirst as any).mockResolvedValue(null);
   });
 
   afterEach(() => {
@@ -99,6 +104,40 @@ describe("GET /api/auth/saml/complete — post-Okta session handoff", () => {
     ).toBeTruthy();
   });
 
+  it("stamps 2FA-required when force2FAAllLogins is on and the user has 2FA", async () => {
+    (db.registrationSettings.findFirst as any).mockResolvedValue({
+      force2FAAllLogins: true,
+    });
+    (db.user.findUnique as any).mockResolvedValue({
+      ...user,
+      twoFactorEnabled: true,
+    });
+
+    await GET(makeReq(tempToken()));
+
+    const decoded = await decode({
+      token: cookieJar.get("next-auth.session-token")!.value,
+      secret: SECRET,
+    });
+    expect(decoded?.twoFactorRequired).toBe(true);
+    expect(decoded?.twoFactorVerified).toBe(false);
+  });
+
+  it("stamps 2FA-setup-required when force2FA is on but the user has no 2FA", async () => {
+    (db.registrationSettings.findFirst as any).mockResolvedValue({
+      force2FAAllLogins: true,
+    });
+    // user.twoFactorEnabled is false by default
+
+    await GET(makeReq(tempToken()));
+
+    const decoded = await decode({
+      token: cookieJar.get("next-auth.session-token")!.value,
+      secret: SECRET,
+    });
+    expect(decoded?.twoFactorSetupRequired).toBe(true);
+  });
+
   it("rejects an invalid/expired token with 401", async () => {
     const res = await GET(makeReq("not-a-valid-jwt"));
     expect(res.status).toBe(401);

testplanit/app/api/auth/saml/complete/route.ts

@@ -1,8 +1,7 @@
-import jwt from "jsonwebtoken";
 import { encode } from "next-auth/jwt";
 import { cookies } from "next/headers";
 import { NextRequest, NextResponse } from "next/server";
-import { getAppBaseUrl } from "~/lib/auth-security";
+import { consumeTempSessionToken, getAppBaseUrl } from "~/lib/auth-security";
 import { db } from "~/server/db";
 
 // SAML completion handler - creates NextAuth session
@@ -16,14 +15,10 @@ export async function GET(request: NextRequest) {
       return NextResponse.json({ error: "Token is required" }, { status: 400 });
     }
 
-    // Verify the JWT token
-    let tokenData: any;
-    try {
-      tokenData = jwt.verify(
-        token,
-        process.env.NEXTAUTH_SECRET || "development-secret"
-      );
-    } catch {
+    // Verify and consume the one-time token (single-use via Valkey when
+    // present, so the token in the redirect URL can't be replayed).
+    const tokenData = await consumeTempSessionToken(token);
+    if (!tokenData) {
       return NextResponse.json(
         { error: "Invalid or expired token" },
         { status: 401 }
@@ -41,13 +36,34 @@ export async function GET(request: NextRequest) {
         isApi: true,
         passwordChangedAt: true,
         mustChangePassword: true,
+        twoFactorEnabled: true,
       },
     });
 
     if (!user) {
       return NextResponse.json({ error: "User not found" }, { status: 404 });
     }
 
+    // Mirror NextAuth's force-2FA-for-SSO gate (the jwt callback in server/auth):
+    // this flow mints the session directly, so without stamping the same flags
+    // here a SAML login would bypass an enforced 2FA policy.
+    const twoFactorClaims: {
+      twoFactorRequired?: boolean;
+      twoFactorVerified?: boolean;
+      twoFactorSetupRequired?: boolean;
+    } = {};
+    const registrationSettings = await db.registrationSettings.findFirst({
+      select: { force2FAAllLogins: true },
+    });
+    if (registrationSettings?.force2FAAllLogins) {
+      if (user.twoFactorEnabled) {
+        twoFactorClaims.twoFactorRequired = true;
+        twoFactorClaims.twoFactorVerified = false;
+      } else {
+        twoFactorClaims.twoFactorSetupRequired = true;
+      }
+    }
+
     // Create NextAuth JWT session token. Seed the same fields the NextAuth `jwt`
     // callback sets at sign-in so middleware (which only decodes the token, and
     // never runs the callback) sees the user's access on the very first request.
@@ -61,6 +77,7 @@ export async function GET(request: NextRequest) {
         isApi: user.isApi,
         passwordChangedAt: user.passwordChangedAt?.toISOString() ?? null,
         mustChangePassword: user.mustChangePassword ?? false,
+        ...twoFactorClaims,
       },
       secret: process.env.NEXTAUTH_SECRET || "development-secret",
     });