add tests for initial keystore state, and after removing and adding

Signed-off-by: Riyaz Faizullabhoy <riyaz.faizullabhoy@docker.com>

client/client_test.go

@@ -391,15 +391,15 @@ func requireRepoHasExpectedKeys(t *testing.T, repo *NotaryRepository,
 	require.NoError(t, err)
 
 	roles := make(map[string]bool)
-	for keyID, role := range ks.ListKeys() {
+	for keyID, keyInfo := range ks.ListKeys() {
-		if role == data.CanonicalRootRole {
+		if keyInfo.Role == data.CanonicalRootRole {
-			require.Equal(t, rootKeyID, keyID, "Unexpected root key ID")
+			assert.Equal(t, rootKeyID, keyID, "Unexpected root key ID")
 		}
 		// just to ensure the content of the key files created are valid
 		_, r, err := ks.GetKey(keyID)
 		require.NoError(t, err)
-		require.Equal(t, role, r)
+		require.Equal(t, keyInfo.Role, r)
-		roles[role] = true
+		roles[keyInfo.Role] = true
 	}
 	// there is a root key and a targets key
 	alwaysThere := []string{data.CanonicalRootRole, data.CanonicalTargetsRole}

cmd/notary/keys.go

@@ -435,10 +435,10 @@ func removeKeyInteractively(keyStores []trustmanager.KeyStore, keyID string,
 	var storesByIndex []trustmanager.KeyStore
 
 	for _, store := range keyStores {
-		for keypath, role := range store.ListKeys() {
+		for keypath, keyInfo := range store.ListKeys() {
 			if filepath.Base(keypath) == keyID {
 				foundKeys = append(foundKeys,
-					[]string{keypath, role, store.Name()})
+					[]string{keypath, keyInfo.Role, store.Name()})
 				storesByIndex = append(storesByIndex, store)
 			}
 		}

cmd/notary/prettyprint.go

@@ -6,7 +6,6 @@ import (
 	"fmt"
 	"io"
 	"math"
-	"path/filepath"
 	"sort"
 	"strings"
 	"time"
@@ -96,19 +95,12 @@ func prettyPrintKeys(keyStores []trustmanager.KeyStore, writer io.Writer) {
 	var info []keyInfo
 
 	for _, store := range keyStores {
-		for keyPath, role := range store.ListKeys() {
+		for keyID, keyIDInfo := range store.ListKeys() {
-			gun := ""
-			if role != data.CanonicalRootRole {
-				dirPath := filepath.Dir(keyPath)
-				if dirPath != "." { // no gun
-					gun = dirPath
-				}
-			}
 			info = append(info, keyInfo{
-				role:     role,
+				role:     keyIDInfo.Role,
 				location: store.Name(),
-				gun:      gun,
+				gun:      keyIDInfo.Gun,
-				keyID:    filepath.Base(keyPath),
+				keyID:    keyID,
 			})
 		}
 	}

cmd/notary/prettyprint_test.go

@@ -93,8 +93,8 @@ func TestPrettyPrintRootAndSigningKeys(t *testing.T) {
 
 	longNameShortened := "..." + strings.Repeat("z", 37)
 
-	keys := make([]data.PrivateKey, 3)
+	keys := make([]data.PrivateKey, 4)
-	for i := 0; i < 3; i++ {
+	for i := 0; i < 4; i++ {
 		key, err := trustmanager.GenerateED25519Key(rand.Reader)
 		assert.NoError(t, err)
 		keys[i] = key
@@ -105,9 +105,9 @@ func TestPrettyPrintRootAndSigningKeys(t *testing.T) {
 	// add keys to the key stores
 	assert.NoError(t, keyStores[0].AddKey(keys[0].ID(), root, keys[0]))
 	assert.NoError(t, keyStores[1].AddKey(keys[0].ID(), root, keys[0]))
-	assert.NoError(t, keyStores[0].AddKey(strings.Repeat("a/", 30)+keys[0].ID(), "targets", keys[0]))
+	assert.NoError(t, keyStores[0].AddKey(strings.Repeat("a/", 30)+keys[1].ID(), "targets", keys[1]))
-	assert.NoError(t, keyStores[1].AddKey("short/gun/"+keys[0].ID(), "snapshot", keys[0]))
+	assert.NoError(t, keyStores[1].AddKey("short/gun/"+keys[1].ID(), "snapshot", keys[1]))
-	assert.NoError(t, keyStores[0].AddKey(keys[1].ID(), "targets/a", keys[1]))
+	assert.NoError(t, keyStores[0].AddKey(keys[3].ID(), "targets/a", keys[3]))
 	assert.NoError(t, keyStores[0].AddKey(keys[2].ID(), "invalidRole", keys[2]))
 
 	expected := [][]string{
@@ -116,10 +116,10 @@ func TestPrettyPrintRootAndSigningKeys(t *testing.T) {
 		{root, keys[0].ID(), longNameShortened},
 		// these have no gun, so they come first
 		{"invalidRole", keys[2].ID(), keyStores[0].Name()},
-		{"targets/a", keys[1].ID(), keyStores[0].Name()},
+		{"targets/a", keys[3].ID(), keyStores[0].Name()},
 		// these have guns, and are sorted then by guns
-		{"targets", "..." + strings.Repeat("/a", 11), keys[0].ID(), keyStores[0].Name()},
+		{"targets", "..." + strings.Repeat("/a", 11), keys[1].ID(), keyStores[0].Name()},
-		{"snapshot", "short/gun", keys[0].ID(), longNameShortened},
+		{"snapshot", "short/gun", keys[1].ID(), longNameShortened},
 	}
 
 	var b bytes.Buffer

cryptoservice/crypto_service.go

@@ -121,7 +121,7 @@ func (cs *CryptoService) ListKeys(role string) []string {
 	var res []string
 	for _, ks := range cs.keyStores {
 		for k, r := range ks.ListKeys() {
-			if r == role {
+			if r.Role == role {
 				res = append(res, k)
 			}
 		}
@@ -134,7 +134,7 @@ func (cs *CryptoService) ListAllKeys() map[string]string {
 	res := make(map[string]string)
 	for _, ks := range cs.keyStores {
 		for k, r := range ks.ListKeys() {
-			res[k] = r // keys are content addressed so don't care about overwrites
+			res[k] = r.Role // keys are content addressed so don't care about overwrites
 		}
 	}
 	return res

cryptoservice/import_export.go

@@ -191,7 +191,6 @@ func (cs *CryptoService) ImportKeysZip(zipReader zip.Reader) error {
 	// Iterate through the files in the archive. Don't add the keys
 	for _, f := range zipReader.File {
 		fNameTrimmed := strings.TrimSuffix(f.Name, filepath.Ext(f.Name))
-
 		rc, err := f.Open()
 		if err != nil {
 			return err
@@ -214,9 +213,6 @@ func (cs *CryptoService) ImportKeysZip(zipReader zip.Reader) error {
 	}
 
 	for keyName, pemBytes := range newKeys {
-		if keyName[len(keyName)-5:] == "_root" {
-			keyName = "root"
-		}
 		// try to import the key to all key stores. As long as one of them
 		// succeeds, consider it a success
 		var tmpErr error
@@ -271,18 +267,18 @@ func (cs *CryptoService) ExportKeysByGUN(dest io.Writer, gun string, passphraseR
 }
 
 func moveKeysByGUN(oldKeyStore, newKeyStore trustmanager.KeyStore, gun string) error {
-	for relKeyPath := range oldKeyStore.ListKeys() {
+	for keyID, keyInfo := range oldKeyStore.ListKeys() {
 		// Skip keys that aren't associated with this GUN
-		if !strings.HasPrefix(relKeyPath, filepath.FromSlash(gun)) {
+		if keyInfo.Gun != gun {
 			continue
 		}
 
-		privKey, alias, err := oldKeyStore.GetKey(relKeyPath)
+		privKey, alias, err := oldKeyStore.GetKey(keyID)
 		if err != nil {
 			return err
 		}
 
-		err = newKeyStore.AddKey(relKeyPath, alias, privKey)
+		err = newKeyStore.AddKey(filepath.Join(keyInfo.Gun, keyID), alias, privKey)
 		if err != nil {
 			return err
 		}
@@ -292,13 +288,13 @@ func moveKeysByGUN(oldKeyStore, newKeyStore trustmanager.KeyStore, gun string) e
 }
 
 func moveKeys(oldKeyStore, newKeyStore trustmanager.KeyStore) error {
-	for f := range oldKeyStore.ListKeys() {
+	for keyID, keyInfo := range oldKeyStore.ListKeys() {
-		privateKey, role, err := oldKeyStore.GetKey(f)
+		privateKey, role, err := oldKeyStore.GetKey(keyID)
 		if err != nil {
 			return err
 		}
 
-		err = newKeyStore.AddKey(f, role, privateKey)
+		err = newKeyStore.AddKey(filepath.Join(keyInfo.Gun, keyID), role, privateKey)
 
 		if err != nil {
 			return err

cryptoservice/import_export_compatibility_test.go

@@ -27,7 +27,10 @@ func TestImport0Dot1Zip(t *testing.T) {
 	zipWriter.Close()
 	zipFile.Close()
 
-	origKeys := ks.ListKeys()
+	origKeys := make(map[string]string)
+	for keyID, keyInfo := range ks.ListKeys() {
+		origKeys[keyID] = keyInfo.Role
+	}
 	assert.Len(t, origKeys, 3)
 
 	// now import the zip file into a new cryptoservice
@@ -89,10 +92,12 @@ func importExportedZip(t *testing.T, original *CryptoService,
 	zipFile, err := ioutil.TempFile("", "notary-test-zipFile")
 	defer os.RemoveAll(zipFile.Name())
 	if gun != "" {
-		original.ExportKeysByGUN(zipFile, gun, ret)
+		err = original.ExportKeysByGUN(zipFile, gun, ret)
+		assert.NoError(t, err)
 		cs = NewCryptoService(gun, ks)
 	} else {
-		original.ExportAllKeys(zipFile, ret)
+		err = original.ExportAllKeys(zipFile, ret)
+		assert.NoError(t, err)
 		cs = NewCryptoService(original.gun, ks)
 	}
 	zipFile.Close()
@@ -122,9 +127,9 @@ func TestImportExport0Dot1GUNKeys(t *testing.T) {
 	// remove root from expected key list, because root is not exported when
 	// we export by gun
 	expectedKeys := make(map[string]string)
-	for keyID, role := range ks.ListKeys() {
+	for keyID, keyInfo := range ks.ListKeys() {
-		if role != data.CanonicalRootRole {
+		if keyInfo.Role != data.CanonicalRootRole {
-			expectedKeys[keyID] = role
+			expectedKeys[keyID] = keyInfo.Role
 		}
 	}
 

cryptoservice/import_export_test.go

@@ -199,7 +199,7 @@ func TestImportExportGUN(t *testing.T) {
 		if alias == data.CanonicalRootRole {
 			continue
 		}
-		relKeyPath := filepath.Join("tuf_keys", privKeyName+".key")
+		relKeyPath := filepath.Join("tuf_keys", gun, privKeyName+".key")
 
 		passphraseByFile[relKeyPath] = exportPassphrase
 	}
@@ -258,7 +258,7 @@ func TestImportExportGUN(t *testing.T) {
 		if alias == data.CanonicalRootRole {
 			continue
 		}
-		relKeyPath := filepath.Join("tuf_keys", privKeyName+".key")
+		relKeyPath := filepath.Join("tuf_keys", gun, privKeyName+".key")
 		privKeyFileName := filepath.Join(tempBaseDir2, "private", relKeyPath)
 		_, err = os.Stat(privKeyFileName)
 		assert.NoError(t, err)

trustmanager/keyfilestore.go

@@ -19,7 +19,7 @@ const (
 	privDir           = "private"
 )
 
-type keyIDMap map[string]KeyInfo
+type keyInfoMap map[string]KeyInfo
 
 // KeyFileStore persists and manages private keys on disk
 type KeyFileStore struct {
@@ -27,7 +27,7 @@ type KeyFileStore struct {
 	SimpleFileStore
 	passphrase.Retriever
 	cachedKeys map[string]*cachedKey
-	keyIDMap
+	keyInfoMap
 }
 
 // KeyMemoryStore manages private keys in memory
@@ -36,13 +36,14 @@ type KeyMemoryStore struct {
 	MemoryFileStore
 	passphrase.Retriever
 	cachedKeys map[string]*cachedKey
-	keyIDMap
+	keyInfoMap
 }
 
-// KeyInfo stores the role and gun for a corresponding private key
+// KeyInfo stores the role, path, and gun for a corresponding private key ID
+// It is assumed that each private key ID is unique
 type KeyInfo struct {
-	role string
+	Gun  string
-	gun  string
+	Role string
 }
 
 // NewKeyFileStore returns a new KeyFileStore creating a private directory to
@@ -54,26 +55,61 @@ func NewKeyFileStore(baseDir string, passphraseRetriever passphrase.Retriever) (
 		return nil, err
 	}
 	cachedKeys := make(map[string]*cachedKey)
-	keyIDMap := make(keyIDMap)
+	keyInfoMap := make(keyInfoMap)
 
 	keyStore := &KeyFileStore{SimpleFileStore: *fileStore,
 		Retriever:  passphraseRetriever,
 		cachedKeys: cachedKeys,
-		keyIDMap:   keyIDMap,
+		keyInfoMap: keyInfoMap,
 	}
 
 	// Load this keystore's ID --> gun/role map
-	keyStore.loadKeyIDInfo()
+	keyStore.loadKeyInfo()
-
 	return keyStore, nil
 }
 
-func generateKeyInfoMap(fullIDToRole map[string]string) map[string]KeyInfo {
+func generateKeyInfoMap(s LimitedFileStore) map[string]KeyInfo {
 	keyInfoMap := make(map[string]KeyInfo)
-	for keyIDAndGun, role := range fullIDToRole {
+	for _, keyPath := range s.ListFiles() {
+		// Remove the prefix of the directory from the filename for GUN/role/ID parsing
+		var keyIDAndGun, keyRole string
+		if strings.HasPrefix(keyPath, rootKeysSubdir+"/") {
+			keyIDAndGun = strings.TrimPrefix(keyPath, rootKeysSubdir+"/")
+		} else {
+			keyIDAndGun = strings.TrimPrefix(keyPath, nonRootKeysSubdir+"/")
+		}
+
+		// Separate the ID and GUN (can be empty) from the filepath
+		keyIDAndGun = strings.TrimSpace(keyIDAndGun)
 		keyGun := getGunFromFullID(keyIDAndGun)
 		keyID := filepath.Base(keyIDAndGun)
-		keyInfoMap[keyID] = KeyInfo{gun: keyGun, role: role}
+
+		// If the key does not have a _, we'll attempt to
+		// read it as a PEM
+		underscoreIndex := strings.LastIndex(keyID, "_")
+		if underscoreIndex == -1 {
+			d, err := s.Get(keyPath)
+			if err != nil {
+				logrus.Error(err)
+				continue
+			}
+			block, _ := pem.Decode(d)
+			if block == nil {
+				continue
+			}
+			if role, ok := block.Headers["role"]; ok {
+				keyRole = role
+			}
+		} else {
+			// This is the legacy KEYID_ROLE filename
+			// The keyID is the first part of the keyname
+			// The keyRole is the second part of the keyname
+			// in a key named abcde_root, abcde is the keyID and root is the KeyAlias
+			legacyID := keyID
+			keyID = legacyID[:underscoreIndex]
+			keyRole = legacyID[underscoreIndex+1:]
+		}
+		keyInfoMap[keyID] = KeyInfo{Gun: keyGun, Role: keyRole}
 	}
 	return keyInfoMap
 }
@@ -87,12 +123,12 @@ func getGunFromFullID(fullKeyID string) string {
 	return keyGun
 }
 
-func (s *KeyFileStore) loadKeyIDInfo() {
+func (s *KeyFileStore) loadKeyInfo() {
-	s.keyIDMap = generateKeyInfoMap(s.ListKeys())
+	s.keyInfoMap = generateKeyInfoMap(s)
 }
 
-func (s *KeyMemoryStore) loadKeyIDInfo() {
+func (s *KeyMemoryStore) loadKeyInfo() {
-	s.keyIDMap = generateKeyInfoMap(s.ListKeys())
+	s.keyInfoMap = generateKeyInfoMap(s)
 }
 
 // Name returns a user friendly name for the location this store
@@ -109,7 +145,7 @@ func (s *KeyFileStore) AddKey(name, role string, privKey data.PrivateKey) error
 	if err != nil {
 		return err
 	}
-	s.keyIDMap[privKey.ID()] = KeyInfo{gun: getGunFromFullID(name), role: role}
+	s.keyInfoMap[privKey.ID()] = KeyInfo{Gun: getGunFromFullID(name), Role: role}
 	return nil
 }
 
@@ -118,15 +154,15 @@ func (s *KeyFileStore) GetKey(name string) (data.PrivateKey, string, error) {
 	s.Lock()
 	defer s.Unlock()
 	// If this is a bare key ID without the gun, prepend the gun so the filestore lookup succeeds
-	if keyInfo, ok := s.keyIDMap[name]; ok {
+	if keyInfo, ok := s.keyInfoMap[name]; ok {
-		name = filepath.Join(keyInfo.gun, name)
+		name = filepath.Join(keyInfo.Gun, name)
 	}
 	return getKey(s, s.Retriever, s.cachedKeys, name)
 }
 
 // ListKeys returns a list of unique PublicKeys present on the KeyFileStore.
-func (s *KeyFileStore) ListKeys() map[string]string {
+func (s *KeyFileStore) ListKeys() map[string]KeyInfo {
-	return listKeys(s)
+	return s.keyInfoMap
 }
 
 // RemoveKey removes the key from the keyfilestore
@@ -134,19 +170,19 @@ func (s *KeyFileStore) RemoveKey(name string) error {
 	s.Lock()
 	defer s.Unlock()
 	// If this is a bare key ID without the gun, prepend the gun so the filestore lookup succeeds
-	if keyInfo, ok := s.keyIDMap[name]; ok {
+	if keyInfo, ok := s.keyInfoMap[name]; ok {
-		name = filepath.Join(keyInfo.gun, name)
+		name = filepath.Join(keyInfo.Gun, name)
 	}
 	err := removeKey(s, s.cachedKeys, name)
 	if err != nil {
 		return err
 	}
 	// Remove this key from our keyInfo map if we removed from our filesystem
-	if _, ok := s.keyIDMap[name]; ok {
+	if _, ok := s.keyInfoMap[name]; ok {
-		delete(s.keyIDMap, name)
+		delete(s.keyInfoMap, name)
 	} else {
 		// This might be of the form GUN/ID  - try to delete without the gun
-		delete(s.keyIDMap, filepath.Base(name))
+		delete(s.keyInfoMap, filepath.Base(name))
 	}
 	return nil
 }
@@ -163,9 +199,15 @@ func (s *KeyFileStore) ExportKey(name string) ([]byte, error) {
 
 // ImportKey imports the private key in the encrypted bytes into the keystore
 // with the given key ID and alias.
-// This is only used for root, so no need to touch the keyIDMap
+// This is only used for root, so no need to touch the keyInfoMap
 func (s *KeyFileStore) ImportKey(pemBytes []byte, alias string) error {
-	return importKey(s, s.Retriever, s.cachedKeys, alias, pemBytes)
+	err := importKey(s, s.Retriever, s.cachedKeys, alias, pemBytes)
+	if err != nil {
+		return err
+	}
+	keyID, keyInfo := getImportedKeyInfo(pemBytes, alias)
+	s.keyInfoMap[keyID] = keyInfo
+	return nil
 }
 
 // NewKeyMemoryStore returns a new KeyMemoryStore which holds keys in memory
@@ -173,17 +215,16 @@ func NewKeyMemoryStore(passphraseRetriever passphrase.Retriever) *KeyMemoryStore
 	memStore := NewMemoryFileStore()
 	cachedKeys := make(map[string]*cachedKey)
 
-	keyIDMap := make(keyIDMap)
+	keyInfoMap := make(keyInfoMap)
 
 	keyStore := &KeyMemoryStore{MemoryFileStore: *memStore,
 		Retriever:  passphraseRetriever,
 		cachedKeys: cachedKeys,
-		keyIDMap:   keyIDMap,
+		keyInfoMap: keyInfoMap,
 	}
 
 	// Load this keystore's ID --> gun/role map
-	keyStore.loadKeyIDInfo()
+	keyStore.loadKeyInfo()
-
 	return keyStore
 }
 
@@ -201,7 +242,7 @@ func (s *KeyMemoryStore) AddKey(name, alias string, privKey data.PrivateKey) err
 	if err != nil {
 		return err
 	}
-	s.keyIDMap[privKey.ID()] = KeyInfo{gun: getGunFromFullID(name), role: alias}
+	s.keyInfoMap[privKey.ID()] = KeyInfo{Gun: getGunFromFullID(name), Role: alias}
 	return nil
 }
 
@@ -210,15 +251,15 @@ func (s *KeyMemoryStore) GetKey(name string) (data.PrivateKey, string, error) {
 	s.Lock()
 	defer s.Unlock()
 	// If this is a bare key ID without the gun, prepend the gun so the filestore lookup succeeds
-	if keyInfo, ok := s.keyIDMap[name]; ok {
+	if keyInfo, ok := s.keyInfoMap[name]; ok {
-		name = filepath.Join(keyInfo.gun, name)
+		name = filepath.Join(keyInfo.Gun, name)
 	}
 	return getKey(s, s.Retriever, s.cachedKeys, name)
 }
 
 // ListKeys returns a list of unique PublicKeys present on the KeyFileStore.
-func (s *KeyMemoryStore) ListKeys() map[string]string {
+func (s *KeyMemoryStore) ListKeys() map[string]KeyInfo {
-	return listKeys(s)
+	return s.keyInfoMap
 }
 
 // RemoveKey removes the key from the keystore
@@ -226,19 +267,19 @@ func (s *KeyMemoryStore) RemoveKey(name string) error {
 	s.Lock()
 	defer s.Unlock()
 	// If this is a bare key ID without the gun, prepend the gun so the filestore lookup succeeds
-	if keyInfo, ok := s.keyIDMap[name]; ok {
+	if keyInfo, ok := s.keyInfoMap[name]; ok {
-		name = filepath.Join(keyInfo.gun, name)
+		name = filepath.Join(keyInfo.Gun, name)
 	}
 	err := removeKey(s, s.cachedKeys, name)
 	if err != nil {
 		return err
 	}
 	// Remove this key from our keyInfo map if we removed from our filesystem
-	if _, ok := s.keyIDMap[name]; ok {
+	if _, ok := s.keyInfoMap[name]; ok {
-		delete(s.keyIDMap, name)
+		delete(s.keyInfoMap, name)
 	} else {
 		// This might be of the form GUN/ID  - try to delete without the gun
-		delete(s.keyIDMap, filepath.Base(name))
+		delete(s.keyInfoMap, filepath.Base(name))
 	}
 	return nil
 }
@@ -255,9 +296,45 @@ func (s *KeyMemoryStore) ExportKey(name string) ([]byte, error) {
 
 // ImportKey imports the private key in the encrypted bytes into the keystore
 // with the given key ID and alias.
-// This is only used for root, so no need to touch the keyIDMap
 func (s *KeyMemoryStore) ImportKey(pemBytes []byte, alias string) error {
-	return importKey(s, s.Retriever, s.cachedKeys, alias, pemBytes)
+	err := importKey(s, s.Retriever, s.cachedKeys, alias, pemBytes)
+	if err != nil {
+		return err
+	}
+	keyID, keyInfo := getImportedKeyInfo(pemBytes, alias)
+	s.keyInfoMap[keyID] = keyInfo
+	return nil
+}
+
+func getImportedKeyInfo(pemBytes []byte, filename string) (string, KeyInfo) {
+	var keyID, gun, role string
+	// Try to read the role and gun from the filepath and PEM headers
+	if filepath.HasPrefix(filename, "root_keys") {
+		role = data.CanonicalRootRole
+		gun = ""
+		filename = strings.TrimPrefix(filename, "root_keys")
+	} else {
+		filename = strings.TrimPrefix(filename, "tuf_keys")
+		gun = getGunFromFullID(filename)
+	}
+	keyIDFull := filepath.Base(filename)
+	underscoreIndex := strings.LastIndex(keyIDFull, "_")
+	if underscoreIndex == -1 {
+		block, _ := pem.Decode(pemBytes)
+		if keyRole, ok := block.Headers["role"]; ok {
+			role = keyRole
+		}
+		keyID = filepath.Base(keyIDFull)
+	} else {
+		// This is the legacy KEYID_ROLE filename
+		// The keyID is the first part of the keyname
+		// The keyRole is the second part of the keyname
+		// in a key named abcde_root, abcde is the keyID and root is the KeyAlias
+		legacyID := keyIDFull
+		keyID = legacyID[:underscoreIndex]
+		role = legacyID[underscoreIndex+1:]
+	}
+	return keyID, KeyInfo{Gun: gun, Role: role}
 }
 
 func addKey(s LimitedFileStore, passphraseRetriever passphrase.Retriever, cachedKeys map[string]*cachedKey, name, role string, privKey data.PrivateKey) error {
@@ -339,50 +416,6 @@ func getKey(s LimitedFileStore, passphraseRetriever passphrase.Retriever, cached
 	return privKey, keyAlias, nil
 }
 
-// ListKeys returns a map of unique PublicKeys present on the KeyFileStore and
-// their corresponding aliases.
-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
-		var keyIDFull string
-		if strings.HasPrefix(f, notary.RootKeysSubdir+"/") {
-			keyIDFull = strings.TrimPrefix(f, notary.RootKeysSubdir+"/")
-		} else {
-			keyIDFull = strings.TrimPrefix(f, notary.NonRootKeysSubdir+"/")
-		}
-
-		keyIDFull = strings.TrimSpace(keyIDFull)
-
-		// If the key does not have a _, we'll attempt to
-		// read it as a PEM
-		underscoreIndex := strings.LastIndex(keyIDFull, "_")
-		if underscoreIndex == -1 {
-			d, err := s.Get(f)
-			if err != nil {
-				logrus.Error(err)
-				continue
-			}
-			block, _ := pem.Decode(d)
-			if block == nil {
-				continue
-			}
-			if role, ok := block.Headers["role"]; ok {
-				keyIDMap[keyIDFull] = role
-			}
-		} else {
-			// 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
-}
-
 // RemoveKey removes the key from the keyfilestore
 func removeKey(s LimitedFileStore, cachedKeys map[string]*cachedKey, name string) error {
 	role, legacy, err := getKeyRole(s, name)

trustmanager/keyfilestore_test.go

@@ -62,14 +62,99 @@ func testAddKeyWithRole(t *testing.T, role, expectedSubdir string) {
 	assert.Contains(t, string(b), "-----BEGIN EC PRIVATE KEY-----")
 
 	// Check that we have the role and gun info for this key's ID
-	keyInfo, ok := store.keyIDMap[privKey.ID()]
+	keyInfo, ok := store.keyInfoMap[privKey.ID()]
 	assert.True(t, ok)
-	assert.Equal(t, role, keyInfo.role)
+	assert.Equal(t, role, keyInfo.Role)
 	if role != data.CanonicalRootRole {
-		assert.Equal(t, filepath.Dir(testName), keyInfo.gun)
+		assert.Equal(t, filepath.Dir(testName), keyInfo.Gun)
 	}
 }
 
+func TestKeyStoreInternalState(t *testing.T) {
+	// 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)
+
+	gun := "docker.com/notary"
+
+	// Mimic a notary repo setup, and test that bringing up a keyfilestore creates the correct keyInfoMap
+	roles := []string{data.CanonicalRootRole, data.CanonicalTargetsRole, data.CanonicalSnapshotRole, "targets/delegation"}
+	// Keep track of the key IDs for each role, so we can validate later against the keystore state
+	roleToID := make(map[string]string)
+	for _, role := range roles {
+		// generate a key for the role
+		privKey, err := GenerateECDSAKey(rand.Reader)
+		assert.NoError(t, err, "could not generate private key")
+
+		// generate the correct PEM role header
+		privKeyPEM, err := KeyToPEM(privKey, role)
+		assert.NoError(t, err, "could not generate PEM")
+
+		// write the key file to the correct location
+		// Pretend our GUN is docker.com/notary
+		keyPath := filepath.Join(tempBaseDir, "private", getSubdir(role))
+		if role == data.CanonicalTargetsRole || role == data.CanonicalSnapshotRole {
+			keyPath = filepath.Join(keyPath, gun)
+		}
+		keyPath = filepath.Join(keyPath, privKey.ID())
+		assert.NoError(t, os.MkdirAll(filepath.Dir(keyPath), 0755))
+		assert.NoError(t, ioutil.WriteFile(keyPath+".key", privKeyPEM, 0755))
+
+		roleToID[role] = privKey.ID()
+	}
+
+	store, err := NewKeyFileStore(tempBaseDir, passphraseRetriever)
+	assert.NoError(t, err)
+	assert.Len(t, store.keyInfoMap, 4)
+	for _, role := range roles {
+		keyID, _ := roleToID[role]
+		// make sure this keyID is the right length
+		assert.Len(t, keyID, notary.Sha256HexSize)
+		assert.Equal(t, role, store.keyInfoMap[keyID].Role)
+		// targets and snapshot keys should have a gun set, root and delegation keys should not
+		if role == data.CanonicalTargetsRole || role == data.CanonicalSnapshotRole {
+			assert.Equal(t, gun, store.keyInfoMap[keyID].Gun)
+		} else {
+			assert.Empty(t, store.keyInfoMap[keyID].Gun)
+		}
+	}
+
+	// Try removing the targets key only by ID (no gun provided)
+	assert.NoError(t, store.RemoveKey(roleToID[data.CanonicalTargetsRole]))
+	// The key file itself should have been removed
+	_, err = os.Stat(filepath.Join(tempBaseDir, "private", "tuf_keys", gun, roleToID[data.CanonicalTargetsRole]+".key"))
+	assert.Error(t, err)
+	// The keyInfoMap should have also updated by deleting the key
+	_, ok := store.keyInfoMap[roleToID[data.CanonicalTargetsRole]]
+	assert.False(t, ok)
+
+	// Try removing the delegation key only by ID (no gun provided)
+	assert.NoError(t, store.RemoveKey(roleToID["targets/delegation"]))
+	// The key file itself should have been removed
+	_, err = os.Stat(filepath.Join(tempBaseDir, "private", "tuf_keys", roleToID["targets/delegation"]+".key"))
+	assert.Error(t, err)
+	// The keyInfoMap should have also updated
+	_, ok = store.keyInfoMap[roleToID["targets/delegation"]]
+	assert.False(t, ok)
+
+	// Try removing the root key only by ID (no gun provided)
+	assert.NoError(t, store.RemoveKey(roleToID[data.CanonicalRootRole]))
+	// The key file itself should have been removed
+	_, err = os.Stat(filepath.Join(tempBaseDir, "private", "root_keys", roleToID[data.CanonicalRootRole]+".key"))
+	assert.Error(t, err)
+	// The keyInfoMap should have also updated_
+	_, ok = store.keyInfoMap[roleToID[data.CanonicalRootRole]]
+	assert.False(t, ok)
+
+	// Generate a new targets key and add it with its gun, check that the map gets updated back
+	privKey, err := GenerateECDSAKey(rand.Reader)
+	assert.NoError(t, err, "could not generate private key")
+	assert.NoError(t, store.AddKey(filepath.Join(gun, privKey.ID()), data.CanonicalTargetsRole, privKey))
+	assert.Equal(t, gun, store.keyInfoMap[privKey.ID()].Gun)
+	assert.Equal(t, data.CanonicalTargetsRole, store.keyInfoMap[privKey.ID()].Role)
+}
+
 func TestGet(t *testing.T) {
 	nonRootRolesToTest := []string{
 		data.CanonicalTargetsRole,
@@ -249,12 +334,13 @@ func TestListKeys(t *testing.T) {
 	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")
-
 	roles := append(data.BaseRoles, "targets/a", "invalidRoleName")
 
 	for i, role := range roles {
+		// Make a new key for each role
+		privKey, err := GenerateECDSAKey(rand.Reader)
+		assert.NoError(t, err, "could not generate private key")
+
 		// Call the AddKey function
 		keyName := fmt.Sprintf("%s%d", testName, i)
 		err = store.AddKey(keyName, role, privKey)
@@ -266,9 +352,9 @@ func TestListKeys(t *testing.T) {
 		// Expect to see exactly one key in the map
 		assert.Len(t, keyMap, i+1)
 		// Expect to see privKeyID inside of the map
-		listedRole, ok := keyMap[keyName]
+		listedInfo, ok := keyMap[privKey.ID()]
 		assert.True(t, ok)
-		assert.Equal(t, role, listedRole)
+		assert.Equal(t, role, listedInfo.Role)
 	}
 
 	// Write an invalid filename to the directory

trustmanager/keystore.go

@@ -44,7 +44,7 @@ type KeyStore interface {
 	// succeeds.  Otherwise, returns an error if it cannot add.
 	AddKey(name, alias string, privKey data.PrivateKey) error
 	GetKey(name string) (data.PrivateKey, string, error)
-	ListKeys() map[string]string
+	ListKeys() map[string]KeyInfo
 	RemoveKey(name string) error
 	ExportKey(name string) ([]byte, error)
 	ImportKey(pemBytes []byte, alias string) error

trustmanager/yubikey/yubikeystore.go

@@ -617,6 +617,7 @@ func (s *YubiKeyStore) setLibLoader(loader pkcs11LibLoader) {
 	s.libLoader = loader
 }
 
+// TODO: yubi key store refactor
 func (s *YubiKeyStore) ListKeys() map[string]string {
 	if len(s.keys) > 0 {
 		return buildKeyMap(s.keys)