From c2ad91bb7470632274ff8b8d7f16fa3de9ed0e65 Mon Sep 17 00:00:00 2001
From: Bryan <bryan@coder.com>
Date: Thu, 17 Feb 2022 08:35:36 -0800
Subject: [PATCH] fix: Redirect to '?redirect' query parameter after successful
 login (#307)

Fixes #304
Unblocks #298

After logging in, the login flow should redirect to a whatever path is specified by the `?redirect` query parameter. This is important for cases like #298 - where we need to set `?redirect=%2Fcli_auth`, but also really any case where the user is linked and might have to go back to the login screen.

The fix is simple - just check if the `redirect` query parameter is set, and if it is, use that as the path to redirect to on success. Also adds a test case - we had one checking that we redirect to the default (root `/`) url, but not one of the `?redirect` param
---
 site/components/SignIn/SignInForm.test.tsx | 24 +++++++++++++++++++++-
 site/components/SignIn/SignInForm.tsx      | 16 +++++++++++++--
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/site/components/SignIn/SignInForm.test.tsx b/site/components/SignIn/SignInForm.test.tsx
index 0f1797de4b..1306fd43fa 100644
--- a/site/components/SignIn/SignInForm.test.tsx
+++ b/site/components/SignIn/SignInForm.test.tsx
@@ -55,7 +55,29 @@ describe("SignInForm", () => {
     act(() => elem.click())
 
     // Then
-    // Should redirect because login was successfully
+    // Should redirect because login was successful
     await waitFor(() => expect(singletonRouter).toMatchObject({ asPath: "/" }))
   })
+
+  it("respects ?redirect query parameter when complete", async () => {
+    // Given
+    const loginHandler = (_email: string, _password: string) => Promise.resolve()
+    // Set a path to redirect to after login is successful
+    mockRouter.setCurrentUrl("/login?redirect=%2Fsome%2Fother%2Fpath")
+
+    // When
+    // Render the component
+    const { container } = render(<SignInForm loginHandler={loginHandler} />)
+    // Set user / password
+    const inputs = container.querySelectorAll("input")
+    fireEvent.change(inputs[0], { target: { value: "test@coder.com" } })
+    fireEvent.change(inputs[1], { target: { value: "password" } })
+    // Click sign-in
+    const elem = await screen.findByText("Sign In")
+    act(() => elem.click())
+
+    // Then
+    // Should redirect to /some/other/path because ?redirect was specified and login was successful
+    await waitFor(() => expect(singletonRouter).toMatchObject({ asPath: "/some/other/path" }))
+  })
 })
diff --git a/site/components/SignIn/SignInForm.tsx b/site/components/SignIn/SignInForm.tsx
index a456cf57ed..fdb0e0c43a 100644
--- a/site/components/SignIn/SignInForm.tsx
+++ b/site/components/SignIn/SignInForm.tsx
@@ -1,6 +1,6 @@
 import { makeStyles } from "@material-ui/core/styles"
 import { FormikContextType, useFormik } from "formik"
-import { useRouter } from "next/router"
+import { NextRouter, useRouter } from "next/router"
 import React from "react"
 import { useSWRConfig } from "swr"
 import * as Yup from "yup"
@@ -9,6 +9,7 @@ import { Welcome } from "./Welcome"
 import { FormTextField } from "../Form"
 import * as API from "./../../api"
 import { LoadingButton } from "./../Button"
+import { firstOrItem } from "../../util/array"
 
 /**
  * BuiltInAuthFormValues describes a form using built-in (email/password)
@@ -61,7 +62,9 @@ export const SignInForm: React.FC<SignInProps> = ({
         await loginHandler(email, password)
         // Tell SWR to invalidate the cache for the user endpoint
         await mutate("/api/v2/users/me")
-        await router.push("/")
+
+        const redirect = getRedirectFromRouter(router)
+        await router.push(redirect)
       } catch (err) {
         helpers.setFieldError("password", "The username or password is incorrect.")
       }
@@ -117,3 +120,12 @@ export const SignInForm: React.FC<SignInProps> = ({
     </>
   )
 }
+
+const getRedirectFromRouter = (router: NextRouter) => {
+  const defaultRedirect = "/"
+  if (router.query?.redirect) {
+    return firstOrItem(router.query.redirect, defaultRedirect)
+  } else {
+    return defaultRedirect
+  }
+}