From b1c93069f77735012f553c761675bba3781be789 Mon Sep 17 00:00:00 2001
From: Chris Chudzicki <christopher.chudzicki@gmail.com>
Date: Thu, 21 Dec 2023 09:28:41 -0500
Subject: [PATCH] preserve current location when redirecting to login

---
 frontends/mit-open/src/common/urls.test.ts    | 19 +++++++++++++++
 frontends/mit-open/src/common/urls.ts         | 23 +++++++++++++++++++
 .../page-components/Header/Header.test.tsx    | 10 ++++++--
 .../src/page-components/Header/UserMenu.tsx   | 12 +++++++---
 4 files changed, 59 insertions(+), 5 deletions(-)
 create mode 100644 frontends/mit-open/src/common/urls.test.ts

diff --git a/frontends/mit-open/src/common/urls.test.ts b/frontends/mit-open/src/common/urls.test.ts
new file mode 100644
index 0000000000..423f44540d
--- /dev/null
+++ b/frontends/mit-open/src/common/urls.test.ts
@@ -0,0 +1,19 @@
+import { login } from "./urls"
+
+test("login encodes the next parameter appropriately", () => {
+  expect(login()).toBe("/login/ol-oidc/?next=/")
+  expect(login({})).toBe("/login/ol-oidc/?next=/")
+
+  expect(
+    login({
+      pathname: "/foo/bar",
+    }),
+  ).toBe("/login/ol-oidc/?next=/foo/bar")
+
+  expect(
+    login({
+      pathname: "/foo/bar",
+      search: "?cat=meow",
+    }),
+  ).toBe("/login/ol-oidc/?next=/foo/bar%3Fcat%3Dmeow")
+})
diff --git a/frontends/mit-open/src/common/urls.ts b/frontends/mit-open/src/common/urls.ts
index 2fb5df13f5..2c1eb9e3e4 100644
--- a/frontends/mit-open/src/common/urls.ts
+++ b/frontends/mit-open/src/common/urls.ts
@@ -18,3 +18,26 @@ export const articlesEditView = (id: number) =>
 
 export const LOGIN = "/login/ol-oidc/"
 export const LOGOUT = "/logout/"
+
+/**
+ * Returns the URL to the login page, with a `next` parameter to redirect back
+ * to the given pathname + search parameters.
+ */
+export const login = ({
+  pathname = "/",
+  search = "",
+}: {
+  pathname?: string
+  search?: string
+} = {}) => {
+  /**
+   * To include search parameters in the next URL, we need to encode them.
+   * If we pass `?next=/foo/bar?cat=meow` directly, Django receives two separate
+   * parameters: `next` and `cat`.
+   *
+   * There's no need to encode the path parameter (it might contain slashes,
+   * but those are allowed in search parameters) so let's keep it readable.
+   */
+  const next = `${pathname}${encodeURIComponent(search)}`
+  return `${LOGIN}?next=${next}`
+}
diff --git a/frontends/mit-open/src/page-components/Header/Header.test.tsx b/frontends/mit-open/src/page-components/Header/Header.test.tsx
index c025f89ac9..dcfc95120b 100644
--- a/frontends/mit-open/src/page-components/Header/Header.test.tsx
+++ b/frontends/mit-open/src/page-components/Header/Header.test.tsx
@@ -76,16 +76,22 @@ describe("UserMenu", () => {
   test.each([
     {
       isAuthenticated: false,
-      expected: { text: "Log in", url: urls.LOGIN },
+      initialUrl: "/foo/bar?cat=meow",
+      expected: {
+        text: "Log in",
+        url: urls.login({ pathname: "/foo/bar", search: "?cat=meow" }),
+      },
     },
     {
       isAuthenticated: true,
+      initialUrl: "/foo/bar?cat=meow",
       expected: { text: "Log out", url: urls.LOGOUT },
     },
   ])(
     "Users (authenticated=$isAuthenticated) see '$expected.text' link",
-    async ({ isAuthenticated, expected }) => {
+    async ({ isAuthenticated, expected, initialUrl }) => {
       renderWithProviders(<Header />, {
+        url: initialUrl,
         user: { is_authenticated: isAuthenticated },
       })
       const menu = await findUserMenu()
diff --git a/frontends/mit-open/src/page-components/Header/UserMenu.tsx b/frontends/mit-open/src/page-components/Header/UserMenu.tsx
index a1f37ba433..e2297a0079 100644
--- a/frontends/mit-open/src/page-components/Header/UserMenu.tsx
+++ b/frontends/mit-open/src/page-components/Header/UserMenu.tsx
@@ -7,6 +7,8 @@ import * as urls from "@/common/urls"
 import { Permissions, hasPermission } from "@/common/permissions"
 import type { User } from "@/types/settings"
 import PersonIcon from "@mui/icons-material/Person"
+import { useLocation } from "react-router"
+import type { Location } from "react-router"
 
 const StyledBadge = styled(Badge)`
   pointer-events: none;
@@ -30,13 +32,16 @@ interface AuthMenuItem extends SimpleMenuItem {
   allow: boolean
 }
 
-const getUserMenuItems = (): SimpleMenuItem[] => {
+const getUserMenuItems = (location: Location): SimpleMenuItem[] => {
   const items: AuthMenuItem[] = [
     {
       label: "Log in",
       key: "login",
       allow: !hasPermission(Permissions.Authenticated),
-      href: urls.LOGIN,
+      href: urls.login({
+        pathname: location.pathname,
+        search: location.search,
+      }),
       LinkComponent: "a",
     },
     {
@@ -66,7 +71,8 @@ const UserIcon: React.FC<{ user: User }> = ({ user }) => {
 
 const UserMenu: React.FC = () => {
   const [visible, setVisible] = useState(false)
-  const items = useMemo(getUserMenuItems, [])
+  const location = useLocation()
+  const items = useMemo(() => getUserMenuItems(location), [location])
 
   return (
     <SimpleMenu