Fix hosts not being added to known hosts file

2 files changed, 7 insertions, 10 deletions

diff --git a/examples/client/client.go b/examples/client/client.go
index e6be51d..387fb47 100644
--- a/examples/client/client.go
+++ b/examples/client/client.go
@@ -37,7 +37,7 @@ func init() {
 					return nil
 				} else if userTrustsCertificatePermanently() {
 					// Add the certificate to the known hosts file
-					knownHosts.Add(cert)
+					knownHosts.Add(req.Hostname(), cert)
 					return nil
 				}
 			}
diff --git a/tofu.go b/tofu.go
index 8143099..946d2b4 100644
--- a/tofu.go
+++ b/tofu.go
@@ -60,15 +60,12 @@ func (k *KnownHosts) LoadFrom(path string) error {
 
 // Add adds a certificate to the list of known hosts.
 // If KnownHosts was loaded from a file, Add will append to the file.
-func (k *KnownHosts) Add(cert *x509.Certificate) {
-	// Add an entry per hostname
-	for _, name := range cert.DNSNames {
-		host := NewKnownHost(name, cert)
-		k.hosts = append(k.hosts, host)
-		// Append to the file
-		if k.file != nil {
-			host.Write(k.file)
-		}
+func (k *KnownHosts) Add(hostname string, cert *x509.Certificate) {
+	host := NewKnownHost(hostname, cert)
+	k.hosts = append(k.hosts, host)
+	// Append to the file
+	if k.file != nil {
+		host.Write(k.file)
 	}
 }