From b2d3a6a4b4d36521cd7f653879cf6981e7c5c340 Mon Sep 17 00:00:00 2001
From: Nicola Murino <nicola.murino@gmail.com>
Date: Sun, 4 Aug 2024 16:01:34 +0200
Subject: [PATCH] ssh/agent: ensure to not add duplicated keys

When adding a new key, if we already have a Signer with the same public
key, we now replace it with the new one instead of duplicating it.

Before this change we had this:

$ ssh-add -l
3072 SHA256:bsBRHC/xgiqBJdSuvSTNpJNLTISP/G356jNMCRYC5Es nicola@p1 (RSA)
3072 SHA256:bsBRHC/xgiqBJdSuvSTNpJNLTISP/G356jNMCRYC5Es nicola@p1 (RSA-CERT)

$ ssh-add /home/nicola/ssh_certs/id_rsa
Identity added: /home/nicola/ssh_certs/id_rsa (nicola@p1)
Certificate added: /home/nicola/ssh_certs/id_rsa-cert.pub (myid)

$ ssh-add -l
3072 SHA256:bsBRHC/xgiqBJdSuvSTNpJNLTISP/G356jNMCRYC5Es nicola@p1 (RSA)
3072 SHA256:bsBRHC/xgiqBJdSuvSTNpJNLTISP/G356jNMCRYC5Es nicola@p1 (RSA-CERT)
3072 SHA256:bsBRHC/xgiqBJdSuvSTNpJNLTISP/G356jNMCRYC5Es nicola@p1 (RSA)
3072 SHA256:bsBRHC/xgiqBJdSuvSTNpJNLTISP/G356jNMCRYC5Es nicola@p1 (RSA-CERT)

Change-Id: Iad1b1a6dc94f68f53f05d7d1172f0017839976fc
Reviewed-on: https://go-review.googlesource.com/c/crypto/+/602955
Reviewed-by: Filippo Valsorda <filippo@golang.org>
Auto-Submit: Nicola Murino <nicola.murino@gmail.com>
Reviewed-by: David Chase <drchase@google.com>
Reviewed-by: Michael Knyszek <mknyszek@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
---
 ssh/agent/keyring.go      |  9 ++++++++
 ssh/agent/keyring_test.go | 46 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/ssh/agent/keyring.go b/ssh/agent/keyring.go
index 21bfa87..c1b4361 100644
--- a/ssh/agent/keyring.go
+++ b/ssh/agent/keyring.go
@@ -175,6 +175,15 @@ func (r *keyring) Add(key AddedKey) error {
 		p.expire = &t
 	}
 
+	// If we already have a Signer with the same public key, replace it with the
+	// new one.
+	for idx, k := range r.keys {
+		if bytes.Equal(k.signer.PublicKey().Marshal(), p.signer.PublicKey().Marshal()) {
+			r.keys[idx] = p
+			return nil
+		}
+	}
+
 	r.keys = append(r.keys, p)
 
 	return nil
diff --git a/ssh/agent/keyring_test.go b/ssh/agent/keyring_test.go
index e5d50e7..e9c90a3 100644
--- a/ssh/agent/keyring_test.go
+++ b/ssh/agent/keyring_test.go
@@ -29,6 +29,10 @@ func validateListedKeys(t *testing.T, a Agent, expectedKeys []string) {
 		t.Fatalf("failed to list keys: %v", err)
 		return
 	}
+	if len(listedKeys) != len(expectedKeys) {
+		t.Fatalf("expeted %d key, got %d", len(expectedKeys), len(listedKeys))
+		return
+	}
 	actualKeys := make(map[string]bool)
 	for _, key := range listedKeys {
 		actualKeys[key.Comment] = true
@@ -74,3 +78,45 @@ func TestKeyringAddingAndRemoving(t *testing.T) {
 	}
 	validateListedKeys(t, k, []string{})
 }
+
+func TestAddDuplicateKey(t *testing.T) {
+	keyNames := []string{"rsa", "user"}
+
+	k := NewKeyring()
+	for _, keyName := range keyNames {
+		addTestKey(t, k, keyName)
+	}
+	validateListedKeys(t, k, keyNames)
+	// Add the keys again.
+	for _, keyName := range keyNames {
+		addTestKey(t, k, keyName)
+	}
+	validateListedKeys(t, k, keyNames)
+	// Add an existing key with an updated comment.
+	keyName := keyNames[0]
+	addedKey := AddedKey{
+		PrivateKey: testPrivateKeys[keyName],
+		Comment:    "comment updated",
+	}
+	err := k.Add(addedKey)
+	if err != nil {
+		t.Fatalf("failed to add key %q: %v", keyName, err)
+	}
+	// Check the that key is found and the comment was updated.
+	keys, err := k.List()
+	if err != nil {
+		t.Fatalf("failed to list keys: %v", err)
+	}
+	if len(keys) != len(keyNames) {
+		t.Fatalf("expected %d keys, got %d", len(keyNames), len(keys))
+	}
+	isFound := false
+	for _, key := range keys {
+		if key.Comment == addedKey.Comment {
+			isFound = true
+		}
+	}
+	if !isFound {
+		t.Fatal("key with the updated comment not found")
+	}
+}