From 53e94cc83dc00911b7add57b2a4dce0e682abe22 Mon Sep 17 00:00:00 2001
From: Alex Gartner <git@agartner.com>
Date: Mon, 25 Nov 2024 21:53:01 -0800
Subject: [PATCH] client: make server downtime more graceful

---
 client/tunnel.go          |  9 +++++----
 cmd/tunnel-client/main.go | 15 +++++++++------
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/client/tunnel.go b/client/tunnel.go
index 1bf3c28..4146141 100644
--- a/client/tunnel.go
+++ b/client/tunnel.go
@@ -5,7 +5,7 @@ import (
 	"context"
 	"crypto/tls"
 	"fmt"
-	"io/ioutil"
+	"io"
 	"net"
 	"net/http"
 	"net/url"
@@ -39,16 +39,17 @@ func NewTunnel(server, hostname, token string, useTLS, tlsSkipVerify, httpTarget
 	}
 }
 
-func (t *Tunnel) Start() {
+func (t *Tunnel) Start() error {
 	conn, err := t.stage1(true)
 	if err != nil {
-		panic(fmt.Errorf("unable to connect to server: %w", err))
+		return fmt.Errorf("unable to complete initial connection to server: %w", err)
 	}
 	go t.stage2(conn)
 
 	for i := 0; i < 20; i++ {
 		go t.both()
 	}
+	return nil
 }
 
 func (t *Tunnel) Shutdown() {
@@ -170,7 +171,7 @@ func (t *Tunnel) stage2(conn net.Conn) {
 		s := fmt.Sprintf("target %s returned error %s", t.target, err)
 		r := http.Response{
 			StatusCode: 500,
-			Body:       ioutil.NopCloser(bytes.NewBufferString(s)),
+			Body:       io.NopCloser(bytes.NewBufferString(s)),
 		}
 		r.Write(conn)
 		return
diff --git a/cmd/tunnel-client/main.go b/cmd/tunnel-client/main.go
index fc7aaff..4cca8e0 100644
--- a/cmd/tunnel-client/main.go
+++ b/cmd/tunnel-client/main.go
@@ -36,10 +36,11 @@ func init() {
 }
 
 var rootCmd = &cobra.Command{
-	Use:     "tunnel <[http://]hostname:port>",
-	Short:   "Start a tunnel to a local server",
-	Example: "tunnel localhost:8888",
-	Args:    cobra.ExactArgs(1),
+	Use:          "tunnel <[http://]hostname:port>",
+	Short:        "Start a tunnel to a local server",
+	Example:      "tunnel localhost:8888",
+	SilenceUsage: true,
+	Args:         cobra.ExactArgs(1),
 	// use envy to parse TUNNEL_ vars into environment
 	PersistentPreRun: func(cmd *cobra.Command, args []string) {
 		envy.ParseCobra(cmd, envy.CobraConfig{
@@ -90,7 +91,10 @@ var rootCmd = &cobra.Command{
 			}
 
 			tunnel := client.NewTunnel(controlName, hostnameFqdn, token, useTLS, tlsSkipVerify, httpTargetHostHeader, target)
-			tunnel.Start()
+			err := tunnel.Start()
+			if err != nil {
+				return fmt.Errorf("start %s: %w", controlName, err)
+			}
 			tunnels = append(tunnels, tunnel)
 		}
 
@@ -107,7 +111,6 @@ var rootCmd = &cobra.Command{
 
 func main() {
 	if err := rootCmd.Execute(); err != nil {
-		fmt.Println(err)
 		os.Exit(1)
 	}
 }