From 2f3083f6163ef51179ad42ed523a18c9a1141467 Mon Sep 17 00:00:00 2001
From: hanwen <hanwen@google.com>
Date: Fri, 31 Jul 2015 22:59:06 +0200
Subject: [PATCH] crypto/ssh: allow client to specify host key algorithms.

Fixes golang/go#11722.

Change-Id: I4fa2a1db14050151f9269427ca35cf7ebd21440a
Reviewed-on: https://go-review.googlesource.com/12907
Reviewed-by: Adam Langley <agl@golang.org>
---
 ssh/client.go         |  7 ++++++
 ssh/handshake.go      | 16 +++++++++++--
 ssh/handshake_test.go |  1 +
 ssh/session_test.go   | 54 +++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 76 insertions(+), 2 deletions(-)

diff --git a/ssh/client.go b/ssh/client.go
index 72bd27f1ca..0b9fbe5002 100644
--- a/ssh/client.go
+++ b/ssh/client.go
@@ -203,4 +203,11 @@ type ClientConfig struct {
 	// ClientVersion contains the version identification string that will
 	// be used for the connection. If empty, a reasonable default is used.
 	ClientVersion string
+
+	// HostKeyAlgorithms lists the key types that the client will
+	// accept from the server as host key, in order of
+	// preference. If empty, a reasonable default is used. Any
+	// string returned from PublicKey.Type method may be used, or
+	// any of the CertAlgoXxxx and KeyAlgoXxxx constants.
+	HostKeyAlgorithms []string
 }
diff --git a/ssh/handshake.go b/ssh/handshake.go
index 4acc1a0a72..50461beb1b 100644
--- a/ssh/handshake.go
+++ b/ssh/handshake.go
@@ -59,7 +59,14 @@ type handshakeTransport struct {
 	serverVersion []byte
 	clientVersion []byte
 
-	hostKeys []Signer // If hostKeys are given, we are the server.
+	// hostKeys is non-empty if we are the server. In that case,
+	// it contains all host keys that can be used to sign the
+	// connection.
+	hostKeys []Signer
+
+	// hostKeyAlgorithms is non-empty if we are the client. In that case,
+	// we accept these key types from the server as host key.
+	hostKeyAlgorithms []string
 
 	// On read error, incoming is closed, and readError is set.
 	incoming  chan []byte
@@ -98,6 +105,11 @@ func newClientTransport(conn keyingTransport, clientVersion, serverVersion []byt
 	t.dialAddress = dialAddr
 	t.remoteAddr = addr
 	t.hostKeyCallback = config.HostKeyCallback
+	if config.HostKeyAlgorithms != nil {
+		t.hostKeyAlgorithms = config.HostKeyAlgorithms
+	} else {
+		t.hostKeyAlgorithms = supportedHostKeyAlgos
+	}
 	go t.readLoop()
 	return t
 }
@@ -234,7 +246,7 @@ func (t *handshakeTransport) sendKexInitLocked() (*kexInitMsg, []byte, error) {
 				msg.ServerHostKeyAlgos, k.PublicKey().Type())
 		}
 	} else {
-		msg.ServerHostKeyAlgos = supportedHostKeyAlgos
+		msg.ServerHostKeyAlgos = t.hostKeyAlgorithms
 	}
 	packet := Marshal(msg)
 
diff --git a/ssh/handshake_test.go b/ssh/handshake_test.go
index f690627572..32023715eb 100644
--- a/ssh/handshake_test.go
+++ b/ssh/handshake_test.go
@@ -69,6 +69,7 @@ func handshakePair(clientConf *ClientConfig, addr string) (client *handshakeTran
 
 	serverConf := &ServerConfig{}
 	serverConf.AddHostKey(testSigners["ecdsa"])
+	serverConf.AddHostKey(testSigners["rsa"])
 	serverConf.SetDefaults()
 	server = newServerTransport(trS, v, v, serverConf)
 
diff --git a/ssh/session_test.go b/ssh/session_test.go
index 7ce44f59c9..f7f0f7642e 100644
--- a/ssh/session_test.go
+++ b/ssh/session_test.go
@@ -718,3 +718,57 @@ func TestInvalidServerConfiguration(t *testing.T) {
 		t.Fatalf("NewServerConn attempted to Read() from Conn while configuration is missing authentication method")
 	}
 }
+
+func TestHostKeyAlgorithms(t *testing.T) {
+	serverConf := &ServerConfig{
+		NoClientAuth: true,
+	}
+	serverConf.AddHostKey(testSigners["rsa"])
+	serverConf.AddHostKey(testSigners["ecdsa"])
+
+	connect := func(clientConf *ClientConfig, want string) {
+		var alg string
+		clientConf.HostKeyCallback = func(h string, a net.Addr, key PublicKey) error {
+			alg = key.Type()
+			return nil
+		}
+		c1, c2, err := netPipe()
+		if err != nil {
+			t.Fatalf("netPipe: %v", err)
+		}
+		defer c1.Close()
+		defer c2.Close()
+
+		go NewServerConn(c1, serverConf)
+		_, _, _, err = NewClientConn(c2, "", clientConf)
+		if err != nil {
+			t.Fatalf("NewClientConn: %v", err)
+		}
+		if alg != want {
+			t.Errorf("selected key algorithm %s, want %s", alg, want)
+		}
+	}
+
+	// By default, we get the preferred algorithm, which is ECDSA 256.
+
+	clientConf := &ClientConfig{}
+	connect(clientConf, KeyAlgoECDSA256)
+
+	// Client asks for RSA explicitly.
+	clientConf.HostKeyAlgorithms = []string{KeyAlgoRSA}
+	connect(clientConf, KeyAlgoRSA)
+
+	c1, c2, err := netPipe()
+	if err != nil {
+		t.Fatalf("netPipe: %v", err)
+	}
+	defer c1.Close()
+	defer c2.Close()
+
+	go NewServerConn(c1, serverConf)
+	clientConf.HostKeyAlgorithms = []string{"nonexistent-hostkey-algo"}
+	_, _, _, err = NewClientConn(c2, "", clientConf)
+	if err == nil {
+		t.Fatal("succeeded connecting with unknown hostkey algorithm")
+	}
+}