From 3a73a327c325304bc57167f0f009dac14c2ba54d Mon Sep 17 00:00:00 2001
From: Arthur Belleville <arthur.belleville@datadoghq.com>
Date: Sat, 5 Jul 2025 21:44:45 +0200
Subject: [PATCH] Fix login with redirect

---
 .../BrandButtons/LoginWithGoogle.test.tsx     |  4 +--
 .../BrandButtons/LoginWithGoogle.tsx          |  8 ++---
 ui/src/components/PublicRoute.tsx             |  9 +++++-
 ui/src/hooks/auth.ts                          | 29 +++++++------------
 ui/src/pages/login.tsx                        | 16 ++++++----
 ui/src/pages/oauth-signin.tsx                 |  6 ++--
 ui/src/pages/signup.tsx                       | 11 +++----
 7 files changed, 41 insertions(+), 42 deletions(-)

diff --git a/ui/src/components/BrandButtons/LoginWithGoogle.test.tsx b/ui/src/components/BrandButtons/LoginWithGoogle.test.tsx
index 78dfd38..3174855 100644
--- a/ui/src/components/BrandButtons/LoginWithGoogle.test.tsx
+++ b/ui/src/components/BrandButtons/LoginWithGoogle.test.tsx
@@ -14,7 +14,7 @@ describe("LoginWithGoogle", () => {
       loginWithGoogle: mockLoginWithGoogle,
     });
 
