From 15e35eaa824f0a4e5427f81a23837745d5f055d5 Mon Sep 17 00:00:00 2001
From: vchske <chskelt@microsoft.com>
Date: Thu, 26 Oct 2023 16:40:59 -0700
Subject: [PATCH] Adds restarting vcore mongo shell if user enters wrong
 password (#1675)

* Adds restarting vcore mongo shell if user enters wrong password

* Deleted unnecessary comment
---
 src/Terminal/JupyterLabAppFactory.ts | 11 +++++++---
 src/Terminal/index.ts                | 33 ++++++++++++++++++----------
 2 files changed, 30 insertions(+), 14 deletions(-)

diff --git a/src/Terminal/JupyterLabAppFactory.ts b/src/Terminal/JupyterLabAppFactory.ts
index c2e48ba16..2ad61b294 100644
--- a/src/Terminal/JupyterLabAppFactory.ts
+++ b/src/Terminal/JupyterLabAppFactory.ts
@@ -10,9 +10,13 @@ import { userContext } from "UserContext";
 export class JupyterLabAppFactory {
   private isShellStarted: boolean | undefined;
   private checkShellStarted: ((content: string | undefined) => void) | undefined;
-  private onShellExited: () => void;
+  private onShellExited: (restartShell: boolean) => void;
+  private restartShell: boolean;
 
   private isShellExited(content: string | undefined) {
+    if (userContext.apiType === "VCoreMongo" && content?.includes("MongoServerError: Invalid key")) {
+      this.restartShell = true;
+    }
     return content?.includes("cosmosuser@");
   }
 
@@ -32,10 +36,11 @@ export class JupyterLabAppFactory {
     this.isShellStarted = content?.includes("Enter password");
   }
 
-  constructor(closeTab: () => void) {
+  constructor(closeTab: (restartShell: boolean) => void) {
     this.onShellExited = closeTab;
     this.isShellStarted = false;
     this.checkShellStarted = undefined;
+    this.restartShell = false;
 
     switch (userContext.apiType) {
       case "Mongo":
@@ -69,7 +74,7 @@ export class JupyterLabAppFactory {
         if (!this.isShellStarted) {
           this.checkShellStarted(content);
         } else if (this.isShellExited(content)) {
-          this.onShellExited();
+          this.onShellExited(this.restartShell);
         }
       }
     }, this);
diff --git a/src/Terminal/index.ts b/src/Terminal/index.ts
index 42faf8f3e..dda1400ed 100644
--- a/src/Terminal/index.ts
+++ b/src/Terminal/index.ts
@@ -11,6 +11,8 @@ import { JupyterLabAppFactory } from "./JupyterLabAppFactory";
 import { TerminalProps } from "./TerminalProps";
 import "./index.css";
 
+let session: ITerminalConnection | undefined;
+
 const createServerSettings = (props: TerminalProps): ServerConnection.ISettings => {
   let body: BodyInit | undefined;
   let headers: HeadersInit | undefined;
@@ -49,7 +51,7 @@ const createServerSettings = (props: TerminalProps): ServerConnection.ISettings
   return ServerConnection.makeSettings(options);
 };
 
-const initTerminal = async (props: TerminalProps): Promise<ITerminalConnection | undefined> => {
+const initTerminal = async (props: TerminalProps): Promise<void> => {
   // Initialize userContext (only properties which are needed by TelemetryProcessor)
   updateUserContext({
     subscriptionId: props.subscriptionId,
@@ -59,28 +61,37 @@ const initTerminal = async (props: TerminalProps): Promise<ITerminalConnection |
   });
 
   const serverSettings = createServerSettings(props);
+
+  createTerminalApp(props, serverSettings);
+};
+
+const createTerminalApp = async (props: TerminalProps, serverSettings: ServerConnection.ISettings) => {
   const data = { baseUrl: serverSettings.baseUrl };
   const startTime = TelemetryProcessor.traceStart(Action.OpenTerminal, data);
 
   try {
-    const session = await new JupyterLabAppFactory(() => closeTab(props.tabId)).createTerminalApp(serverSettings);
+    session = await new JupyterLabAppFactory((restartShell: boolean) =>
+      closeTab(props, serverSettings, restartShell),
+    ).createTerminalApp(serverSettings);
     TelemetryProcessor.traceSuccess(Action.OpenTerminal, data, startTime);
-    return session;
   } catch (error) {
     TelemetryProcessor.traceFailure(Action.OpenTerminal, data, startTime);
-    return undefined;
+    session = undefined;
   }
 };
 
-const closeTab = (tabId: string): void => {
-  window.parent.postMessage(
-    { type: MessageTypes.CloseTab, data: { tabId: tabId }, signature: "pcIframe" },
-    window.document.referrer,
-  );
+const closeTab = (props: TerminalProps, serverSettings: ServerConnection.ISettings, restartShell: boolean): void => {
+  if (restartShell) {
+    createTerminalApp(props, serverSettings);
+  } else {
+    window.parent.postMessage(
+      { type: MessageTypes.CloseTab, data: { tabId: props.tabId }, signature: "pcIframe" },
+      window.document.referrer,
+    );
+  }
 };
 
 const main = async (): Promise<void> => {
-  let session: ITerminalConnection | undefined;
   postRobot.on(
     "props",
     {
@@ -91,7 +102,7 @@ const main = async (): Promise<void> => {
       // Typescript definition for event is wrong. So read props by casting to <any>
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
       const props = (event as any).data as TerminalProps;
-      session = await initTerminal(props);
+      await initTerminal(props);
     },
   );