Keep auth params when onboarding (#6182)

src/auth/ha-authorize.ts

@@ -6,19 +6,18 @@ import {
   property,
   PropertyValues,
 } from "lit-element";
-import { AuthProvider, fetchAuthProviders } from "../data/auth";
+import {
+  AuthProvider,
+  fetchAuthProviders,
+  AuthUrlSearchParams,
+} from "../data/auth";
 import { litLocalizeLiteMixin } from "../mixins/lit-localize-lite-mixin";
 import { registerServiceWorker } from "../util/register-service-worker";
 import "./ha-auth-flow";
+import { extractSearchParamsObject } from "../common/url/search-params";
 
 import(/* webpackChunkName: "pick-auth-provider" */ "./ha-pick-auth-provider");
 
-interface QueryParams {
-  client_id?: string;
-  redirect_uri?: string;
-  state?: string;
-}
-
 class HaAuthorize extends litLocalizeLiteMixin(LitElement) {
   @property() public clientId?: string;
 
@@ -33,14 +32,7 @@ class HaAuthorize extends litLocalizeLiteMixin(LitElement) {
   constructor() {
     super();
     this.translationFragment = "page-authorize";
-    const query: QueryParams = {};
-    const values = location.search.substr(1).split("&");
-    for (const item of values) {
-      const value = item.split("=");
-      if (value.length > 1) {
-        query[decodeURIComponent(value[0])] = decodeURIComponent(value[1]);
-      }
-    }
+    const query = extractSearchParamsObject() as AuthUrlSearchParams;
     if (query.client_id) {
       this.clientId = query.client_id;
     }
@@ -145,7 +137,7 @@ class HaAuthorize extends litLocalizeLiteMixin(LitElement) {
         response.status === 400 &&
         authProviders.code === "onboarding_required"
       ) {
-        location.href = "/?";
+        location.href = `/onboarding.html${location.search}`;
         return;
       }
 

src/common/url/search-params.ts

@@ -0,0 +1,8 @@
+export const extractSearchParamsObject = (): { [key: string]: string } => {
+  const query = {};
+  const searchParams = new URLSearchParams(location.search);
+  for (const [key, value] of searchParams.entries()) {
+    query[key] = value;
+  }
+  return query;
+};

src/data/auth.ts

@@ -1,5 +1,11 @@
 import { HomeAssistant } from "../types";
 
+export interface AuthUrlSearchParams {
+  client_id?: string;
+  redirect_uri?: string;
+  state?: string;
+}
+
 export interface AuthProvider {
   name: string;
   id: string;

src/data/onboarding.ts

@@ -51,7 +51,7 @@ export const onboardCoreConfigStep = (hass: HomeAssistant) =>
 
 export const onboardIntegrationStep = (
   hass: HomeAssistant,
-  params: { client_id: string }
+  params: { client_id: string; redirect_uri: string }
 ) =>
   hass.callApi<OnboardingIntegrationStepResponse>(
     "POST",

src/onboarding/ha-onboarding.ts

@@ -1,9 +1,9 @@
 import {
   Auth,
   createConnection,
-  genClientId,
   getAuth,
   subscribeConfig,
+  genClientId,
 } from "home-assistant-js-websocket";
 import {
   customElement,
@@ -14,12 +14,12 @@ import {
 } from "lit-element";
 import { HASSDomEvent } from "../common/dom/fire_event";
 import { subscribeOne } from "../common/util/subscribe-one";
-import { hassUrl } from "../data/auth";
+import { hassUrl, AuthUrlSearchParams } from "../data/auth";
 import {
   fetchOnboardingOverview,
   OnboardingResponses,
   OnboardingStep,
-  ValidOnboardingStep,
+  onboardIntegrationStep,
 } from "../data/onboarding";
 import { subscribeUser } from "../data/ws-user";
 import { litLocalizeLiteMixin } from "../mixins/lit-localize-lite-mixin";
@@ -28,19 +28,28 @@ import { HomeAssistant } from "../types";
 import { registerServiceWorker } from "../util/register-service-worker";
 import "./onboarding-create-user";
 import "./onboarding-loading";
+import { extractSearchParamsObject } from "../common/url/search-params";
 
-interface OnboardingEvent<T extends ValidOnboardingStep> {
-  type: T;
-  result: OnboardingResponses[T];
-}
+type OnboardingEvent =
+  | {
+      type: "user";
+      result: OnboardingResponses["user"];
+    }
+  | {
+      type: "core_config";
+      result: OnboardingResponses["core_config"];
+    }
+  | {
+      type: "integration";
+    };
 
 declare global {
   interface HASSDomEvents {
-    "onboarding-step": OnboardingEvent<ValidOnboardingStep>;
+    "onboarding-step": OnboardingEvent;
   }
 
   interface GlobalEventHandlersEventMap {
-    "onboarding-step": HASSDomEvent<OnboardingEvent<ValidOnboardingStep>>;
+    "onboarding-step": HASSDomEvent<OnboardingEvent>;
   }
 }
 
@@ -150,9 +159,7 @@ class HaOnboarding extends litLocalizeLiteMixin(HassElement) {
     }
   }
 
-  private async _handleStepDone(
-    ev: HASSDomEvent<OnboardingEvent<ValidOnboardingStep>>
-  ) {
+  private async _handleStepDone(ev: HASSDomEvent<OnboardingEvent>) {
     const stepResult = ev.detail;
     this._steps = this._steps!.map((step) =>
       step.step === stepResult.type ? { ...step, done: true } : step
@@ -176,9 +183,41 @@ class HaOnboarding extends litLocalizeLiteMixin(HassElement) {
     } else if (stepResult.type === "core_config") {
       // We do nothing
     } else if (stepResult.type === "integration") {
-      const result = stepResult.result as OnboardingResponses["integration"];
       this._loading = true;
 
+      // Determine if oauth redirect has been provided
+      const externalAuthParams = extractSearchParamsObject() as AuthUrlSearchParams;
+      const authParams =
+        externalAuthParams.client_id && externalAuthParams.redirect_uri
+          ? externalAuthParams
+          : {
+              client_id: genClientId(),
+              redirect_uri: `${location.protocol}//${location.host}/?auth_callback=1`,
+              state: btoa(
+                JSON.stringify({
+                  hassUrl: `${location.protocol}//${location.host}`,
+                  clientId: genClientId(),
+                })
+              ),
+            };
+
+      let result: OnboardingResponses["integration"];
+
+      try {
+        result = await onboardIntegrationStep(this.hass!, {
+          client_id: authParams.client_id!,
+          redirect_uri: authParams.redirect_uri!,
+        });
+      } catch (err) {
+        this.hass!.connection.close();
+        await this.hass!.auth.revoke();
+
+        alert(`Unable to finish onboarding: ${err.message}`);
+
+        document.location.assign("/?");
+        return;
+      }
+
       // If we don't close the connection manually, the connection will be
       // closed when we navigate away from the page. Firefox allows JS to
       // continue to execute, and so HAWS will automatically reconnect once
@@ -191,17 +230,17 @@ class HaOnboarding extends litLocalizeLiteMixin(HassElement) {
       // Revoke current auth token.
       await this.hass!.auth.revoke();
 
-      const state = btoa(
-        JSON.stringify({
-          hassUrl: `${location.protocol}//${location.host}`,
-          clientId: genClientId(),
-        })
-      );
-      document.location.assign(
-        `/?auth_callback=1&code=${encodeURIComponent(
-          result.auth_code
-        )}&state=${state}`
-      );
+      // Build up the url to redirect to
+      let redirectUrl = authParams.redirect_uri!;
+      redirectUrl +=
+        (redirectUrl.includes("?") ? "&" : "?") +
+        `code=${encodeURIComponent(result.auth_code)}`;
+
+      if (authParams.state) {
+        redirectUrl += `&state=${encodeURIComponent(authParams.state)}`;
+      }
+
+      document.location.assign(redirectUrl);
     }
   }
 

src/onboarding/onboarding-integrations.ts

@@ -1,5 +1,4 @@
 import "@material/mwc-button/mwc-button";
-import { genClientId } from "home-assistant-js-websocket";
 import {
   css,
   CSSResult,
@@ -21,7 +20,6 @@ import {
 } from "../data/config_flow";
 import { DataEntryFlowProgress } from "../data/data_entry_flow";
 import { domainToName } from "../data/integration";
-import { onboardIntegrationStep } from "../data/onboarding";
 import {
   loadConfigFlowDialog,
   showConfigFlowDialog,
@@ -169,12 +167,8 @@ class OnboardingIntegrations extends LitElement {
   }
 
   private async _finish() {
-    const result = await onboardIntegrationStep(this.hass, {
-      client_id: genClientId(),
-    });
     fireEvent(this, "onboarding-step", {
       type: "integration",
-      result,
     });
   }