From d2a34d5c8225d6aaaee287ce3ea8b218fbe210d4 Mon Sep 17 00:00:00 2001
From: Pavel Repin <prepin@gmail.com>
Date: Tue, 11 Jul 2023 18:17:13 -0700
Subject: [PATCH] ssh: wrap errors from client handshake

When an error is returned by a user defined host key callback,
it is now possible to handle it using standard Go mechanisms
such as errors.Is or errors.As.

Fixes golang/go#61309
---
 ssh/client.go      |  2 +-
 ssh/client_test.go | 25 +++++++++++++++++++++++--
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/ssh/client.go b/ssh/client.go
index bdc356cbdf..fd8c49749e 100644
--- a/ssh/client.go
+++ b/ssh/client.go
@@ -82,7 +82,7 @@ func NewClientConn(c net.Conn, addr string, config *ClientConfig) (Conn, <-chan
 
 	if err := conn.clientHandshake(addr, &fullConf); err != nil {
 		c.Close()
-		return nil, nil, nil, fmt.Errorf("ssh: handshake failed: %v", err)
+		return nil, nil, nil, fmt.Errorf("ssh: handshake failed: %w", err)
 	}
 	conn.mux = newMux(conn.transport)
 	return conn, conn.mux.incomingChannels, conn.mux.incomingRequests, nil
diff --git a/ssh/client_test.go b/ssh/client_test.go
index 281475596c..6618e4924f 100644
--- a/ssh/client_test.go
+++ b/ssh/client_test.go
@@ -7,6 +7,9 @@ package ssh
 import (
 	"bytes"
 	"crypto/rand"
+	"errors"
+	"fmt"
+	"net"
 	"strings"
 	"testing"
 )
@@ -207,9 +210,12 @@ func TestBannerCallback(t *testing.T) {
 }
 
 func TestNewClientConn(t *testing.T) {
+	errHostKeyMismatch := errors.New("host key mismatch")
+
 	for _, tt := range []struct {
-		name string
-		user string
+		name                    string
+		user                    string
+		simulateHostKeyMismatch HostKeyCallback
 	}{
 		{
 			name: "good user field for ConnMetadata",
@@ -219,6 +225,13 @@ func TestNewClientConn(t *testing.T) {
 			name: "empty user field for ConnMetadata",
 			user: "",
 		},
+		{
+			name: "host key mismatch",
+			user: "testuser",
+			simulateHostKeyMismatch: func(hostname string, remote net.Addr, key PublicKey) error {
+				return fmt.Errorf("%w: %s", errHostKeyMismatch, bytes.TrimSpace(MarshalAuthorizedKey(key)))
+			},
+		},
 	} {
 		t.Run(tt.name, func(t *testing.T) {
 			c1, c2, err := netPipe()
@@ -243,8 +256,16 @@ func TestNewClientConn(t *testing.T) {
 				},
 				HostKeyCallback: InsecureIgnoreHostKey(),
 			}
+
+			if tt.simulateHostKeyMismatch != nil {
+				clientConf.HostKeyCallback = tt.simulateHostKeyMismatch
+			}
+
 			clientConn, _, _, err := NewClientConn(c2, "", clientConf)
 			if err != nil {
+				if tt.simulateHostKeyMismatch != nil && errors.Is(err, errHostKeyMismatch) {
+					return
+				}
 				t.Fatal(err)
 			}