From 93c16ab145680a5ba636620ee1d1ae68849b4a75 Mon Sep 17 00:00:00 2001
From: Michael Vogt <mvo@ubuntu.com>
Date: Tue, 7 Feb 2023 14:54:08 +0100
Subject: [PATCH] client: fix hang on ssh.Dial() on reboot

This commit adds a thin wrapper around the real ssh.Dial() that
additionally sets a deadline on the underlying connection.

It is needed because a ssh.Dial() can happens right after the
reboot command is issued. The net.Dial() itself is successful but
then then during the ssh session setup the TCP connection ends
because of the reboot. The golang "ssh" package has no concpt of
"ssh -o ServerAliveInterval=10" or simialr so the code will
just hang in a read forever. This was observed running the
spread "cerberus" tests on ubuntu 23.04.

Note that half of the function is just a copy of
golang.org/x/crypto/ssh/client.go:func Dial()
and only the conn.SetDeadline() bits are new.

See also e.g. https://github.com/golang/go/issues/51926 for
various bugreports about the golang "ssh" package and hangs.
---
 spread/client.go | 40 ++++++++++++++++++++++++++++++++++++++--
 1 file changed, 38 insertions(+), 2 deletions(-)

diff --git a/spread/client.go b/spread/client.go
index f0f51082..ba983e77 100644
--- a/spread/client.go
+++ b/spread/client.go
@@ -22,7 +22,43 @@ import (
 	"syscall"
 )
 
-var sshDial = ssh.Dial
+// sshDial is a copy/thin wrapper around the real ssh.Dial() that
+// additionally sets a deadline on the underlying connection.
+//
+// It is needed because a ssh.Dial() can happens right after the
+// reboot command is issued. The net.Dial() itself is successful but
+// then then during the ssh session setup the TCP connection ends
+// because of the reboot. The golang "ssh" package has no concpt of
+// "ssh -o ServerAliveInterval=10" or simialr so the code will
+// just hang in a read forever. This was observed running the
+// spread "cerberus" tests on ubuntu 23.04.
+//
+// Note that half of the function is just a copy of
+// golang.org/x/crypto/ssh/client.go:func Dial()
+var sshDial = func(network, addr string, config *ssh.ClientConfig) (*ssh.Client, error) {
+	conn, err := net.DialTimeout(network, addr, config.Timeout)
+	if err != nil {
+		return nil, err
+	}
+	// Ensure that if the connection goes away during e.g. a reboot
+	// the code does not hang forever.
+	//
+	// See e.g. https://github.com/golang/go/issues/51926
+	if config.Timeout > 0 {
+		if err := conn.SetDeadline(time.Now().Add(config.Timeout)); err != nil {
+			return nil, err
+		}
+		defer func() {
+			conn.SetDeadline(time.Time{})
+		}()
+	}
+	// end of the new code
+	c, chans, reqs, err := ssh.NewClientConn(conn, addr, config)
+	if err != nil {
+		return nil, err
+	}
+	return ssh.NewClient(c, chans, reqs), nil
+}
 
 type Client struct {
 	server Server
@@ -91,7 +127,7 @@ func (c *Client) dialOnReboot(prevUptime time.Time) error {
 		// waitConfig is not well honored by golang, it is
 		// set to 5sec above but in reality it takes ~60sec
 		// before the code times out.
-		sshc, err := ssh.Dial("tcp", c.addr, &waitConfig)
+		sshc, err := sshDial("tcp", c.addr, &waitConfig)
 		if err == nil {
 			// once successfully connected, check uptime to
 			// see if the reboot actually happend