From bb734be4bcdfc4f5e0f82b899fe5904ef3b573ed Mon Sep 17 00:00:00 2001
From: Bram Kragten <mail@bramkragten.nl>
Date: Thu, 28 Oct 2021 19:05:30 +0200
Subject: [PATCH] Remove keep logged in query string after login, dont show on
 select_mfa_module step (#10439)

---
 src/auth/ha-auth-flow.ts         | 3 ++-
 src/common/auth/token_storage.ts | 6 ++----
 src/entrypoints/core.ts          | 4 +++-
 3 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/src/auth/ha-auth-flow.ts b/src/auth/ha-auth-flow.ts
index aa02e61d58..d1b678e31e 100644
--- a/src/auth/ha-auth-flow.ts
+++ b/src/auth/ha-auth-flow.ts
@@ -206,7 +206,8 @@ class HaAuthFlow extends litLocalizeLiteMixin(LitElement) {
             .computeError=${this._computeErrorCallback(step)}
             @value-changed=${this._stepDataChanged}
           ></ha-form>
-          ${this.clientId === genClientId() && step.step_id !== "mfa"
+          ${this.clientId === genClientId() &&
+          !["select_mfa_module", "mfa"].includes(step.step_id)
             ? html`
                 <ha-formfield
                   class="store-token"
diff --git a/src/common/auth/token_storage.ts b/src/common/auth/token_storage.ts
index 6b432a7029..302da3d7fd 100644
--- a/src/common/auth/token_storage.ts
+++ b/src/common/auth/token_storage.ts
@@ -1,4 +1,5 @@
 import { AuthData } from "home-assistant-js-websocket";
+import { extractSearchParam } from "../url/search-params";
 
 const storage = window.localStorage || {};
 
@@ -31,10 +32,7 @@ export function askWrite() {
 export function saveTokens(tokens: AuthData | null) {
   tokenCache.tokens = tokens;
 
-  if (
-    !tokenCache.writeEnabled &&
-    new URLSearchParams(window.location.search).get("storeToken") === "true"
-  ) {
+  if (!tokenCache.writeEnabled && extractSearchParam("storeToken") === "true") {
     tokenCache.writeEnabled = true;
   }
 
diff --git a/src/entrypoints/core.ts b/src/entrypoints/core.ts
index 26986822f1..03639417ea 100644
--- a/src/entrypoints/core.ts
+++ b/src/entrypoints/core.ts
@@ -60,10 +60,12 @@ const connProm = async (auth) => {
       searchParams.delete("auth_callback");
       searchParams.delete("code");
       searchParams.delete("state");
+      searchParams.delete("storeToken");
+      const search = searchParams.toString();
       history.replaceState(
         null,
         "",
-        `${location.pathname}?${searchParams.toString()}`
+        `${location.pathname}${search ? `?${search}` : ""}`
       );
     }