From 9ffa39141695d62c7efa8f59e7356b0e5c7e6ca8 Mon Sep 17 00:00:00 2001
From: miloschwartz <mschwartz10612@gmail.com>
Date: Sat, 20 Dec 2025 12:00:58 -0500
Subject: [PATCH] improve clean redirects

---
 src/app/auth/login/page.tsx                |  1 +
 src/components/private/IdpLoginButtons.tsx |  4 +-
 src/lib/cleanRedirect.ts                   | 96 ++++++++++++++++++----
 3 files changed, 84 insertions(+), 17 deletions(-)

diff --git a/src/app/auth/login/page.tsx b/src/app/auth/login/page.tsx
index 7ef77807..bd6327fd 100644
--- a/src/app/auth/login/page.tsx
+++ b/src/app/auth/login/page.tsx
@@ -66,6 +66,7 @@ export default async function Page(props: {
     let redirectUrl: string | undefined = undefined;
     if (searchParams.redirect) {
         redirectUrl = cleanRedirect(searchParams.redirect as string);
+        searchParams.redirect = redirectUrl;
     }
 
     let loginIdps: LoginFormIDP[] = [];
diff --git a/src/components/private/IdpLoginButtons.tsx b/src/components/private/IdpLoginButtons.tsx
index e3af2d06..b855683a 100644
--- a/src/components/private/IdpLoginButtons.tsx
+++ b/src/components/private/IdpLoginButtons.tsx
@@ -15,6 +15,7 @@ import {
     useSearchParams
 } from "next/navigation";
 import { useRouter } from "next/navigation";
+import { cleanRedirect } from "@app/lib/cleanRedirect";
 
 export type LoginFormIDP = {
     idpId: number;
@@ -63,9 +64,10 @@ export default function IdpLoginButtons({
                 redirect || "/",
                 orgId
             );
+            const safeRedirect = cleanRedirect(redirect || "/");
             const response = await generateOidcUrlProxy(
                 idpId,
-                redirect || "/",
+                safeRedirect,
                 orgId
             );
 
diff --git a/src/lib/cleanRedirect.ts b/src/lib/cleanRedirect.ts
index 048b3bdc..02a8dde1 100644
--- a/src/lib/cleanRedirect.ts
+++ b/src/lib/cleanRedirect.ts
@@ -1,22 +1,86 @@
-type PatternConfig = {
-    name: string;
-    regex: RegExp;
+type CleanRedirectOptions = {
+    fallback?: string;
+    maxRedirectDepth?: number;
 };
 
-const patterns: PatternConfig[] = [
-    { name: "Invite Token", regex: /^\/invite\?token=[a-zA-Z0-9-]+$/ },
-    { name: "Setup", regex: /^\/setup$/ },
-    { name: "Resource Auth Portal", regex: /^\/auth\/resource\/\d+$/ },
-    {
-        name: "Device Login",
-        regex: /^\/auth\/login\/device(\?code=[a-zA-Z0-9-]+)?$/
-    }
-];
+const ALLOWED_QUERY_PARAMS = new Set([
+    "forceLogin",
+    "code",
+    "token",
+    "redirect"
+]);
+
+const DUMMY_BASE = "https://internal.local";
+
+export function cleanRedirect(
+    input: string,
+    options: CleanRedirectOptions = {}
+): string {
+    const { fallback = "/", maxRedirectDepth = 2 } = options;
 
-export function cleanRedirect(input: string, fallback?: string): string {
     if (!input || typeof input !== "string") {
-        return "/";
+        return fallback;
+    }
+
+    try {
+        return sanitizeUrl(input, fallback, maxRedirectDepth);
+    } catch {
+        return fallback;
     }
-    const isAccepted = patterns.some((pattern) => pattern.regex.test(input));
-    return isAccepted ? input : fallback || "/";
+}
+
+function sanitizeUrl(
+    input: string,
+    fallback: string,
+    remainingRedirectDepth: number
+): string {
+    if (
+        input.startsWith("javascript:") ||
+        input.startsWith("data:") ||
+        input.startsWith("//")
+    ) {
+        return fallback;
+    }
+
+    const url = new URL(input, DUMMY_BASE);
+
+    // Must be a relative/internal path
+    if (url.origin !== DUMMY_BASE) {
+        return fallback;
+    }
+
+    if (!url.pathname.startsWith("/")) {
+        return fallback;
+    }
+
+    const cleanParams = new URLSearchParams();
+
+    for (const [key, value] of url.searchParams.entries()) {
+        if (!ALLOWED_QUERY_PARAMS.has(key)) {
+            continue;
+        }
+
+        if (key === "redirect") {
+            if (remainingRedirectDepth <= 0) {
+                continue;
+            }
+
+            const cleanedRedirect = sanitizeUrl(
+                value,
+                "",
+                remainingRedirectDepth - 1
+            );
+
+            if (cleanedRedirect) {
+                cleanParams.set("redirect", cleanedRedirect);
+            }
+
+            continue;
+        }
+
+        cleanParams.set(key, value);
+    }
+
+    const queryString = cleanParams.toString();
+    return queryString ? `${url.pathname}?${queryString}` : url.pathname;
 }