Fix wush ssh

Author: coadler

cmd/wush/ssh.go

@@ -40,7 +40,7 @@ func sshCmd() *serpent.Command {
 				if sshStdio {
 					return
 				}
-				fmt.Fprintf(inv.Stderr, str, args...)
+				fmt.Fprintf(inv.Stderr, str+"\n", args...)
 			}
 			if authID == "" {
 				err := huh.NewInput().
@@ -82,9 +82,9 @@ func sshCmd() *serpent.Command {
 			if send.Auth.ReceiverDERPRegionID > 0 {
 				derpStr = dm.Regions[int(send.Auth.ReceiverDERPRegionID)].RegionName
 			}
-			logF("\t> Server overlay DERP home:   %s", cliui.Code(derpStr))
-			logF("\t> Server overlay public key:  %s", cliui.Code(send.Auth.ReceiverPublicKey.ShortString()))
-			logF("\t> Server overlay auth key:    %s", cliui.Code(send.Auth.OverlayPrivateKey.Public().ShortString()))
+			logF("\t> Server overlay DERP home:    %s", cliui.Code(derpStr))
+			logF("\t> Server overlay public key:   %s", cliui.Code(send.Auth.ReceiverPublicKey.ShortString()))
+			logF("\t> Server overlay auth key:     %s", cliui.Code(send.Auth.OverlayPrivateKey.Public().ShortString()))
 
 			s, err := tsserver.NewServer(ctx, logger, send)
 			if err != nil {

xssh/client.go

@@ -2,20 +2,15 @@ package xssh
 
 import (
 	"context"
-	"io"
-	"log/slog"
 	"os"
 	"strings"
-	"sync"
-	"time"
 
 	"github.com/coder/coder/v2/pty"
 	"github.com/coder/serpent"
 	"github.com/mattn/go-isatty"
 	"golang.org/x/crypto/ssh"
 	"golang.org/x/term"
 	"golang.org/x/xerrors"
-	"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
 	"tailscale.com/tsnet"
 )
 
@@ -49,7 +44,7 @@ func TailnetSSH(ctx context.Context, inv *serpent.Invocation, ts *tsnet.Server,
 	sshSession.Stdout = inv.Stdout
 	sshSession.Stderr = inv.Stderr
 
-	if len(inv.Args) > 0 {
+	if len(inv.Args) > 1 {
 		return sshSession.Run(strings.Join(inv.Args, " "))
 	}
 
@@ -108,66 +103,3 @@ func TailnetSSH(ctx context.Context, inv *serpent.Invocation, ts *tsnet.Server,
 
 	return sshSession.Wait()
 }
-
-type rawSSHCopier struct {
-	conn   *gonet.TCPConn
-	logger *slog.Logger
-	r      io.Reader
-	w      io.Writer
-
-	done chan struct{}
-}
-
-func newRawSSHCopier(logger *slog.Logger, conn *gonet.TCPConn, r io.Reader, w io.Writer) *rawSSHCopier {
-	return &rawSSHCopier{conn: conn, logger: logger, r: r, w: w, done: make(chan struct{})}
-}
-
-func (c *rawSSHCopier) copy(wg *sync.WaitGroup) {
-	defer close(c.done)
-	logCtx := context.Background()
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
-		// We close connections using CloseWrite instead of Close, so that the SSH server sees the
-		// closed connection while reading, and shuts down cleanly.  This will trigger the io.Copy
-		// in the server-to-client direction to also be closed and the copy() routine will exit.
-		// This ensures that we don't leave any state in the server, like forwarded ports if
-		// copy() were to return and the underlying tailnet connection torn down before the TCP
-		// session exits. This is a bit of a hack to block shut down at the application layer, since
-		// we can't serialize the TCP and tailnet layers shutting down.
-		//
-		// Of course, if the underlying transport is broken, io.Copy will still return.
-		defer func() {
-			cwErr := c.conn.CloseWrite()
-			c.logger.DebugContext(logCtx, "closed raw SSH connection for writing", "err", cwErr)
-		}()
-
-		_, err := io.Copy(c.conn, c.r)
-		if err != nil {
-			c.logger.ErrorContext(logCtx, "copy stdin error", "err", err)
-		} else {
-			c.logger.DebugContext(logCtx, "copy stdin complete")
-		}
-	}()
-	_, err := io.Copy(c.w, c.conn)
-	if err != nil {
-		c.logger.ErrorContext(logCtx, "copy stdout error", "err", err)
-	} else {
-		c.logger.DebugContext(logCtx, "copy stdout complete")
-	}
-}
-
-func (c *rawSSHCopier) Close() error {
-	err := c.conn.CloseWrite()
-
-	// give the copy() call a chance to return on a timeout, so that we don't
-	// continue tearing down and close the underlying netstack before the SSH
-	// session has a chance to gracefully shut down.
-	t := time.NewTimer(5 * time.Second)
-	defer t.Stop()
-	select {
-	case <-c.done:
-	case <-t.C:
-	}
-	return err
-}