From cf3900a9067ddb4e6921538b9d520bd54fe5fb19 Mon Sep 17 00:00:00 2001
From: Jeff Mitchell <jeffrey.mitchell@gmail.com>
Date: Wed, 8 Nov 2023 12:30:54 -0500
Subject: [PATCH] Fix three minor issues: (#4022)

* If passing a private key without specifying it in URL syntax, give a
  proper error message (previously we were only making sure it was a
  scheme we knew how to handle, but we need to also be more informative
  if it's not a URL at all)

* If the handshake from the proxy fails, don't hang, just exit

* Update some error messages to use capitalization in line with our
  standard CLI error format
---
 internal/cmd/commands/connect/connect.go           | 14 ++++++++------
 .../credentialscmd/ssh-private-key_funcs.go        |  2 +-
 2 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/internal/cmd/commands/connect/connect.go b/internal/cmd/commands/connect/connect.go
index a521732e35..092fa7bf4f 100644
--- a/internal/cmd/commands/connect/connect.go
+++ b/internal/cmd/commands/connect/connect.go
@@ -629,10 +629,10 @@ func (c *Command) Run(args []string) (retCode int) {
 		ctx, cancel := context.WithTimeout(context.Background(), sessionCancelTimeout)
 		wsConn, err := c.getWsConn(ctx, workerAddr, transport)
 		if err != nil {
-			c.PrintCliError(fmt.Errorf("error fetching connection to send session teardown request to worker: %w", err))
+			c.PrintCliError(fmt.Errorf("Error fetching connection to send session teardown request to worker: %w", err))
 		} else {
 			if err := c.sendSessionTeardown(ctx, wsConn, tofuToken); err != nil {
-				c.PrintCliError(fmt.Errorf("error sending session teardown request to worker: %w", err))
+				c.PrintCliError(fmt.Errorf("Error sending session teardown request to worker: %w", err))
 			}
 		}
 		cancel()
@@ -651,7 +651,7 @@ func (c *Command) Run(args []string) (retCode int) {
 		case "json":
 			out, err := json.Marshal(&termInfo)
 			if err != nil {
-				c.PrintCliError(fmt.Errorf("error marshaling termination information: %w", err))
+				c.PrintCliError(fmt.Errorf("Error marshaling termination information: %w", err))
 				return base.CommandCliError
 			}
 			c.UI.Output(string(out))
@@ -675,7 +675,7 @@ func (c *Command) printCredentials(creds []*targets.SessionCredential) error {
 			Credentials: creds,
 		})
 		if err != nil {
-			return fmt.Errorf("error marshaling credential information: %w", err)
+			return fmt.Errorf("Error marshaling credential information: %w", err)
 		}
 		c.UI.Output(string(out))
 	}
@@ -744,7 +744,7 @@ func (c *Command) runTcpProxyV1(
 ) error {
 	handshake := proxy.ClientHandshake{TofuToken: tofuToken}
 	if err := wspb.Write(c.proxyCtx, wsConn, &handshake); err != nil {
-		return fmt.Errorf("error sending handshake to worker: %w", err)
+		return fmt.Errorf("Error sending handshake to worker: %w", err)
 	}
 	var handshakeResult proxy.HandshakeResult
 	if err := wspb.Read(c.proxyCtx, wsConn, &handshakeResult); err != nil {
@@ -761,7 +761,9 @@ func (c *Command) runTcpProxyV1(
 			c.proxyCancel()
 			return errors.New("Session token has already been used")
 		default:
-			return fmt.Errorf("error reading handshake result: %w", err)
+			// If we can't handshake we can't do anything, so quit out
+			c.proxyCancel()
+			return fmt.Errorf("Error reading handshake result: %w", err)
 		}
 	}
 
diff --git a/internal/cmd/commands/credentialscmd/ssh-private-key_funcs.go b/internal/cmd/commands/credentialscmd/ssh-private-key_funcs.go
index 31329b29f6..230efa92bc 100644
--- a/internal/cmd/commands/credentialscmd/ssh-private-key_funcs.go
+++ b/internal/cmd/commands/credentialscmd/ssh-private-key_funcs.go
@@ -84,7 +84,7 @@ func extraSshPrivateKeyFlagHandlingFuncImpl(c *SshPrivateKeyCommand, _ *base.Fla
 	privateKey, err := parseutil.MustParsePath(c.flagPrivateKey)
 	switch {
 	case err == nil:
-	case errors.Is(err, parseutil.ErrNotParsed):
+	case errors.Is(err, parseutil.ErrNotAUrl), errors.Is(err, parseutil.ErrNotParsed):
 		c.UI.Error("Private key flag must be used with env:// or file:// syntax")
 		return false
 	default: