fix: write i/o timeout on ssh connection

Author: HalloTschuess

client.go

@@ -211,7 +211,6 @@ func (c *Client) messageHandler() {
 	for {
 		if c.scanner.Scan() {
 			line := c.scanner.Text()
-			//nolint: gocritic
 			if line == "error id=0 msg=ok" {
 				c.err <- nil
 			} else if matches := respTrailerRe.FindStringSubmatch(line); len(matches) == 4 {
@@ -253,9 +252,6 @@ func (c *Client) workHandler() {
 }
 
 func (c *Client) process(data string) {
-	if err := c.conn.SetWriteDeadline(time.Now().Add(c.timeout)); err != nil {
-		c.err <- err
-	}
 	if _, err := c.conn.Write([]byte(data)); err != nil {
 		c.err <- err
 	}

client_test.go

@@ -128,7 +128,7 @@ func TestClientWriteFail(t *testing.T) {
 	if !assert.NoError(t, err) {
 		return
 	}
-	assert.NoError(t, c.conn.(*legacyConnection).Conn.(*net.TCPConn).CloseWrite())
+	assert.NoError(t, c.conn.(*legacyConnection).Conn.(*writeTimeoutConn).Conn.(*net.TCPConn).CloseWrite())
 
 	_, err = c.Exec("version")
 	assert.Error(t, err)

connection.go

@@ -20,6 +20,21 @@ const (
 	DefaultSSHPort = 10022
 )
 
+type writeTimeoutConn struct {
+	net.Conn
+	timeout time.Duration
+}
+
+func (c *writeTimeoutConn) Write(p []byte) (n int, err error) {
+	if err = c.Conn.SetWriteDeadline(time.Now().Add(c.timeout)); err != nil {
+		return 0, fmt.Errorf("writeTimeoutConn: SetWriteDeadline: %w", err)
+	}
+	if n, err = c.Conn.Write(p); err != nil {
+		return n, fmt.Errorf("writeTimeoutConn: write: %w", err)
+	}
+	return n, nil
+}
+
 // legacyConnection is an insecure TCP connection.
 type legacyConnection struct {
 	net.Conn
@@ -32,10 +47,16 @@ func (c *legacyConnection) Connect(addr string, timeout time.Duration) error {
 		return err
 	}
 
-	c.Conn, err = net.DialTimeout("tcp", addr, timeout)
+	conn, err := net.DialTimeout("tcp", addr, timeout)
 	if err != nil {
 		return fmt.Errorf("legacy connection: dial: %w", err)
 	}
+
+	c.Conn = &writeTimeoutConn{
+		Conn:    conn,
+		timeout: timeout,
+	}
+
 	return nil
 }
 
@@ -53,10 +74,16 @@ func (c *sshConnection) Connect(addr string, timeout time.Duration) error {
 		return err
 	}
 
-	if c.Conn, err = net.DialTimeout("tcp", addr, timeout); err != nil {
+	conn, err := net.DialTimeout("tcp", addr, timeout)
+	if err != nil {
 		return fmt.Errorf("ssh connection: dial: %w", err)
 	}
 
+	c.Conn = &writeTimeoutConn{
+		Conn:    conn,
+		timeout: timeout,
+	}
+
 	clientConn, chans, reqs, err := ssh.NewClientConn(c.Conn, addr, c.config)
 	if err != nil {
 		return fmt.Errorf("ssh connecion: ssh client conn: %w", err)