refactor ImportRoleKey, simplify integration test code, update constants

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

cmd/notary/integration_test.go

@@ -20,6 +20,7 @@ import (
 
 	"github.com/Sirupsen/logrus"
 	ctxu "github.com/docker/distribution/context"
+	"github.com/docker/notary"
 	"github.com/docker/notary/cryptoservice"
 	"github.com/docker/notary/passphrase"
 	"github.com/docker/notary/server"
@@ -1069,6 +1070,21 @@ func TestClientKeyImportExportRootOnly(t *testing.T) {
 		t, tempDir, server.URL, "gun", target, tempFile.Name())
 }
 
+// Helper method to get the subdirectory for TUF keys
+func getKeySubdir(role, gun string) string {
+	subdir := notary.PrivDir
+	switch role {
+	case data.CanonicalRootRole:
+		return filepath.Join(subdir, notary.RootKeysSubdir)
+	case data.CanonicalTargetsRole:
+		return filepath.Join(subdir, notary.NonRootKeysSubdir, gun)
+	case data.CanonicalSnapshotRole:
+		return filepath.Join(subdir, notary.NonRootKeysSubdir, gun)
+	default:
+		return filepath.Join(subdir, notary.NonRootKeysSubdir)
+	}
+}
+
 // Tests importing and exporting keys for all different roles and GUNs
 func TestClientKeyImportExportAllRoles(t *testing.T) {
 	// -- setup --
@@ -1080,222 +1096,76 @@ func TestClientKeyImportExportAllRoles(t *testing.T) {
 	server := setupServer()
 	defer server.Close()
 
-	tempFile, err := ioutil.TempFile("", "pemfile")
-	assert.NoError(t, err)
-	// close later, because we might need to write to it
-	defer os.Remove(tempFile.Name())
-
-	privKey1, err := trustmanager.GenerateECDSAKey(rand.Reader)
-	assert.NoError(t, err)
-
-	privKey2, err := trustmanager.GenerateECDSAKey(rand.Reader)
-	assert.NoError(t, err)
-
-	privKey3, err := trustmanager.GenerateECDSAKey(rand.Reader)
-	assert.NoError(t, err)
-
-	privKey4, err := trustmanager.GenerateECDSAKey(rand.Reader)
-	assert.NoError(t, err)
-
 	// -- tests --
-	_, err = runCommand(t, tempDir, "-s", server.URL, "init", "gun")
+	_, err := runCommand(t, tempDir, "-s", server.URL, "init", "gun")
 	assert.NoError(t, err)
 
-	rootPemBytes, err := trustmanager.EncryptPrivateKey(privKey1, "root", testPassphrase)
+	testRoles := append(data.BaseRoles, "targets/releases")
-	assert.NoError(t, err)
+	// Test importing and exporting keys to all base roles and delegation role
-	ioutil.WriteFile(tempFile.Name(), rootPemBytes, 0644)
+	for _, role := range testRoles {
+		// Do this while importing keys that have the PEM header role set or have --role set on import
+		for _, setKeyRole := range []bool{true, false} {
+			// Make a new key for this role
+			privKey, err := trustmanager.GenerateECDSAKey(rand.Reader)
+			assert.NoError(t, err)
 
-	// Import from root, specified in PEM
+			// Make a tempfile for importing
-	_, err = runCommand(t, tempDir, "key", "import", tempFile.Name())
+			tempFile, err := ioutil.TempFile("", "pemfile")
-	assert.NoError(t, err)
+			assert.NoError(t, err)
-	_, err = os.Stat(filepath.Join(tempDir, "private", "root_keys", privKey1.ID()+".key"))
-	assert.Nil(t, err)
 
-	// Ensure exporting this key by ID gets the same key
+			// Specify the role in the PEM header
-	_, err = runCommand(t, tempDir, "key", "export", privKey1.ID(), tempFile.Name())
+			pemBytes, err := trustmanager.EncryptPrivateKey(privKey, role, testPassphrase)
-	assert.NoError(t, err)
+			assert.NoError(t, err)
-	// Compare the bytes of the exported file and the root key file in the repo
+			ioutil.WriteFile(tempFile.Name(), pemBytes, 0644)
-	exportedBytes, err := ioutil.ReadFile(tempFile.Name())
-	assert.NoError(t, err)
-	repoBytes, err := ioutil.ReadFile(filepath.Join(tempDir, "private", "root_keys", privKey1.ID()+".key"))
-	assert.NoError(t, err)
-	assert.Equal(t, repoBytes, exportedBytes)
 
-	targetsPemBytes, err := trustmanager.EncryptPrivateKey(privKey2, "targets", testPassphrase)
+			// If we need to set the key role with the --role flag, do so on import
-	assert.NoError(t, err)
+			if setKeyRole {
-	ioutil.WriteFile(tempFile.Name(), targetsPemBytes, 0644)
+				// If it's targets/snapshot we must specify the GUN
+				if role == data.CanonicalTargetsRole || role == data.CanonicalSnapshotRole {
+					_, err = runCommand(t, tempDir, "key", "import", tempFile.Name(), "--gun", "gun", "--role", role)
+				} else {
+					_, err = runCommand(t, tempDir, "key", "import", tempFile.Name(), "--role", role)
+				}
+			} else {
+				// If it's targets/snapshot we must specify the GUN
+				if role == data.CanonicalTargetsRole || role == data.CanonicalSnapshotRole {
+					_, err = runCommand(t, tempDir, "key", "import", tempFile.Name(), "--gun", "gun")
+				} else {
+					_, err = runCommand(t, tempDir, "key", "import", tempFile.Name())
+				}
+			}
+			assert.NoError(t, err)
 
-	// Import from snapshot, specified in PEM.  Must supply GUN
+			// Test that we imported correctly
-	_, err = runCommand(t, tempDir, "key", "import", tempFile.Name(), "--gun", "gun")
+			keySubdir := getKeySubdir(role, "gun")
-	assert.NoError(t, err)
+			_, err = os.Stat(filepath.Join(tempDir, keySubdir, privKey.ID()+".key"))
-	_, err = os.Stat(filepath.Join(tempDir, "private", "tuf_keys", "gun", privKey2.ID()+".key"))
+			assert.Nil(t, err)
-	assert.Nil(t, err)
 
-	// Ensure exporting this key by ID gets the same key
+			// Remove the input file so we can test exporting
-	_, err = runCommand(t, tempDir, "key", "export", privKey2.ID(), tempFile.Name())
+			assert.NoError(t, os.Remove(tempFile.Name()))
-	assert.NoError(t, err)
-	// Compare the bytes of the exported file and the targets key file in the repo
-	exportedBytes, err = ioutil.ReadFile(tempFile.Name())
-	assert.NoError(t, err)
-	repoBytes, err = ioutil.ReadFile(filepath.Join(tempDir, "private", "tuf_keys", "gun", privKey2.ID()+".key"))
-	assert.NoError(t, err)
-	assert.Equal(t, repoBytes, exportedBytes)
 
-	snapshotPemBytes, err := trustmanager.EncryptPrivateKey(privKey3, "snapshot", testPassphrase)
+			// Make a tempfile for exporting to
-	assert.NoError(t, err)
+			tempFile, err = ioutil.TempFile("", "pemfile")
-	ioutil.WriteFile(tempFile.Name(), snapshotPemBytes, 0644)
+			assert.NoError(t, err)
 
-	// Import from snapshot, specified in PEM.  Must supply GUN
+			// Ensure exporting this key by ID gets the same key
-	_, err = runCommand(t, tempDir, "key", "import", tempFile.Name(), "--gun", "gun")
+			_, err = runCommand(t, tempDir, "key", "export", privKey.ID(), tempFile.Name())
-	assert.NoError(t, err)
+			assert.NoError(t, err)
-	_, err = os.Stat(filepath.Join(tempDir, "private", "tuf_keys", "gun", privKey3.ID()+".key"))
+			// Compare the bytes of the exported file and the root key file in the repo
-	assert.Nil(t, err)
+			exportedBytes, err := ioutil.ReadFile(tempFile.Name())
+			assert.NoError(t, err)
+			repoBytes, err := ioutil.ReadFile(filepath.Join(tempDir, keySubdir, privKey.ID()+".key"))
+			assert.NoError(t, err)
+			assert.Equal(t, repoBytes, exportedBytes)
 
-	// Ensure exporting this key by ID gets the same key
+			// Ensure exporting this key and changing the passphrase works
-	_, err = runCommand(t, tempDir, "key", "export", privKey3.ID(), tempFile.Name())
+			_, err = runCommand(t, tempDir, "key", "export", privKey.ID(), tempFile.Name(), "-p")
-	assert.NoError(t, err)
+			assert.NoError(t, err)
-	// Compare the bytes of the exported file and the snapshot key file in the repo
-	exportedBytes, err = ioutil.ReadFile(tempFile.Name())
-	assert.NoError(t, err)
-	repoBytes, err = ioutil.ReadFile(filepath.Join(tempDir, "private", "tuf_keys", "gun", privKey3.ID()+".key"))
-	assert.NoError(t, err)
-	assert.Equal(t, repoBytes, exportedBytes)
 
-	delegationPemBytes, err := trustmanager.EncryptPrivateKey(privKey4, "targets/releases", testPassphrase)
+			// Remove the export file for cleanup
-	assert.NoError(t, err)
+			assert.NoError(t, os.Remove(tempFile.Name()))
-	ioutil.WriteFile(tempFile.Name(), delegationPemBytes, 0644)
+		}
-
+	}
-	// Import from delegation key, specified in PEM.  No GUN needed
-	_, err = runCommand(t, tempDir, "key", "import", tempFile.Name())
-	assert.NoError(t, err)
-	_, err = os.Stat(filepath.Join(tempDir, "private", "tuf_keys", privKey4.ID()+".key"))
-	assert.Nil(t, err)
-
-	// Ensure exporting this key by ID gets the same key
-	_, err = runCommand(t, tempDir, "key", "export", privKey4.ID(), tempFile.Name())
-	assert.NoError(t, err)
-	// Compare the bytes of the exported file and the delegation key file in the repo
-	exportedBytes, err = ioutil.ReadFile(tempFile.Name())
-	assert.NoError(t, err)
-	repoBytes, err = ioutil.ReadFile(filepath.Join(tempDir, "private", "tuf_keys", privKey4.ID()+".key"))
-	assert.NoError(t, err)
-	assert.Equal(t, repoBytes, exportedBytes)
-
-	rootPemBytes, err = trustmanager.EncryptPrivateKey(privKey1, "", testPassphrase)
-	assert.NoError(t, err)
-	ioutil.WriteFile(tempFile.Name(), rootPemBytes, 0644)
-
-	// Import from root, specified in flag only
-	_, err = runCommand(t, tempDir, "key", "import", tempFile.Name(), "--role", "root")
-	assert.NoError(t, err)
-	_, err = os.Stat(filepath.Join(tempDir, "private", "root_keys", privKey1.ID()+".key"))
-	assert.NoError(t, err)
-
-	// Assert the PEM role header is "root"
-	pemBytes, err := ioutil.ReadFile(filepath.Join(tempDir, "private", "root_keys", privKey1.ID()+".key"))
-	assert.NoError(t, err)
-	assert.Equal(t, "root", trustmanager.ReadRoleFromPEM(pemBytes))
-
-	// Ensure exporting this key by ID gets the same key
-	_, err = runCommand(t, tempDir, "key", "export", privKey1.ID(), tempFile.Name())
-	assert.NoError(t, err)
-	// Compare the bytes of the exported file and the root key file in the repo
-	exportedBytes, err = ioutil.ReadFile(tempFile.Name())
-	assert.NoError(t, err)
-	repoBytes, err = ioutil.ReadFile(filepath.Join(tempDir, "private", "root_keys", privKey1.ID()+".key"))
-	assert.NoError(t, err)
-	assert.Equal(t, repoBytes, exportedBytes)
-
-	// Ensure exporting this key and changing the passphrase works
-	_, err = runCommand(t, tempDir, "key", "export", privKey1.ID(), tempFile.Name(), "-p")
-	assert.NoError(t, err)
-
-	targetsPemBytes, err = trustmanager.EncryptPrivateKey(privKey2, "", testPassphrase)
-	assert.NoError(t, err)
-	ioutil.WriteFile(tempFile.Name(), targetsPemBytes, 0644)
-
-	// Import from snapshot, specified in flag.  Must supply GUN
-	_, err = runCommand(t, tempDir, "key", "import", tempFile.Name(), "--gun", "gun", "--role", "targets")
-	assert.NoError(t, err)
-	_, err = os.Stat(filepath.Join(tempDir, "private", "tuf_keys", "gun", privKey2.ID()+".key"))
-	assert.NoError(t, err)
-
-	// Assert the PEM role header is "targets"
-	pemBytes, err = ioutil.ReadFile(filepath.Join(tempDir, "private", "tuf_keys", "gun", privKey2.ID()+".key"))
-	assert.NoError(t, err)
-	assert.Equal(t, "targets", trustmanager.ReadRoleFromPEM(pemBytes))
-
-	// Ensure exporting this key by ID gets the same key
-	_, err = runCommand(t, tempDir, "key", "export", privKey2.ID(), tempFile.Name())
-	assert.NoError(t, err)
-	// Compare the bytes of the exported file and the targets key file in the repo
-	exportedBytes, err = ioutil.ReadFile(tempFile.Name())
-	assert.NoError(t, err)
-	repoBytes, err = ioutil.ReadFile(filepath.Join(tempDir, "private", "tuf_keys", "gun", privKey2.ID()+".key"))
-	assert.NoError(t, err)
-	assert.Equal(t, repoBytes, exportedBytes)
-
-	// Ensure exporting this key and changing the passphrase works
-	_, err = runCommand(t, tempDir, "key", "export", privKey2.ID(), tempFile.Name(), "-p")
-	assert.NoError(t, err)
-
-	snapshotPemBytes, err = trustmanager.EncryptPrivateKey(privKey3, "", testPassphrase)
-	assert.NoError(t, err)
-	ioutil.WriteFile(tempFile.Name(), snapshotPemBytes, 0644)
-
-	// Import from snapshot, specified in flags.  Must supply GUN
-	_, err = runCommand(t, tempDir, "key", "import", tempFile.Name(), "--gun", "gun", "--role", "snapshot")
-	assert.NoError(t, err)
-	_, err = os.Stat(filepath.Join(tempDir, "private", "tuf_keys", "gun", privKey3.ID()+".key"))
-	assert.NoError(t, err)
-
-	// Assert the PEM role header is "snapshot"
-	pemBytes, err = ioutil.ReadFile(filepath.Join(tempDir, "private", "tuf_keys", "gun", privKey3.ID()+".key"))
-	assert.NoError(t, err)
-	assert.Equal(t, "snapshot", trustmanager.ReadRoleFromPEM(pemBytes))
-
-	// Ensure exporting this key by ID gets the same key
-	_, err = runCommand(t, tempDir, "key", "export", privKey3.ID(), tempFile.Name())
-	assert.NoError(t, err)
-	// Compare the bytes of the exported file and the snapshot key file in the repo
-	exportedBytes, err = ioutil.ReadFile(tempFile.Name())
-	assert.NoError(t, err)
-	repoBytes, err = ioutil.ReadFile(filepath.Join(tempDir, "private", "tuf_keys", "gun", privKey3.ID()+".key"))
-	assert.NoError(t, err)
-	assert.Equal(t, repoBytes, exportedBytes)
-
-	// Ensure exporting this key and changing the passphrase works
-	_, err = runCommand(t, tempDir, "key", "export", privKey3.ID(), tempFile.Name(), "-p")
-	assert.NoError(t, err)
-
-	delegationPemBytes, err = trustmanager.EncryptPrivateKey(privKey4, "", testPassphrase)
-	assert.NoError(t, err)
-	ioutil.WriteFile(tempFile.Name(), delegationPemBytes, 0644)
-
-	// Import from delegation key, specified in flag.  No GUN needed
-	_, err = runCommand(t, tempDir, "key", "import", tempFile.Name(), "--role", "targets/delegation")
-	assert.NoError(t, err)
-	_, err = os.Stat(filepath.Join(tempDir, "private", "tuf_keys", privKey4.ID()+".key"))
-	assert.NoError(t, err)
-
-	// Assert the PEM role header is "targets/delegation"
-	pemBytes, err = ioutil.ReadFile(filepath.Join(tempDir, "private", "tuf_keys", privKey4.ID()+".key"))
-	assert.NoError(t, err)
-	assert.Equal(t, "targets/delegation", trustmanager.ReadRoleFromPEM(pemBytes))
-
-	// Ensure exporting this key by ID gets the same key
-	_, err = runCommand(t, tempDir, "key", "export", privKey4.ID(), tempFile.Name())
-	assert.NoError(t, err)
-	// Compare the bytes of the exported file and the delegation key file in the repo
-	exportedBytes, err = ioutil.ReadFile(tempFile.Name())
-	assert.NoError(t, err)
-	repoBytes, err = ioutil.ReadFile(filepath.Join(tempDir, "private", "tuf_keys", privKey4.ID()+".key"))
-	assert.NoError(t, err)
-	assert.Equal(t, repoBytes, exportedBytes)
-
-	// Ensure exporting this key and changing the passphrase works
-	_, err = runCommand(t, tempDir, "key", "export", privKey4.ID(), tempFile.Name(), "-p")
-	assert.NoError(t, err)
 }
 
 func assertNumCerts(t *testing.T, tempDir string, expectedNum int) []string {

cmd/notary/keys.go

@@ -275,11 +275,12 @@ func (k *keyCommander) keysExport(cmd *cobra.Command, args []string) error {
 		for keypath, role := range store.ListKeys() {
 			if filepath.Base(keypath) == keyID {
 				keyRole = role
-				if role != data.CanonicalRootRole {
+				if role == data.CanonicalRootRole {
-					dirPath := filepath.Dir(keypath)
+					continue
-					if dirPath != "." { // no gun
+				}
-						keyGun = dirPath
+				dirPath := filepath.Dir(keypath)
-					}
+				if dirPath != "." { // no gun
+					keyGun = dirPath
 				}
 				break
 			}
@@ -372,20 +373,6 @@ func (k *keyCommander) keysImport(cmd *cobra.Command, args []string) error {
 
 	pemRole := trustmanager.ReadRoleFromPEM(pemBytes)
 
-	// Rewind after reading the first time
-	_, err = importFile.Seek(0, 0)
-	if err != nil {
-		return fmt.Errorf("Error reading input file: %v", err)
-	}
-
-	if pemRole != "" && !data.ValidRole(pemRole) {
-		return fmt.Errorf("Invalid role specified for key: %s", pemRole)
-	}
-
-	if k.keysImportRole != "" && !data.ValidRole(k.keysImportRole) {
-		return fmt.Errorf("Invalid role specified for key: %s", k.keysImportRole)
-	}
-
 	// If the PEM key doesn't have a role in it, we must have --role set
 	if pemRole == "" && k.keysImportRole == "" {
 		return fmt.Errorf("Could not infer role, and no role was specified for key")
@@ -410,11 +397,7 @@ func (k *keyCommander) keysImport(cmd *cobra.Command, args []string) error {
 	}
 
 	cs := cryptoservice.NewCryptoService(k.keysImportGUN, ks...)
-	if importRole == data.CanonicalRootRole {
+	err = cs.ImportRoleKey(pemBytes, importRole, k.getRetriever())
-		err = cs.ImportRootKey(importFile)
-	} else {
-		err = cs.ImportRoleKey(importFile, importRole, k.getRetriever())
-	}
 
 	if err != nil {
 		return fmt.Errorf("Error importing root key: %v", err)

cmd/notary/keys_test.go

@@ -434,33 +434,6 @@ func TestChangeKeyPassphraseNonexistentID(t *testing.T) {
 	assert.Contains(t, err.Error(), "could not retrieve local key for key ID provided")
 }
 
-func TestKeyImportInvalidFlagRole(t *testing.T) {
-	k := &keyCommander{
-		configGetter:   func() (*viper.Viper, error) { return viper.New(), nil },
-		getRetriever:   func() passphrase.Retriever { return passphrase.ConstantRetriever("pass") },
-		keysImportRole: "invalid",
-	}
-	tempFileName := generateTempTestKeyFile(t, "")
-	defer os.Remove(tempFileName)
-
-	err := k.keysImport(&cobra.Command{}, []string{tempFileName})
-	assert.Error(t, err)
-	assert.Contains(t, err.Error(), "Invalid role specified for key:")
-}
-
-func TestKeyImportInvalidPEMRole(t *testing.T) {
-	k := &keyCommander{
-		configGetter: func() (*viper.Viper, error) { return viper.New(), nil },
-		getRetriever: func() passphrase.Retriever { return passphrase.ConstantRetriever("pass") },
-	}
-	tempFileName := generateTempTestKeyFile(t, "invalid")
-	defer os.Remove(tempFileName)
-
-	err := k.keysImport(&cobra.Command{}, []string{tempFileName})
-	assert.Error(t, err)
-	assert.Contains(t, err.Error(), "Invalid role specified for key:")
-}
-
 func TestKeyImportMismatchingRoles(t *testing.T) {
 	k := &keyCommander{
 		configGetter:   func() (*viper.Viper, error) { return viper.New(), nil },

cryptoservice/import_export.go

@@ -104,19 +104,19 @@ func (cs *CryptoService) ExportKeyReencrypt(dest io.Writer, keyID string, newPas
 // It prompts for the key's passphrase to verify the data and to determine
 // the key ID.
 func (cs *CryptoService) ImportRootKey(source io.Reader) error {
-	return cs.ImportRoleKey(source, data.CanonicalRootRole, nil)
-}
-
-// ImportRoleKey imports a private key in PEM format key from an io.Reader
-// It prompts for the key's passphrase to verify the data and to determine
-// the key ID.
-func (cs *CryptoService) ImportRoleKey(source io.Reader, role string, newPassphraseRetriever passphrase.Retriever) error {
 	pemBytes, err := ioutil.ReadAll(source)
-
 	if err != nil {
 		return err
 	}
+	return cs.ImportRoleKey(pemBytes, data.CanonicalRootRole, nil)
+}
+
+// ImportRoleKey imports a private key in PEM format key from a byte array
+// It prompts for the key's passphrase to verify the data and to determine
+// the key ID.
+func (cs *CryptoService) ImportRoleKey(pemBytes []byte, role string, newPassphraseRetriever passphrase.Retriever) error {
 	var alias string
+	var err error
 	if role == data.CanonicalRootRole {
 		alias = role
 		if err = checkRootKeyIsEncrypted(pemBytes); err != nil {

cryptoservice/import_export_test.go

@@ -77,7 +77,7 @@ func TestImportExportZip(t *testing.T) {
 		_, alias, err := cs.GetPrivateKey(privKeyName)
 		assert.NoError(t, err, "privKey %s has no alias", privKeyName)
 
-		if alias == "root" {
+		if alias == data.CanonicalRootRole {
 			continue
 		}
 		relKeyPath := filepath.Join("tuf_keys", privKeyName+".key")
@@ -138,7 +138,7 @@ func TestImportExportZip(t *testing.T) {
 		_, alias, err := cs2.GetPrivateKey(privKeyName)
 		assert.NoError(t, err, "privKey %s has no alias", privKeyName)
 
-		if alias == "root" {
+		if alias == data.CanonicalRootRole {
 			continue
 		}
 		relKeyPath := filepath.Join("tuf_keys", privKeyName+".key")
@@ -196,7 +196,7 @@ func TestImportExportGUN(t *testing.T) {
 	for privKeyName := range privKeyMap {
 		_, alias, err := cs.GetPrivateKey(privKeyName)
 		assert.NoError(t, err, "privKey %s has no alias", privKeyName)
-		if alias == "root" {
+		if alias == data.CanonicalRootRole {
 			continue
 		}
 		relKeyPath := filepath.Join("tuf_keys", privKeyName+".key")
@@ -250,12 +250,12 @@ func TestImportExportGUN(t *testing.T) {
 	// Look for keys in private. The filenames should match the key IDs
 	// in the repo's private key store.
 	for privKeyName, role := range privKeyMap {
-		if role == "root" {
+		if role == data.CanonicalRootRole {
 			continue
 		}
 		_, alias, err := cs2.GetPrivateKey(privKeyName)
 		assert.NoError(t, err, "privKey %s has no alias", privKeyName)
-		if alias == "root" {
+		if alias == data.CanonicalRootRole {
 			continue
 		}
 		relKeyPath := filepath.Join("tuf_keys", privKeyName+".key")
@@ -329,7 +329,7 @@ func TestImportExportRootKey(t *testing.T) {
 	// Should be able to unlock the root key with the old password
 	key, alias, err := cs2.GetPrivateKey(rootKeyID)
 	assert.NoError(t, err, "could not unlock root key")
-	assert.Equal(t, "root", alias)
+	assert.Equal(t, data.CanonicalRootRole, alias)
 	assert.Equal(t, rootKeyID, key.ID())
 }
 
@@ -381,7 +381,7 @@ func TestImportExportRootKeyReencrypt(t *testing.T) {
 	// Should be able to unlock the root key with the new password
 	key, alias, err := cs2.GetPrivateKey(rootKeyID)
 	assert.NoError(t, err, "could not unlock root key")
-	assert.Equal(t, "root", alias)
+	assert.Equal(t, data.CanonicalRootRole, alias)
 	assert.Equal(t, rootKeyID, key.ID())
 }
 
@@ -419,7 +419,10 @@ func TestImportExportNonRootKey(t *testing.T) {
 	keyReader, err := os.Open(tempKeyFilePath)
 	assert.NoError(t, err, "could not open key file")
 
-	err = cs2.ImportRoleKey(keyReader, data.CanonicalTargetsRole, oldPassphraseRetriever)
+	pemBytes, err := ioutil.ReadAll(keyReader)
+	assert.NoError(t, err, "could not read key file")
+
+	err = cs2.ImportRoleKey(pemBytes, data.CanonicalTargetsRole, oldPassphraseRetriever)
 	assert.NoError(t, err)
 	keyReader.Close()
 
@@ -433,7 +436,7 @@ func TestImportExportNonRootKey(t *testing.T) {
 	// Check that the key is the same
 	key, alias, err := cs2.GetPrivateKey(targetsKeyID)
 	assert.NoError(t, err, "could not unlock targets key")
-	assert.Equal(t, "targets", alias)
+	assert.Equal(t, data.CanonicalTargetsRole, alias)
 	assert.Equal(t, targetsKeyID, key.ID())
 }
 
@@ -471,7 +474,10 @@ func TestImportExportNonRootKeyReencrypt(t *testing.T) {
 	keyReader, err := os.Open(tempKeyFilePath)
 	assert.NoError(t, err, "could not open key file")
 
-	err = cs2.ImportRoleKey(keyReader, "snapshot", newPassphraseRetriever)
+	pemBytes, err := ioutil.ReadAll(keyReader)
+	assert.NoError(t, err, "could not read key file")
+
+	err = cs2.ImportRoleKey(pemBytes, data.CanonicalSnapshotRole, newPassphraseRetriever)
 	assert.NoError(t, err)
 	keyReader.Close()
 
@@ -485,6 +491,6 @@ func TestImportExportNonRootKeyReencrypt(t *testing.T) {
 	// Should be able to unlock the root key with the new password
 	key, alias, err := cs2.GetPrivateKey(snapshotKeyID)
 	assert.NoError(t, err, "could not unlock snapshot key")
-	assert.Equal(t, "snapshot", alias)
+	assert.Equal(t, data.CanonicalSnapshotRole, alias)
 	assert.Equal(t, snapshotKeyID, key.ID())
 }