From 7741e9feb02c1bff989f55ec7b9e29bbbf576bb6 Mon Sep 17 00:00:00 2001
From: Mihai Parparita <mihai@tailscale.com>
Date: Mon, 17 Oct 2022 15:02:59 -0700
Subject: [PATCH] cmd/tsconnect: add progress and connection callbacks

Allows UI to display slightly more fine-grained progress when the SSH
connection is being established.

Updates tailscale/corp#7186

Signed-off-by: Mihai Parparita <mihai@tailscale.com>
---
 cmd/tsconnect/src/app/ssh.tsx        |  7 ++++++-
 cmd/tsconnect/src/lib/ssh.ts         | 16 ++++++++++++----
 cmd/tsconnect/src/types/wasm_js.d.ts |  2 ++
 cmd/tsconnect/wasm/wasm_js.go        | 17 +++++++++++++++--
 4 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/cmd/tsconnect/src/app/ssh.tsx b/cmd/tsconnect/src/app/ssh.tsx
index 8048d9b2a..05ce130dc 100644
--- a/cmd/tsconnect/src/app/ssh.tsx
+++ b/cmd/tsconnect/src/app/ssh.tsx
@@ -46,7 +46,12 @@ function SSHSession({
   const ref = useRef<HTMLDivElement>(null)
   useEffect(() => {
     if (ref.current) {
-      runSSHSession(ref.current, def, ipn, onDone, (err) => console.error(err))
+      runSSHSession(ref.current, def, ipn, {
+        onConnectionProgress: (p) => console.log("Connection progress", p),
+        onConnected() {},
+        onError: (err) => console.error(err),
+        onDone,
+      })
     }
   }, [ref])
 
diff --git a/cmd/tsconnect/src/lib/ssh.ts b/cmd/tsconnect/src/lib/ssh.ts
index 3c7552c68..afbcdb895 100644
--- a/cmd/tsconnect/src/lib/ssh.ts
+++ b/cmd/tsconnect/src/lib/ssh.ts
@@ -9,12 +9,18 @@ export type SSHSessionDef = {
   timeoutSeconds?: number
 }
 
+export type SSHSessionCallbacks = {
+  onConnectionProgress: (messsage: string) => void
+  onConnected: () => void
+  onDone: () => void
+  onError?: (err: string) => void
+}
+
 export function runSSHSession(
   termContainerNode: HTMLDivElement,
   def: SSHSessionDef,
   ipn: IPN,
-  onDone: () => void,
-  onError?: (err: string) => void,
+  callbacks: SSHSessionCallbacks,
   terminalOptions?: ITerminalOptions
 ) {
   const parentWindow = termContainerNode.ownerDocument.defaultView ?? window
@@ -49,7 +55,7 @@ export function runSSHSession(
       term.write(input)
     },
     writeErrorFn(err) {
-      onError?.(err)
+      callbacks.onError?.(err)
       term.write(err)
     },
     setReadFn(hook) {
@@ -57,13 +63,15 @@ export function runSSHSession(
     },
     rows: term.rows,
     cols: term.cols,
+    onConnectionProgress: callbacks.onConnectionProgress,
+    onConnected: callbacks.onConnected,
     onDone() {
       resizeObserver?.disconnect()
       term.dispose()
       if (handleUnload) {
         parentWindow.removeEventListener("unload", handleUnload)
       }
-      onDone()
+      callbacks.onDone()
     },
     timeoutSeconds: def.timeoutSeconds,
   })
diff --git a/cmd/tsconnect/src/types/wasm_js.d.ts b/cmd/tsconnect/src/types/wasm_js.d.ts
index 86d5cb245..1802014fc 100644
--- a/cmd/tsconnect/src/types/wasm_js.d.ts
+++ b/cmd/tsconnect/src/types/wasm_js.d.ts
@@ -25,6 +25,8 @@ declare global {
         cols: number
         /** Defaults to 5 seconds */
         timeoutSeconds?: number
+        onConnectionProgress: (message: string) => void
+        onConnected: () => void
         onDone: () => void
       }
     ): IPNSSHSession
diff --git a/cmd/tsconnect/wasm/wasm_js.go b/cmd/tsconnect/wasm/wasm_js.go
index b62485ce8..e33a5797d 100644
--- a/cmd/tsconnect/wasm/wasm_js.go
+++ b/cmd/tsconnect/wasm/wasm_js.go
@@ -364,15 +364,21 @@ func (s *jsSSHSession) Run() {
 	if jsTimeoutSeconds := s.termConfig.Get("timeoutSeconds"); jsTimeoutSeconds.Type() == js.TypeNumber {
 		timeoutSeconds = jsTimeoutSeconds.Float()
 	}
+	onConnectionProgress := s.termConfig.Get("onConnectionProgress")
+	onConnected := s.termConfig.Get("onConnected")
 	onDone := s.termConfig.Get("onDone")
 	defer onDone.Invoke()
 
 	writeError := func(label string, err error) {
 		writeErrorFn.Invoke(fmt.Sprintf("%s Error: %v\r\n", label, err))
 	}
+	reportProgress := func(message string) {
+		onConnectionProgress.Invoke(message)
+	}
 
 	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeoutSeconds*float64(time.Second)))
 	defer cancel()
+	reportProgress(fmt.Sprintf("Connecting to %s…", strings.Split(s.host, ".")[0]))
 	c, err := s.jsIPN.dialer.UserDial(ctx, "tcp", net.JoinHostPort(s.host, "22"))
 	if err != nil {
 		writeError("Dial", err)
@@ -381,10 +387,16 @@ func (s *jsSSHSession) Run() {
 	defer c.Close()
 
 	config := &ssh.ClientConfig{
-		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
-		User:            s.username,
+		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
+			// Host keys are not used with Tailscale SSH, but we can use this
+			// callback to know that the connection has been established.
+			reportProgress("SSH connection established…")
+			return nil
+		},
+		User: s.username,
 	}
 
+	reportProgress("Starting SSH client…")
 	sshConn, _, _, err := ssh.NewClientConn(c, s.host, config)
 	if err != nil {
 		writeError("SSH Connection", err)
@@ -442,6 +454,7 @@ func (s *jsSSHSession) Run() {
 		return
 	}
 
+	onConnected.Invoke()
 	err = session.Wait()
 	if err != nil {
 		writeError("Wait", err)