-    render(<LoginWithGoogle redirectUrl={null} />);
+    render(<LoginWithGoogle />);
 
     const button = screen.getByRole("button", {
       name: /Continuer avec Google/i,
@@ -28,7 +28,7 @@ describe("LoginWithGoogle", () => {
       loginWithGoogle: mockLoginWithGoogle,
     });
 
-    render(<LoginWithGoogle redirectUrl={null} />);
+    render(<LoginWithGoogle />);
 
     const button = screen.getByRole("button", {
       name: /Continuer avec Google/i,
diff --git a/ui/src/components/BrandButtons/LoginWithGoogle.tsx b/ui/src/components/BrandButtons/LoginWithGoogle.tsx
index 7417dee..0dedcd0 100644
--- a/ui/src/components/BrandButtons/LoginWithGoogle.tsx
+++ b/ui/src/components/BrandButtons/LoginWithGoogle.tsx
@@ -1,12 +1,8 @@
 import "./login-with-google.css";
 import { useLoginGoogle } from "../../hooks/auth";
 
-export function LoginWithGoogle({
-  redirectUrl,
-}: {
-  redirectUrl: string | null;
-}) {
-  const { loginWithGoogle } = useLoginGoogle({ redirectUrl });
+export function LoginWithGoogle() {
+  const { loginWithGoogle } = useLoginGoogle();
 
   return (
     <button className="gsi-material-button" onClick={() => loginWithGoogle()}>
diff --git a/ui/src/components/PublicRoute.tsx b/ui/src/components/PublicRoute.tsx
index 5f16e2e..146333b 100644
--- a/ui/src/components/PublicRoute.tsx
+++ b/ui/src/components/PublicRoute.tsx
@@ -1,12 +1,19 @@
 import { useEffect, useState } from "react";
 import { useSession } from "../contexts/SessionContext";
-import { Navigate, Outlet } from "react-router-dom";
+import { Navigate, Outlet, useSearchParams } from "react-router-dom";
 import { match } from "ts-pattern";
 
 export const PublicRoute = () => {
   const { session } = useSession();
   const [isLoading, setIsLoading] = useState(true);
 
+  const [searchParams] = useSearchParams();
+  const redirectUrl = searchParams.get("redirect");
+
+  if (redirectUrl) {
+    localStorage.setItem("redirectUrl", redirectUrl);
+  }
+
   useEffect(() => {
     const timer = setTimeout(() => {
       setIsLoading(false);
diff --git a/ui/src/hooks/auth.ts b/ui/src/hooks/auth.ts
index 95db513..44cd4d6 100644
--- a/ui/src/hooks/auth.ts
+++ b/ui/src/hooks/auth.ts
@@ -48,7 +48,7 @@ interface AuthResponse {
   session: Session | null;
 }
 
-export function useSignUp() {
+export function useSignUp({ redirectUrl }: { redirectUrl: string | null }) {
   const navigate = useNavigate();
   const [errors, setErrors] = useState<Record<string, string>>({});
   const { signUpToStream } = useSignUpToStream();
@@ -76,7 +76,12 @@ export function useSignUp() {
       return response;
     },
     onSuccess: () => {
-      navigate("/");
+      if (redirectUrl) {
+        localStorage.removeItem("redirectUrl");
+        navigate(decodeURIComponent(redirectUrl));
+      } else {
+        navigate("/");
+      }
     },
     onError: (error) => {
       const errMap: Record<string, string> = {};
@@ -144,6 +149,7 @@ export function useLoginEmail({ redirectUrl }: { redirectUrl: string | null }) {
     },
     onSuccess: () => {
       if (redirectUrl) {
+        localStorage.removeItem("redirectUrl");
         navigate(decodeURIComponent(redirectUrl));
       } else {
         navigate("/");
@@ -172,28 +178,13 @@ export function useLoginEmail({ redirectUrl }: { redirectUrl: string | null }) {
   return { mutate, isPending, errors };
 }
 
-export function useLoginGoogle({
-  redirectUrl,
-}: {
-  redirectUrl: string | null;
-}) {
-  console.log({
-    redirectTo: redirectUrl
-      ? `${
-          window.location.origin
-        }/login-with-oauth?redirect=${encodeURIComponent(redirectUrl)}`
-      : `${window.location.origin}/login-with-oauth`,
-  });
+export function useLoginGoogle() {
   const { mutate } = useMutation({
     mutationFn: async () => {
       const { data, error } = await supabase.auth.signInWithOAuth({
         provider: "google",
         options: {
-          redirectTo: redirectUrl
-            ? `${
-                window.location.origin
-              }/login-with-oauth?redirect=${encodeURIComponent(redirectUrl)}`
-            : `${window.location.origin}/login-with-oauth`,
+          redirectTo: `${window.location.origin}/login-with-oauth`,
         },
       });
       if (error) throw error;
diff --git a/ui/src/pages/login.tsx b/ui/src/pages/login.tsx
index 87a576e..2c3d884 100644
--- a/ui/src/pages/login.tsx
+++ b/ui/src/pages/login.tsx
@@ -5,13 +5,17 @@ import { Label, Input, TextField, FieldError } from "@ui/ui-library/field";
 import { useLoginEmail } from "@ui/hooks/auth";
 import { Form } from "@ui/ui-library/form";
 import { LoginWithGoogle } from "@ui/components/BrandButtons/LoginWithGoogle";
-import { Link, useSearchParams } from "react-router-dom";
+import { Link } from "react-router-dom";
 
 export function LoginPage() {
-  const [searchParams] = useSearchParams();
-  const redirectUrl = searchParams.get("redirect");
-
-  const { mutate: login, isPending, errors } = useLoginEmail({ redirectUrl });
+  const redirectUrl = localStorage.getItem("redirectUrl");
+  const {
+    mutate: login,
+    isPending,
+    errors,
+  } = useLoginEmail({
+    redirectUrl: redirectUrl ?? null,
+  });
   const [formData, setFormData] = useState({
     email: "",
     password: "",
@@ -116,7 +120,7 @@ export function LoginPage() {
             </div>
           </div>
 
-          <LoginWithGoogle redirectUrl={redirectUrl} />
+          <LoginWithGoogle />
 
           <p className="text-center text-sm text-slate-600 dark:text-slate-400">
             Pas encore de compte ?{" "}
diff --git a/ui/src/pages/oauth-signin.tsx b/ui/src/pages/oauth-signin.tsx
index 81dd804..f4794f5 100644
--- a/ui/src/pages/oauth-signin.tsx
+++ b/ui/src/pages/oauth-signin.tsx
@@ -1,5 +1,5 @@
 import { useEffect } from "react";
-import { useNavigate, useSearchParams } from "react-router-dom";
+import { useNavigate } from "react-router-dom";
 import { useSession } from "@ui/contexts/SessionContext";
 import { useSignUpToStream } from "@ui/hooks/auth";
 
@@ -7,13 +7,13 @@ export const OAuthSigninPage = () => {
   const navigate = useNavigate();
   const { session } = useSession();
   const { signUpToStream } = useSignUpToStream();
-  const [searchParams] = useSearchParams();
-  const redirectUrl = searchParams.get("redirect");
+  const redirectUrl = localStorage.getItem("redirectUrl");
   useEffect(() => {
     const interval = setInterval(() => {
       if (session) {
         signUpToStream(session.access_token);
         if (redirectUrl) {
+          localStorage.removeItem("redirectUrl");
           navigate(decodeURIComponent(redirectUrl));
         } else {
           navigate("/");
diff --git a/ui/src/pages/signup.tsx b/ui/src/pages/signup.tsx
index 1d1bd8e..a953cca 100644
--- a/ui/src/pages/signup.tsx
+++ b/ui/src/pages/signup.tsx
@@ -1,6 +1,6 @@
 import { Button } from "@ui/ui-library/button";
 import { twMerge } from "tailwind-merge";
-import { Link, useNavigate, useSearchParams } from "react-router-dom";
+import { Link, useNavigate } from "react-router-dom";
 import { useState } from "react";
 import { Label, Input, TextField, FieldError } from "@ui/ui-library/field";
 import { useSignUp } from "@ui/hooks/auth";
@@ -10,9 +10,10 @@ import { LoginWithGoogle } from "@ui/components/BrandButtons/LoginWithGoogle";
 
 export function SignUpPage() {
   const navigate = useNavigate();
-  const [searchParams] = useSearchParams();
-  const redirectUrl = searchParams.get("redirect");
-  const { mutate: signUp, isPending } = useSignUp();
+  const redirectUrl = localStorage.getItem("redirectUrl");
+  const { mutate: signUp, isPending } = useSignUp({
+    redirectUrl: redirectUrl ?? null,
+  });
   const [errors, setErrors] = useState<Record<string, string>>({});
 
   const [formData, setFormData] = useState({
@@ -248,7 +249,7 @@ export function SignUpPage() {
             </div>
           </div>
 
-          <LoginWithGoogle redirectUrl={redirectUrl} />
+          <LoginWithGoogle />
 
           <p className="text-center text-sm text-slate-600 dark:text-slate-400">
             Déjà un compte ?{" "}