improve robustness and diagnostics

Author: lonnywong

tsshd/client.go

@@ -29,6 +29,7 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -135,7 +136,20 @@ func NewSshUdpClient(opts *UdpClientOptions) (*SshUdpClient, error) {
 	if err != nil {
 		return nil, err
 	}
+	beginTime := time.Now()
 	if err := udpClient.networkProxy.renewTransportPath(opts.ProxyClient, opts.ConnectTimeout); err != nil {
+		if opts.ConnectTimeout > 2*time.Second && time.Since(beginTime) > (opts.ConnectTimeout-time.Second) {
+			net := "UDP"
+			if opts.ServerInfo.ProxyMode == kProxyModeTCP {
+				net = "TCP"
+			}
+			port := opts.TsshdAddr
+			if pos := strings.LastIndex(opts.TsshdAddr, ":"); pos >= 0 {
+				port = opts.TsshdAddr[pos+1:]
+			}
+			return nil, fmt.Errorf("%v\r\n%s", err, fmt.Sprintf(
+				"\033[0;36mHint:\033[0m This may be caused by a firewall blocking the %s port (%s) that tsshd is listening on.", net, port))
+		}
 		return nil, err
 	}
 
@@ -348,7 +362,7 @@ func (c *SshUdpClient) DialUDP(network, addr string, timeout time.Duration) (Pac
 	}
 
 	c.exitWG.Add(1)
-	return &sshUdpPacketConn{conn, c}, nil
+	return &sshUdpPacketConn{packetConn: conn, client: c}, nil
 }
 
 // Listen requests the remote peer open a listening socket on addr
@@ -1224,9 +1238,13 @@ func (c *sshUdpChannel) Stderr() io.ReadWriter {
 type sshUdpPacketConn struct {
 	*packetConn
 	client *SshUdpClient
+	closed atomic.Bool
 }
 
 func (c *sshUdpPacketConn) Close() error {
+	if !c.closed.CompareAndSwap(false, true) {
+		return nil
+	}
 	err := c.packetConn.Close()
 	c.client.exitWG.Done()
 	return err

tsshd/rekey.go

@@ -188,6 +188,10 @@ func (r *rotatingCrypto) startRekey() {
 		return
 	}
 
+	if r.client.IsClosed() {
+		return
+	}
+
 	err := func() error {
 		curve := ecdh.P256()
 		clientPriKey, err := curve.GenerateKey(crypto_rand.Reader)

tsshd/session.go

@@ -90,7 +90,7 @@ type sessionContext struct {
 	rows    int
 	cmd     *exec.Cmd
 	pty     *tsshdPty
-	wg      sync.WaitGroup
+	outWG   sync.WaitGroup
 	stdin   io.WriteCloser
 	stdout  io.ReadCloser
 	stderr  io.ReadCloser
@@ -452,11 +452,11 @@ func (c *sessionContext) forwardIO(stream Stream) {
 	}
 
 	if c.stdout != nil {
-		c.wg.Go(func() { c.forwardOutput("stdout", c.stdout, stream) })
+		c.outWG.Go(func() { c.forwardOutput("stdout", c.stdout, stream) })
 	}
 
 	if c.stderr != nil {
-		c.wg.Go(func() {
+		c.outWG.Go(func() {
 			if stderr := getStderrStream(c.id); stderr != nil {
 				c.forwardOutput("stderr", c.stderr, stderr.stream)
 				stderr.Close()
@@ -475,17 +475,34 @@ func (c *sessionContext) Wait() {
 	// windows pty only close the stdout in pty.Wait
 	if runtime.GOOS == "windows" && c.pty != nil {
 		_ = c.pty.Wait()
-		c.wg.Wait()
+		c.outWG.Wait()
 		debug("session [%d] wait completed", c.id)
 		return
 	}
 
-	c.wg.Wait() // wait for the output done first to prevent cmd.Wait close output too early
+	done := make(chan struct{})
+	go func() {
+		c.outWG.Wait() // wait for the output first to prevent cmd.Wait close output too early
+		close(done)
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+	}
+
 	if c.pty != nil {
 		_ = c.pty.Wait()
 	} else {
 		_ = c.cmd.Wait()
 	}
+
+	select {
+	case <-done:
+	case <-time.After(3 * time.Second):
+		warning("child process has exited, but output streams did not close in time")
+	}
+
 	debug("session [%d] wait completed", c.id)
 }