From baa92cefa3a90f7fbc33b7b920b758658f0067a3 Mon Sep 17 00:00:00 2001
From: Diogo Monica <diogo@docker.com>
Date: Sun, 8 Nov 2015 17:20:11 -0800
Subject: [PATCH] Fixed panic on listKeys with invalid keys, added tests

Signed-off-by: David Lawrence <david.lawrence@docker.com>

Signed-off-by: Diogo Monica <diogo@docker.com> (github: endophage)
---
 trustmanager/keyfilestore.go      | 18 ++++++++--
 trustmanager/keyfilestore_test.go | 55 +++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/trustmanager/keyfilestore.go b/trustmanager/keyfilestore.go
index ed8339cfe1..977e3e6983 100644
--- a/trustmanager/keyfilestore.go
+++ b/trustmanager/keyfilestore.go
@@ -212,14 +212,28 @@ func listKeys(s LimitedFileStore) map[string]string {
 	keyIDMap := make(map[string]string)
 
 	for _, f := range s.ListFiles() {
+		// Remove the prefix of the directory from the filename
 		if f[:len(rootKeysSubdir)] == rootKeysSubdir {
 			f = strings.TrimPrefix(f, rootKeysSubdir+"/")
 		} else {
 			f = strings.TrimPrefix(f, nonRootKeysSubdir+"/")
 		}
+
+		// Remove the extension from the full filename
+		// abcde_root.key becomes abcde_root
 		keyIDFull := strings.TrimSpace(strings.TrimSuffix(f, filepath.Ext(f)))
-		keyID := keyIDFull[:strings.LastIndex(keyIDFull, "_")]
-		keyAlias := keyIDFull[strings.LastIndex(keyIDFull, "_")+1:]
+
+		// If the key does not have a _, it is malformed
+		underscoreIndex := strings.LastIndex(keyIDFull, "_")
+		if underscoreIndex == -1 {
+			continue
+		}
+
+		// The keyID is the first part of the keyname
+		// The KeyAlias is the second part of the keyname
+		// in a key named abcde_root, abcde is the keyID and root is the KeyAlias
+		keyID := keyIDFull[:underscoreIndex]
+		keyAlias := keyIDFull[underscoreIndex+1:]
 		keyIDMap[keyID] = keyAlias
 	}
 	return keyIDMap
diff --git a/trustmanager/keyfilestore_test.go b/trustmanager/keyfilestore_test.go
index fa418bf5a6..16e39f31f5 100644
--- a/trustmanager/keyfilestore_test.go
+++ b/trustmanager/keyfilestore_test.go
@@ -114,6 +114,60 @@ EMl3eFOJXjIch/wIesRSN+2dGOsl7neercjMh1i9RvpCwHDx/E0=
 	assert.Equal(t, testData, pemPrivKey)
 }
 
+func TestListKeys(t *testing.T) {
+	testName := "docker.com/notary/root"
+	perms := os.FileMode(0755)
+
+	// Temporary directory where test files will be created
+	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
+	assert.NoError(t, err, "failed to create a temporary directory")
+	defer os.RemoveAll(tempBaseDir)
+
+	// Create our store
+	store, err := NewKeyFileStore(tempBaseDir, passphraseRetriever)
+	assert.NoError(t, err, "failed to create new key filestore")
+
+	privKey, err := GenerateECDSAKey(rand.Reader)
+	assert.NoError(t, err, "could not generate private key")
+
+	// Call the AddKey function
+	err = store.AddKey(testName, "root", privKey)
+	assert.NoError(t, err, "failed to add key to store")
+
+	// Check to see if the keystore lists this key
+	keyMap := store.ListKeys()
+
+	// Expect to see exactly one key in the map
+	assert.Len(t, keyMap, 1)
+	// Expect to see privKeyID inside of the map
+	role, ok := keyMap[testName]
+	assert.True(t, ok)
+	assert.Equal(t, role, "root")
+
+	// Call the AddKey function for the second key
+	err = store.AddKey(testName+"2", "targets", privKey)
+	assert.NoError(t, err, "failed to add key to store")
+
+	// Check to see if the keystore lists this key
+	keyMap = store.ListKeys()
+
+	// Expect to see exactly two keys in the map
+	assert.Len(t, keyMap, 2)
+	// Expect to see privKeyID2 inside of the map
+	role, ok = keyMap[testName+"2"]
+	assert.True(t, ok)
+	assert.Equal(t, role, "targets")
+
+	// Write an invalid filename to the directory
+	filePath := filepath.Join(tempBaseDir, rootKeysSubdir, "fakekeyname.key")
+	err = ioutil.WriteFile(filePath, []byte("data"), perms)
+	assert.NoError(t, err, "failed to write test file")
+
+	// Check to see if the keystore still lists two keys
+	keyMap = store.ListKeys()
+	assert.Len(t, keyMap, 2)
+}
+
 func TestAddGetKeyMemStore(t *testing.T) {
 	testName := "docker.com/notary/root"
 	testAlias := "root"
@@ -136,6 +190,7 @@ func TestAddGetKeyMemStore(t *testing.T) {
 	assert.Equal(t, retrievedKey.Public(), privKey.Public())
 	assert.Equal(t, retrievedKey.Private(), privKey.Private())
 }
+
 func TestGetDecryptedWithTamperedCipherText(t *testing.T) {
 	testExt := "key"
 	testAlias := "root"