Gracefully close SSH server, add verbose option to wush receive

Author: coadler

cmd/wush/receive.go

@@ -15,6 +15,7 @@ import (
 	"tailscale.com/tsnet"
 
 	cslog "cdr.dev/slog"
+	csloghuman "cdr.dev/slog/sloggers/sloghuman"
 	"github.com/coder/coder/v2/agent/agentssh"
 	"github.com/coder/serpent"
 	"github.com/coder/wush/cliui"
@@ -23,14 +24,21 @@ import (
 )
 
 func receiveCmd() *serpent.Command {
-	var overlayType string
+	var (
+		overlayType string
+		verbose     bool
+	)
 	return &serpent.Command{
 		Use:     "receive",
 		Aliases: []string{"host"},
 		Long:    "Runs the wush server. Allows other wush CLIs to connect to this computer.",
 		Handler: func(inv *serpent.Invocation) error {
 			ctx := inv.Context()
-			logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+			var logSink io.Writer = io.Discard
+			if verbose {
+				logSink = inv.Stderr
+			}
+			logger := slog.New(slog.NewTextHandler(logSink, nil))
 			dm, err := tsserver.DERPMapTailscale(ctx)
 			if err != nil {
 				return err
@@ -77,7 +85,12 @@ func receiveCmd() *serpent.Command {
 
 			fmt.Println(cliui.Timestamp(time.Now()), "Wireguard is ready")
 
-			sshSrv, err := agentssh.NewServer(ctx, cslog.Make( /* sloghuman.Sink(os.Stderr)*/ ), prometheus.NewRegistry(), fs, nil)
+			sshSrv, err := agentssh.NewServer(ctx,
+				cslog.Make(csloghuman.Sink(logSink)),
+				prometheus.NewRegistry(),
+				fs,
+				nil,
+			)
 			if err != nil {
 				return err
 			}
@@ -87,14 +100,33 @@ func receiveCmd() *serpent.Command {
 				return err
 			}
 
-			return sshSrv.Serve(ls)
+			go func() {
+				fmt.Println(cliui.Timestamp(time.Now()), "SSH server listening")
+				err := sshSrv.Serve(ls)
+				if err != nil {
+					logger.Info("ssh server exited", "err", err)
+				}
+			}()
+
+			ctx, ctxCancel := inv.SignalNotifyContext(ctx, os.Interrupt)
+			defer ctxCancel()
+
+			<-ctx.Done()
+			return sshSrv.Close()
 		},
 		Options: []serpent.Option{
 			{
 				Flag:    "overlay-type",
 				Default: "derp",
 				Value:   serpent.EnumOf(&overlayType, "derp", "stun"),
 			},
+			{
+				Flag:          "verbose",
+				FlagShorthand: "v",
+				Description:   "Enable verbose logging.",
+				Default:       "false",
+				Value:         serpent.BoolOf(&verbose),
+			},
 		},
 	}
 }

cmd/wush/ssh.go

@@ -23,17 +23,19 @@ import (
 
 func sshCmd() *serpent.Command {
 	var (
-		authID             string
+		authKey            string
 		waitP2P            bool
 		stunAddrOverride   string
 		stunAddrOverrideIP netip.Addr
 		sshStdio           bool
 	)
 	return &serpent.Command{
-		Use: "wush",
+		Use:     "wush",
+		Aliases: []string{"ssh"},
 		Long: "Opens an SSH connection to a " + cliui.Code("wush") + " peer. " +
 			"Use " + cliui.Code("wush receive") + " on the computer you would like to connect to.",
 		Handler: func(inv *serpent.Invocation) error {
+
 			ctx := inv.Context()
 			logger := slog.New(slog.NewTextHandler(io.Discard, nil))
 			logF := func(str string, args ...any) {
@@ -42,10 +44,10 @@ func sshCmd() *serpent.Command {
 				}
 				fmt.Fprintf(inv.Stderr, str+"\n", args...)
 			}
-			if authID == "" {
+			if authKey == "" {
 				err := huh.NewInput().
-					Title("Enter the receiver's Auth ID:").
-					Value(&authID).
+					Title("Enter the receiver's Auth key:").
+					Value(&authKey).
 					Run()
 				if err != nil {
 					return fmt.Errorf("get auth id: %w", err)
@@ -67,7 +69,7 @@ func sshCmd() *serpent.Command {
 			send := overlay.NewSendOverlay(logger, dm)
 			send.STUNIPOverride = stunAddrOverrideIP
 
-			err = send.Auth.Parse(authID)
+			err = send.Auth.Parse(authKey)
 			if err != nil {
 				return fmt.Errorf("parse auth key: %w", err)
 			}
@@ -133,11 +135,11 @@ func sshCmd() *serpent.Command {
 		},
 		Options: []serpent.Option{
 			{
-				Flag:        "auth-id",
-				Env:         "WUSH_AUTH_ID",
-				Description: "The auth id returned by " + cliui.Code("wush receive") + ". If not provided, it will be asked for on startup.",
+				Flag:        "auth",
+				Env:         "WUSH_AUTH",
+				Description: "The auth key returned by " + cliui.Code("wush receive") + ". If not provided, it will be asked for on startup.",
 				Default:     "",
-				Value:       serpent.StringOf(&authID),
+				Value:       serpent.StringOf(&authKey),
 			},
 			{
 				Flag:    "stun-ip-override",