From faccb1243072ba315d0caad1e7169ee65be909d4 Mon Sep 17 00:00:00 2001
From: Bram Kragten <mail@bramkragten.nl>
Date: Thu, 9 Dec 2021 18:57:11 +0100
Subject: [PATCH] Fix keep me logged in (#10835)

---
 src/entrypoints/core.ts         | 38 ++++++++++++++++++---------------
 src/layouts/hass-router-page.ts |  4 +++-
 src/layouts/home-assistant.ts   |  4 +++-
 src/panels/lovelace/hui-root.ts |  2 +-
 4 files changed, 28 insertions(+), 20 deletions(-)

diff --git a/src/entrypoints/core.ts b/src/entrypoints/core.ts
index 03639417ea..2875b3a8e4 100644
--- a/src/entrypoints/core.ts
+++ b/src/entrypoints/core.ts
@@ -37,6 +37,25 @@ declare global {
   }
 }
 
+const clearUrlParams = () => {
+  // Clear auth data from url if we have been able to establish a connection
+  if (location.search.includes("auth_callback=1")) {
+    const searchParams = new URLSearchParams(location.search);
+    // https://github.com/home-assistant/home-assistant-js-websocket/blob/master/lib/auth.ts
+    // Remove all data from QueryCallbackData type
+    searchParams.delete("auth_callback");
+    searchParams.delete("code");
+    searchParams.delete("state");
+    searchParams.delete("storeToken");
+    const search = searchParams.toString();
+    history.replaceState(
+      null,
+      "",
+      `${location.pathname}${search ? `?${search}` : ""}`
+    );
+  }
+};
+
 const authProm = isExternal
   ? () =>
       import("../external_app/external_auth").then(({ createExternalAuth }) =>
@@ -52,23 +71,7 @@ const authProm = isExternal
 const connProm = async (auth) => {
   try {
     const conn = await createConnection({ auth });
-    // Clear auth data from url if we have been able to establish a connection
-    if (location.search.includes("auth_callback=1")) {
-      const searchParams = new URLSearchParams(location.search);
-      // https://github.com/home-assistant/home-assistant-js-websocket/blob/master/lib/auth.ts
-      // Remove all data from QueryCallbackData type
-      searchParams.delete("auth_callback");
-      searchParams.delete("code");
-      searchParams.delete("state");
-      searchParams.delete("storeToken");
-      const search = searchParams.toString();
-      history.replaceState(
-        null,
-        "",
-        `${location.pathname}${search ? `?${search}` : ""}`
-      );
-    }
-
+    clearUrlParams();
     return { auth, conn };
   } catch (err: any) {
     if (err !== ERR_INVALID_AUTH) {
@@ -85,6 +88,7 @@ const connProm = async (auth) => {
     }
     auth = await authProm();
     const conn = await createConnection({ auth });
+    clearUrlParams();
     return { auth, conn };
   }
 };
diff --git a/src/layouts/hass-router-page.ts b/src/layouts/hass-router-page.ts
index 5c37cd577c..6fb01e521f 100644
--- a/src/layouts/hass-router-page.ts
+++ b/src/layouts/hass-router-page.ts
@@ -127,7 +127,9 @@ export class HassRouterPage extends ReactiveElement {
 
         // Update the url if we know where we're mounted.
         if (route) {
-          navigate(`${route.prefix}/${result}`, { replace: true });
+          navigate(`${route.prefix}/${result}${location.search}`, {
+            replace: true,
+          });
         }
       }
     }
diff --git a/src/layouts/home-assistant.ts b/src/layouts/home-assistant.ts
index bb452d1781..e3bf2f9d8b 100644
--- a/src/layouts/home-assistant.ts
+++ b/src/layouts/home-assistant.ts
@@ -51,7 +51,9 @@ export class HomeAssistantAppEl extends QuickBarMixin(HassElement) {
     const path = curPath();
 
     if (["", "/"].includes(path)) {
-      navigate(`/${getStorageDefaultPanelUrlPath()}`, { replace: true });
+      navigate(`/${getStorageDefaultPanelUrlPath()}${location.search}`, {
+        replace: true,
+      });
     }
     this._route = {
       prefix: "",
diff --git a/src/panels/lovelace/hui-root.ts b/src/panels/lovelace/hui-root.ts
index 61445901d7..fa572585cb 100644
--- a/src/panels/lovelace/hui-root.ts
+++ b/src/panels/lovelace/hui-root.ts
@@ -698,7 +698,7 @@ class HUIRoot extends LitElement {
 
   private _navigateToView(path: string | number, replace?: boolean) {
     if (!this.lovelace!.editMode) {
-      navigate(`${this.route!.prefix}/${path}`, { replace });
+      navigate(`${this.route!.prefix}/${path}${location.search}`, { replace });
       return;
     }
     navigate(`${this.route!.prefix}/${path}?${addSearchParam({ edit: "1" })}`, {