Update AddHostKey to avoid always appending

Author: belak

server.go

@@ -315,6 +315,18 @@ func (srv *Server) ListenAndServe() error {
 func (srv *Server) AddHostKey(key Signer) {
 	// these are later added via AddHostKey on ServerConfig, which performs the
 	// check for one of every algorithm.
+
+	// This check is based on the AddHostKey method from the x/crypto/ssh
+	// library. This allows us to only keep one active key for each type on a
+	// server at once. So, if you're dynamically updating keys at runtime, this
+	// list will not keep growing.
+	for i, k := range srv.HostSigners {
+		if k.PublicKey().Type() == key.PublicKey().Type() {
+			srv.HostSigners[i] = key
+			return
+		}
+	}
+
 	srv.HostSigners = append(srv.HostSigners, key)
 }
 

server_test.go

@@ -8,6 +8,26 @@ import (
 	"time"
 )
 
+func TestAddHostKey(t *testing.T) {
+	s := Server{}
+	signer, err := generateSigner()
+	if err != nil {
+		t.Fatal(err)
+	}
+	s.AddHostKey(signer)
+	if len(s.HostSigners) != 1 {
+		t.Fatal("Key was not properly added")
+	}
+	signer, err = generateSigner()
+	if err != nil {
+		t.Fatal(err)
+	}
+	s.AddHostKey(signer)
+	if len(s.HostSigners) != 1 {
+		t.Fatal("Key was not properly replaced")
+	}
+}
+
 func TestServerShutdown(t *testing.T) {
 	l := newLocalListener()
 	testBytes := []byte("Hello world\n")