Merge pull request #2 from trzsz/main

Implement workaround in WriteKnownHost for IPv6 addresses. Closes #1

knownhosts.go

@@ -3,10 +3,12 @@
 package knownhosts
 
 import (
+	"encoding/base64"
 	"errors"
 	"io"
 	"net"
 	"sort"
+	"strings"
 
 	"golang.org/x/crypto/ssh"
 	xknownhosts "golang.org/x/crypto/ssh/knownhosts"
@@ -42,9 +44,7 @@ func (hkcb HostKeyCallback) HostKeys(hostWithPort string) (keys []ssh.PublicKey)
 	placeholderPubKey := &fakePublicKey{}
 	var kkeys []xknownhosts.KnownKey
 	if hkcbErr := hkcb(hostWithPort, placeholderAddr, placeholderPubKey); errors.As(hkcbErr, &keyErr) {
-		for _, knownKey := range keyErr.Want {
-			kkeys = append(kkeys, knownKey)
-		}
+		kkeys = append(kkeys, keyErr.Want...)
 		knownKeyLess := func(i, j int) bool {
 			if kkeys[i].Filename < kkeys[j].Filename {
 				return true
@@ -98,6 +98,36 @@ func IsHostUnknown(err error) bool {
 	return errors.As(err, &keyErr) && len(keyErr.Want) == 0
 }
 
+// Normalize normalizes an address into the form used in known_hosts
+func Normalize(address string) string {
+	host, port, err := net.SplitHostPort(address)
+	if err != nil {
+		host = address
+		port = "22"
+	}
+	entry := host
+	if port != "22" {
+		entry = "[" + entry + "]:" + port
+	} else if strings.HasPrefix(host, "[") && strings.HasSuffix(host, "]") {
+		entry = entry[1 : len(entry)-1]
+	}
+	return entry
+}
+
+// Line returns a line to append to the known_hosts files.
+func Line(addresses []string, key ssh.PublicKey) string {
+	var trimmed []string
+	for _, a := range addresses {
+		trimmed = append(trimmed, Normalize(a))
+	}
+
+	return strings.Join([]string{
+		strings.Join(trimmed, ","),
+		key.Type(),
+		base64.StdEncoding.EncodeToString(key.Marshal()),
+	}, " ")
+}
+
 // WriteKnownHost writes a known_hosts line to writer for the supplied hostname,
 // remote, and key. This is useful when writing a custom hostkey callback which
 // wraps a callback obtained from knownhosts.New to provide additional
@@ -108,11 +138,11 @@ func WriteKnownHost(w io.Writer, hostname string, remote net.Addr, key ssh.Publi
 	// and doesn't normalize to the same string as hostname.
 	addresses := []string{hostname}
 	remoteStr := remote.String()
-	remoteStrNormalized := xknownhosts.Normalize(remoteStr)
-	if remoteStrNormalized != "[0.0.0.0]:0" && remoteStrNormalized != xknownhosts.Normalize(hostname) {
+	remoteStrNormalized := Normalize(remoteStr)
+	if remoteStrNormalized != "[0.0.0.0]:0" && remoteStrNormalized != Normalize(hostname) {
 		addresses = append(addresses, remoteStr)
 	}
-	line := xknownhosts.Line(addresses, key) + "\n"
+	line := Line(addresses, key) + "\n"
 	_, err := w.Write([]byte(line))
 	return err
 }

knownhosts_test.go

@@ -1,6 +1,7 @@
 package knownhosts
 
 import (
+	"bytes"
 	"crypto/ecdsa"
 	"crypto/ed25519"
 	"crypto/elliptic"
@@ -139,6 +140,75 @@ func TestIsHostUnknown(t *testing.T) {
 	}
 }
 
+func TestNormalize(t *testing.T) {
+	for in, want := range map[string]string{
+		"127.0.0.1":                 "127.0.0.1",
+		"127.0.0.1:22":              "127.0.0.1",
+		"[127.0.0.1]:22":            "127.0.0.1",
+		"[127.0.0.1]:23":            "[127.0.0.1]:23",
+		"127.0.0.1:23":              "[127.0.0.1]:23",
+		"[a.b.c]:22":                "a.b.c",
+		"abcd::abcd:abcd:abcd":      "abcd::abcd:abcd:abcd",
+		"[abcd::abcd:abcd:abcd]":    "abcd::abcd:abcd:abcd",
+		"[abcd::abcd:abcd:abcd]:22": "abcd::abcd:abcd:abcd",
+		"[abcd::abcd:abcd:abcd]:23": "[abcd::abcd:abcd:abcd]:23",
+	} {
+		got := Normalize(in)
+		if got != want {
+			t.Errorf("Normalize(%q) = %q, want %q", in, got, want)
+		}
+	}
+}
+
+func TestLine(t *testing.T) {
+	edKeyStr := "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIF9Wn63tLEhSWl9Ye+4x2GnruH8cq0LIh2vum/fUHrFQ"
+	edKey, _, _, _, err := ssh.ParseAuthorizedKey([]byte(edKeyStr))
+	if err != nil {
+		t.Fatalf("Unable to parse authorized key: %v", err)
+	}
+	for in, want := range map[string]string{
+		"server.org":                             "server.org " + edKeyStr,
+		"server.org:22":                          "server.org " + edKeyStr,
+		"server.org:23":                          "[server.org]:23 " + edKeyStr,
+		"[c629:1ec4:102:304:102:304:102:304]:22": "c629:1ec4:102:304:102:304:102:304 " + edKeyStr,
+		"[c629:1ec4:102:304:102:304:102:304]:23": "[c629:1ec4:102:304:102:304:102:304]:23 " + edKeyStr,
+	} {
+		if got := Line([]string{in}, edKey); got != want {
+			t.Errorf("Line(%q) = %q, want %q", in, got, want)
+		}
+	}
+}
+
+func TestWriteKnownHost(t *testing.T) {
+	edKeyStr := "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIF9Wn63tLEhSWl9Ye+4x2GnruH8cq0LIh2vum/fUHrFQ"
+	edKey, _, _, _, err := ssh.ParseAuthorizedKey([]byte(edKeyStr))
+	if err != nil {
+		t.Fatalf("Unable to parse authorized key: %v", err)
+	}
+	for _, m := range []struct {
+		hostname   string
+		remoteAddr string
+		want       string
+	}{
+		{hostname: "::1", remoteAddr: "[::1]:22", want: "::1 " + edKeyStr + "\n"},
+		{hostname: "127.0.0.1", remoteAddr: "127.0.0.1:22", want: "127.0.0.1 " + edKeyStr + "\n"},
+		{hostname: "ipv4.test", remoteAddr: "192.168.0.1:23", want: "ipv4.test,[192.168.0.1]:23 " + edKeyStr + "\n"},
+		{hostname: "ipv6.test", remoteAddr: "[ff01::1234]:23", want: "ipv6.test,[ff01::1234]:23 " + edKeyStr + "\n"},
+	} {
+		remote, err := net.ResolveTCPAddr("tcp", m.remoteAddr)
+		if err != nil {
+			t.Fatalf("Unable to resolve tcp addr: %v", err)
+		}
+		var got bytes.Buffer
+		if err = WriteKnownHost(&got, m.hostname, remote, edKey); err != nil {
+			t.Fatalf("Unable to write known host: %v", err)
+		}
+		if got.String() != m.want {
+			t.Errorf("WriteKnownHost(%q) = %q, want %q", m.hostname, got.String(), m.want)
+		}
+	}
+}
+
 // writeTestKnownHosts generates the test known_hosts file and returns the
 // file path to it. The generated file contains several hosts with a mix of
 // key types; each known host has between 1 and 3 different known host keys.
@@ -151,6 +221,7 @@ func writeTestKnownHosts(t *testing.T) string {
 		"only-ed25519.example.test:22": {generagePubKeyEd25519(t)},
 		"multi.example.test:2233":      {generatePubKeyRSA(t), generatePubKeyECDSA(t), generagePubKeyEd25519(t)},
 		"192.168.1.102:2222":           {generatePubKeyECDSA(t), generagePubKeyEd25519(t)},
+		"[fe80::abc:abc:abcd:abcd]:22": {generagePubKeyEd25519(t), generatePubKeyRSA(t)},
 	}
 
 	dir := t.TempDir()