From 7108450a21578e435c6e0912e9d998bcf937dddb Mon Sep 17 00:00:00 2001
From: Ying Li <ying.li@docker.com>
Date: Mon, 9 Nov 2015 19:32:55 -0800
Subject: [PATCH 1/7] Add more unit tests for the YubiKeyStore.

Including how it interacts with the backup key store, and with more
assertions against a new YubiKeyStore so that we won't get false
positives or negatives from the cache.

Signed-off-by: Ying Li <ying.li@docker.com>
Signed-off-by: David Lawrence <david.lawrence@docker.com>

Signed-off-by: Ying Li <ying.li@docker.com> (github: endophage)
---
 trustmanager/yubikey/yubikeystore_test.go | 355 ++++++++++++++++++----
 1 file changed, 300 insertions(+), 55 deletions(-)

diff --git a/trustmanager/yubikey/yubikeystore_test.go b/trustmanager/yubikey/yubikeystore_test.go
index 9d7863d62b..a5c10a0375 100644
--- a/trustmanager/yubikey/yubikeystore_test.go
+++ b/trustmanager/yubikey/yubikeystore_test.go
@@ -3,6 +3,7 @@
 package yubikey
 
 import (
+	"bytes"
 	"crypto/rand"
 	"reflect"
 	"testing"
@@ -13,11 +14,11 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+var ret = passphrase.ConstantRetriever("passphrase")
+
+// create a new store for clearing out keys, because we don't want to pollute
+// any cache
 func clearAllKeys(t *testing.T) {
-	// TODO(cyli): this is creating a new yubikey store because for some reason,
-	// removing and then adding with the same YubiKeyStore causes
-	// non-deterministic failures at least on Mac OS
-	ret := passphrase.ConstantRetriever("passphrase")
 	store, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
 	assert.NoError(t, err)
 
@@ -53,7 +54,7 @@ func TestEnsurePrivateKeySizePadsLessThanRequiredSizeArrays(t *testing.T) {
 	assert.True(t, reflect.DeepEqual(expected, result))
 }
 
-func testAddKey(t *testing.T, store *YubiKeyStore) (data.PrivateKey, error) {
+func testAddKey(t *testing.T, store trustmanager.KeyStore) (data.PrivateKey, error) {
 	privKey, err := trustmanager.GenerateECDSAKey(rand.Reader)
 	assert.NoError(t, err)
 
@@ -61,68 +62,207 @@ func testAddKey(t *testing.T, store *YubiKeyStore) (data.PrivateKey, error) {
 	return privKey, err
 }
 
-func TestAddKeyToNextEmptyYubikeySlot(t *testing.T) {
-	if !YubikeyAccessible() {
-		t.Skip("Must have Yubikey access.")
-	}
-	clearAllKeys(t)
-
-	ret := passphrase.ConstantRetriever("passphrase")
-	store, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
-	assert.NoError(t, err)
-	SetYubikeyKeyMode(KeymodeNone)
-	defer func() {
-		SetYubikeyKeyMode(KeymodeTouch | KeymodePinOnce)
-	}()
-
+func addMaxKeys(t *testing.T, store trustmanager.KeyStore) []string {
 	keys := make([]string, 0, numSlots)
-
 	// create the maximum number of keys
 	for i := 0; i < numSlots; i++ {
 		privKey, err := testAddKey(t, store)
 		assert.NoError(t, err)
 		keys = append(keys, privKey.ID())
 	}
-
-	// create a new store, to make sure we're not just using the keys cache
-	store, err = NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
-	assert.NoError(t, err)
-	listedKeys := store.ListKeys()
-	assert.Len(t, listedKeys, numSlots)
-	for _, k := range keys {
-		r, ok := listedKeys[k]
-		assert.True(t, ok)
-		assert.Equal(t, data.CanonicalRootRole, r)
-	}
-
-	// add another key - should fail because there are no more slots
-	_, err = testAddKey(t, store)
-	assert.Error(t, err)
-
-	// delete one of the middle keys, and assert we can still create a new key
-	err = store.RemoveKey(keys[numSlots/2])
-	assert.NoError(t, err)
-
-	_, err = testAddKey(t, store)
-	assert.NoError(t, err)
+	return keys
 }
 
-// ImportKey imports a key as root without adding it to the backup store
-func TestImportKey(t *testing.T) {
+// We can add keys enough times to fill up all the slots in the Yubikey.
+// They are backed up, and we can then list them and get the keys.
+func TestYubiAddKeysAndRetrieve(t *testing.T) {
 	if !YubikeyAccessible() {
 		t.Skip("Must have Yubikey access.")
 	}
 	clearAllKeys(t)
 
-	ret := passphrase.ConstantRetriever("passphrase")
-	backup := trustmanager.NewKeyMemoryStore(ret)
-	store, err := NewYubiKeyStore(backup, ret)
-	assert.NoError(t, err)
 	SetYubikeyKeyMode(KeymodeNone)
 	defer func() {
 		SetYubikeyKeyMode(KeymodeTouch | KeymodePinOnce)
 	}()
 
+	// create 4 keys on the original store
+	backup := trustmanager.NewKeyMemoryStore(ret)
+	origStore, err := NewYubiKeyStore(backup, ret)
+	assert.NoError(t, err)
+	keys := addMaxKeys(t, origStore)
+
+	// create a new store, since we want to be sure the original store's cache
+	// is not masking any issues
+	cleanStore, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
+	assert.NoError(t, err)
+
+	// All 4 keys should be in the original store, in the clean store (which
+	// makes sure the keys are actually on the Yubikey and not on the original
+	// store's cache, and on the backup store)
+	for _, store := range []trustmanager.KeyStore{origStore, cleanStore, backup} {
+		listedKeys := store.ListKeys()
+		assert.Len(t, listedKeys, numSlots)
+		for _, k := range keys {
+			r, ok := listedKeys[k]
+			assert.True(t, ok)
+			assert.Equal(t, data.CanonicalRootRole, r)
+
+			_, _, err := store.GetKey(k)
+			assert.NoError(t, err)
+		}
+	}
+}
+
+// We can't add a key if there are no more slots
+func TestYubiAddKeyFailureIfNoMoreSlots(t *testing.T) {
+	if !YubikeyAccessible() {
+		t.Skip("Must have Yubikey access.")
+	}
+	clearAllKeys(t)
+
+	SetYubikeyKeyMode(KeymodeNone)
+	defer func() {
+		SetYubikeyKeyMode(KeymodeTouch | KeymodePinOnce)
+	}()
+
+	// create 4 keys on the original store
+	backup := trustmanager.NewKeyMemoryStore(ret)
+	origStore, err := NewYubiKeyStore(backup, ret)
+	assert.NoError(t, err)
+	addMaxKeys(t, origStore)
+
+	// add another key - should fail because there are no more slots
+	badKey, err := testAddKey(t, origStore)
+	assert.Error(t, err)
+
+	// create a new store, since we want to be sure the original store's cache
+	// is not masking any issues
+	cleanStore, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
+	assert.NoError(t, err)
+
+	// The key should not be in the original store, in the new clean store, or
+	// in teh backup store.
+	for _, store := range []trustmanager.KeyStore{origStore, cleanStore, backup} {
+		// the key that wasn't created should not appear in ListKeys or GetKey
+		_, _, err := store.GetKey(badKey.ID())
+		assert.Error(t, err)
+		for k := range store.ListKeys() {
+			assert.NotEqual(t, badKey, k)
+		}
+	}
+}
+
+// If some random key in the middle was removed, adding a key will work (keys
+// do not have to be deleted/added in order)
+func TestYubiAddKeyCanAddToMiddleSlot(t *testing.T) {
+	if !YubikeyAccessible() {
+		t.Skip("Must have Yubikey access.")
+	}
+	clearAllKeys(t)
+
+	SetYubikeyKeyMode(KeymodeNone)
+	defer func() {
+		SetYubikeyKeyMode(KeymodeTouch | KeymodePinOnce)
+	}()
+
+	// create 4 keys on the original store
+	backup := trustmanager.NewKeyMemoryStore(ret)
+	origStore, err := NewYubiKeyStore(backup, ret)
+	assert.NoError(t, err)
+	keys := addMaxKeys(t, origStore)
+
+	// delete one of the middle keys, and assert we can still create a new key
+	keyIDToDelete := keys[numSlots/2]
+	err = origStore.RemoveKey(keyIDToDelete)
+	assert.NoError(t, err)
+
+	newKey, err := testAddKey(t, origStore)
+	assert.NoError(t, err)
+
+	// create a new store, since we want to be sure the original store's cache
+	// is not masking any issues
+	cleanStore, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
+	assert.NoError(t, err)
+
+	// The new key should be in the original store, in the new clean store, and
+	// in the backup store.  The old key should not be in the original store,
+	// or the new clean store.
+	for _, store := range []trustmanager.KeyStore{origStore, cleanStore, backup} {
+		// new key should appear in all stores
+		gottenKey, _, err := store.GetKey(newKey.ID())
+		assert.NoError(t, err)
+		assert.Equal(t, gottenKey.ID(), newKey.ID())
+
+		listedKeys := store.ListKeys()
+		_, ok := listedKeys[newKey.ID()]
+		assert.True(t, ok)
+
+		// old key should not be in the non-backup stores
+		if store != backup {
+			_, _, err := store.GetKey(keyIDToDelete)
+			assert.Error(t, err)
+			_, ok = listedKeys[keyIDToDelete]
+			assert.False(t, ok)
+		}
+	}
+}
+
+// RemoveKey removes a key from the yubikey, but not from the backup store.
+func TestYubiRemoveKey(t *testing.T) {
+	if !YubikeyAccessible() {
+		t.Skip("Must have Yubikey access.")
+	}
+	clearAllKeys(t)
+
+	SetYubikeyKeyMode(KeymodeNone)
+	defer func() {
+		SetYubikeyKeyMode(KeymodeTouch | KeymodePinOnce)
+	}()
+
+	backup := trustmanager.NewKeyMemoryStore(ret)
+	origStore, err := NewYubiKeyStore(backup, ret)
+	assert.NoError(t, err)
+
+	key, err := testAddKey(t, origStore)
+	assert.NoError(t, err)
+	err = origStore.RemoveKey(key.ID())
+	assert.NoError(t, err)
+
+	// key remains in the backup store
+	backupKey, role, err := backup.GetKey(key.ID())
+	assert.NoError(t, err)
+	assert.Equal(t, data.CanonicalRootRole, role)
+	assert.Equal(t, key.ID(), backupKey.ID())
+
+	// create a new store, since we want to be sure the original store's cache
+	// is not masking any issues
+	cleanStore, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
+	assert.NoError(t, err)
+
+	// key is not in either the original store or the clean store
+	for _, store := range []*YubiKeyStore{origStore, cleanStore} {
+		_, _, err := store.GetKey(key.ID())
+		assert.Error(t, err)
+	}
+}
+
+// ImportKey imports a key as root without adding it to the backup store
+func TestYubiImportKey(t *testing.T) {
+	if !YubikeyAccessible() {
+		t.Skip("Must have Yubikey access.")
+	}
+	clearAllKeys(t)
+
+	SetYubikeyKeyMode(KeymodeNone)
+	defer func() {
+		SetYubikeyKeyMode(KeymodeTouch | KeymodePinOnce)
+	}()
+
+	backup := trustmanager.NewKeyMemoryStore(ret)
+	origStore, err := NewYubiKeyStore(backup, ret)
+	assert.NoError(t, err)
+
 	// generate key and import it
 	privKey, err := trustmanager.GenerateECDSAKey(rand.Reader)
 	assert.NoError(t, err)
@@ -130,18 +270,123 @@ func TestImportKey(t *testing.T) {
 	pemBytes, err := trustmanager.EncryptPrivateKey(privKey, "passphrase")
 	assert.NoError(t, err)
 
-	err = store.ImportKey(pemBytes, "root")
+	err = origStore.ImportKey(pemBytes, "root")
 	assert.NoError(t, err)
 
 	// key is not in backup store
 	_, _, err = backup.GetKey(privKey.ID())
 	assert.Error(t, err)
 
-	// ensure key is in Yubikey -  create a new store, to make sure we're not
-	// just using the keys cache
-	store, err = NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
-	gottenKey, role, err := store.GetKey(privKey.ID())
+	// create a new store, since we want to be sure the original store's cache
+	// is not masking any issues
+	cleanStore, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
+	assert.NoError(t, err)
+	for _, store := range []*YubiKeyStore{origStore, cleanStore} {
+		gottenKey, role, err := store.GetKey(privKey.ID())
+		assert.NoError(t, err)
+		assert.Equal(t, data.CanonicalRootRole, role)
+		assert.Equal(t, privKey.Public(), gottenKey.Public())
+	}
+}
+
+// One cannot export from hardware - it will not export from the backup
+func TestYubiExportKeyFails(t *testing.T) {
+	if !YubikeyAccessible() {
+		t.Skip("Must have Yubikey access.")
+	}
+	clearAllKeys(t)
+
+	SetYubikeyKeyMode(KeymodeNone)
+	defer func() {
+		SetYubikeyKeyMode(KeymodeTouch | KeymodePinOnce)
+	}()
+
+	store, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
+	assert.NoError(t, err)
+
+	key, err := testAddKey(t, store)
+	assert.NoError(t, err)
+
+	_, err = store.ExportKey(key.ID())
+	assert.Error(t, err)
+}
+
+// If there are keys in the backup store but no keys in the Yubikey,
+// listing and getting cannot access the keys in the backup store
+func TestYubiListAndGetKeysIgnoresBackup(t *testing.T) {
+	if !YubikeyAccessible() {
+		t.Skip("Must have Yubikey access.")
+	}
+	clearAllKeys(t)
+
+	SetYubikeyKeyMode(KeymodeNone)
+	defer func() {
+		SetYubikeyKeyMode(KeymodeTouch | KeymodePinOnce)
+	}()
+
+	backup := trustmanager.NewKeyMemoryStore(ret)
+	key, err := testAddKey(t, backup)
+	assert.NoError(t, err)
+
+	store, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
+	assert.Len(t, store.ListKeys(), 0)
+	_, _, err = store.GetKey(key.ID())
+	assert.Error(t, err)
+}
+
+// Get a YubiPrivateKey.  Check that it has the right algorithm, etc, and
+// specifically that you cannot get the private bytes out.
+func TestYubiKey(t *testing.T) {
+	if !YubikeyAccessible() {
+		t.Skip("Must have Yubikey access.")
+	}
+	clearAllKeys(t)
+
+	SetYubikeyKeyMode(KeymodeNone)
+	defer func() {
+		SetYubikeyKeyMode(KeymodeTouch | KeymodePinOnce)
+	}()
+
+	store, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
+	assert.NoError(t, err)
+
+	ecdsaPrivateKey, err := testAddKey(t, store)
+	assert.NoError(t, err)
+
+	yubiPrivateKey, _, err := store.GetKey(ecdsaPrivateKey.ID())
+	assert.NoError(t, err)
+
+	assert.Equal(t, data.ECDSAKey, yubiPrivateKey.Algorithm())
+	assert.Equal(t, data.ECDSASignature, yubiPrivateKey.SignatureAlgorithm())
+	assert.Equal(t, ecdsaPrivateKey.Public(), yubiPrivateKey.Public())
+	assert.Nil(t, yubiPrivateKey.Private())
+}
+
+// Get a YubiPrivateKey.  Sign something with it.
+func TestYubiSigning(t *testing.T) {
+	// TODO(cyli): the signature should be verified, but the importing the
+	// verifiers causes an import cycle.  A bigger refactor needs to be done
+	// to fix it.
+	if !YubikeyAccessible() {
+		t.Skip("Must have Yubikey access.")
+	}
+	clearAllKeys(t)
+
+	SetYubikeyKeyMode(KeymodeNone)
+	defer func() {
+		SetYubikeyKeyMode(KeymodeTouch | KeymodePinOnce)
+	}()
+
+	store, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
+	assert.NoError(t, err)
+
+	ecdsaPrivateKey, err := testAddKey(t, store)
+	assert.NoError(t, err)
+
+	yubiPrivateKey, _, err := store.GetKey(ecdsaPrivateKey.ID())
+	assert.NoError(t, err)
+
+	msg := []byte("Hello there")
+	_, err = yubiPrivateKey.Sign(bytes.NewBuffer(msg), msg, nil)
 	assert.NoError(t, err)
-	assert.Equal(t, data.CanonicalRootRole, role)
-	assert.Equal(t, privKey.Public(), gottenKey.Public())
 }

From 09c0f9d05bbf3752f6515036309450ca5705e074 Mon Sep 17 00:00:00 2001
From: Ying Li <ying.li@docker.com>
Date: Tue, 10 Nov 2015 00:53:11 -0800
Subject: [PATCH 2/7] Replace the pkcs11 library with interfaces for easier
 testing.

Signed-off-by: Ying Li <ying.li@docker.com>
Signed-off-by: David Lawrence <david.lawrence@docker.com>

Signed-off-by: Ying Li <ying.li@docker.com> (github: endophage)
---
 trustmanager/yubikey/pkcs11_interface.go | 40 +++++++++++++++++++
 trustmanager/yubikey/yubikeystore.go     | 50 ++++++++++++++++--------
 2 files changed, 73 insertions(+), 17 deletions(-)
 create mode 100644 trustmanager/yubikey/pkcs11_interface.go

diff --git a/trustmanager/yubikey/pkcs11_interface.go b/trustmanager/yubikey/pkcs11_interface.go
new file mode 100644
index 0000000000..4a46e05f0c
--- /dev/null
+++ b/trustmanager/yubikey/pkcs11_interface.go
@@ -0,0 +1,40 @@
+// +build pkcs11
+
+// an interface around the pkcs11 library, so that things can be mocked out
+// for testing
+
+package yubikey
+
+import "github.com/miekg/pkcs11"
+
+// IPKCS11 is an interface for wrapping github.com/miekg/pkcs11
+type pkcs11LibLoader func(module string) IPKCS11Ctx
+
+func defaultLoader(module string) IPKCS11Ctx {
+	return pkcs11.New(module)
+}
+
+// IPKCS11Ctx is an interface for wrapping the parts of
+// github.com/miekg/pkcs11.Ctx that yubikeystore requires
+type IPKCS11Ctx interface {
+	Destroy()
+	Initialize() error
+	Finalize() error
+	GetSlotList(tokenPresent bool) ([]uint, error)
+	OpenSession(slotID uint, flags uint) (pkcs11.SessionHandle, error)
+	CloseSession(sh pkcs11.SessionHandle) error
+	Login(sh pkcs11.SessionHandle, userType uint, pin string) error
+	Logout(sh pkcs11.SessionHandle) error
+	CreateObject(sh pkcs11.SessionHandle, temp []*pkcs11.Attribute) (
+		pkcs11.ObjectHandle, error)
+	DestroyObject(sh pkcs11.SessionHandle, oh pkcs11.ObjectHandle) error
+	GetAttributeValue(sh pkcs11.SessionHandle, o pkcs11.ObjectHandle,
+		a []*pkcs11.Attribute) ([]*pkcs11.Attribute, error)
+	FindObjectsInit(sh pkcs11.SessionHandle, temp []*pkcs11.Attribute) error
+	FindObjects(sh pkcs11.SessionHandle, max int) (
+		[]pkcs11.ObjectHandle, bool, error)
+	FindObjectsFinal(sh pkcs11.SessionHandle) error
+	SignInit(sh pkcs11.SessionHandle, m []*pkcs11.Mechanism,
+		o pkcs11.ObjectHandle) error
+	Sign(sh pkcs11.SessionHandle, message []byte) ([]byte, error)
+}
diff --git a/trustmanager/yubikey/yubikeystore.go b/trustmanager/yubikey/yubikeystore.go
index c3e55da0b7..fae3be6659 100644
--- a/trustmanager/yubikey/yubikeystore.go
+++ b/trustmanager/yubikey/yubikeystore.go
@@ -112,17 +112,21 @@ type YubiPrivateKey struct {
 	data.ECDSAPublicKey
 	passRetriever passphrase.Retriever
 	slot          []byte
+	libLoader     pkcs11LibLoader
 }
 
 type YubikeySigner struct {
 	YubiPrivateKey
 }
 
-func NewYubiPrivateKey(slot []byte, pubKey data.ECDSAPublicKey, passRetriever passphrase.Retriever) *YubiPrivateKey {
+func NewYubiPrivateKey(slot []byte, pubKey data.ECDSAPublicKey,
+	passRetriever passphrase.Retriever) *YubiPrivateKey {
+
 	return &YubiPrivateKey{
 		ECDSAPublicKey: pubKey,
 		passRetriever:  passRetriever,
 		slot:           slot,
+		libLoader:      defaultLoader,
 	}
 }
 
@@ -135,6 +139,10 @@ func (ys *YubikeySigner) Public() crypto.PublicKey {
 	return publicKey
 }
 
+func (y *YubiPrivateKey) setLibLoader(loader pkcs11LibLoader) {
+	y.libLoader = loader
+}
+
 // CryptoSigner returns a crypto.Signer tha wraps the YubiPrivateKey. Needed for
 // Certificate generation only
 func (y *YubiPrivateKey) CryptoSigner() crypto.Signer {
@@ -153,7 +161,7 @@ func (y YubiPrivateKey) SignatureAlgorithm() data.SigAlgorithm {
 }
 
 func (y *YubiPrivateKey) Sign(rand io.Reader, msg []byte, opts crypto.SignerOpts) ([]byte, error) {
-	ctx, session, err := SetupHSMEnv(pkcs11Lib)
+	ctx, session, err := SetupHSMEnv(pkcs11Lib, y.libLoader)
 	if err != nil {
 		return nil, err
 	}
@@ -186,7 +194,7 @@ func ensurePrivateKeySize(payload []byte) []byte {
 
 // addECDSAKey adds a key to the yubikey
 func addECDSAKey(
-	ctx *pkcs11.Ctx,
+	ctx IPKCS11Ctx,
 	session pkcs11.SessionHandle,
 	privKey data.PrivateKey,
 	pkcs11KeyID []byte,
@@ -255,7 +263,7 @@ func addECDSAKey(
 	return nil
 }
 
-func getECDSAKey(ctx *pkcs11.Ctx, session pkcs11.SessionHandle, pkcs11KeyID []byte) (*data.ECDSAPublicKey, string, error) {
+func getECDSAKey(ctx IPKCS11Ctx, session pkcs11.SessionHandle, pkcs11KeyID []byte) (*data.ECDSAPublicKey, string, error) {
 	findTemplate := []*pkcs11.Attribute{
 		pkcs11.NewAttribute(pkcs11.CKA_TOKEN, true),
 		pkcs11.NewAttribute(pkcs11.CKA_ID, pkcs11KeyID),
@@ -313,7 +321,7 @@ func getECDSAKey(ctx *pkcs11.Ctx, session pkcs11.SessionHandle, pkcs11KeyID []by
 }
 
 // Sign returns a signature for a given signature request
-func sign(ctx *pkcs11.Ctx, session pkcs11.SessionHandle, pkcs11KeyID []byte, passRetriever passphrase.Retriever, payload []byte) ([]byte, error) {
+func sign(ctx IPKCS11Ctx, session pkcs11.SessionHandle, pkcs11KeyID []byte, passRetriever passphrase.Retriever, payload []byte) ([]byte, error) {
 	err := login(ctx, session, passRetriever, pkcs11.CKU_USER, USER_PIN)
 	if err != nil {
 		return nil, fmt.Errorf("error logging in: %v", err)
@@ -367,7 +375,7 @@ func sign(ctx *pkcs11.Ctx, session pkcs11.SessionHandle, pkcs11KeyID []byte, pas
 	return sig[:], nil
 }
 
-func yubiRemoveKey(ctx *pkcs11.Ctx, session pkcs11.SessionHandle, pkcs11KeyID []byte, passRetriever passphrase.Retriever, keyID string) error {
+func yubiRemoveKey(ctx IPKCS11Ctx, session pkcs11.SessionHandle, pkcs11KeyID []byte, passRetriever passphrase.Retriever, keyID string) error {
 	err := login(ctx, session, passRetriever, pkcs11.CKU_SO, SO_USER_PIN)
 	if err != nil {
 		return err
@@ -408,7 +416,7 @@ func yubiRemoveKey(ctx *pkcs11.Ctx, session pkcs11.SessionHandle, pkcs11KeyID []
 	return nil
 }
 
-func yubiListKeys(ctx *pkcs11.Ctx, session pkcs11.SessionHandle) (keys map[string]yubiSlot, err error) {
+func yubiListKeys(ctx IPKCS11Ctx, session pkcs11.SessionHandle) (keys map[string]yubiSlot, err error) {
 	keys = make(map[string]yubiSlot)
 	findTemplate := []*pkcs11.Attribute{
 		pkcs11.NewAttribute(pkcs11.CKA_TOKEN, true),
@@ -501,7 +509,7 @@ func yubiListKeys(ctx *pkcs11.Ctx, session pkcs11.SessionHandle) (keys map[strin
 	return
 }
 
-func getNextEmptySlot(ctx *pkcs11.Ctx, session pkcs11.SessionHandle) ([]byte, error) {
+func getNextEmptySlot(ctx IPKCS11Ctx, session pkcs11.SessionHandle) ([]byte, error) {
 	findTemplate := []*pkcs11.Attribute{
 		pkcs11.NewAttribute(pkcs11.CKA_TOKEN, true),
 	}
@@ -571,6 +579,7 @@ type YubiKeyStore struct {
 	passRetriever passphrase.Retriever
 	keys          map[string]yubiSlot
 	backupStore   trustmanager.KeyStore
+	libLoader     pkcs11LibLoader
 }
 
 // NewYubiKeyStore returns a YubiKeyStore, given a backup key store to write any
@@ -582,6 +591,7 @@ func NewYubiKeyStore(backupStore trustmanager.KeyStore, passphraseRetriever pass
 		passRetriever: passphraseRetriever,
 		keys:          make(map[string]yubiSlot),
 		backupStore:   backupStore,
+		libLoader:     defaultLoader,
 	}
 	s.ListKeys() // populate keys field
 	return s, nil
@@ -593,11 +603,15 @@ func (s YubiKeyStore) Name() string {
 	return "yubikey"
 }
 
+func (s *YubiKeyStore) setLibLoader(loader pkcs11LibLoader) {
+	s.libLoader = loader
+}
+
 func (s *YubiKeyStore) ListKeys() map[string]string {
 	if len(s.keys) > 0 {
 		return buildKeyMap(s.keys)
 	}
-	ctx, session, err := SetupHSMEnv(pkcs11Lib)
+	ctx, session, err := SetupHSMEnv(pkcs11Lib, s.libLoader)
 	if err != nil {
 		logrus.Debugf("Failed to initialize PKCS11 environment: %s", err.Error())
 		return nil
@@ -625,7 +639,7 @@ func (s *YubiKeyStore) addKey(
 		return fmt.Errorf("yubikey only supports storing root keys, got %s for key: %s", role, keyID)
 	}
 
-	ctx, session, err := SetupHSMEnv(pkcs11Lib)
+	ctx, session, err := SetupHSMEnv(pkcs11Lib, s.libLoader)
 	if err != nil {
 		logrus.Debugf("Failed to initialize PKCS11 environment: %s", err.Error())
 		return err
@@ -668,7 +682,7 @@ func (s *YubiKeyStore) addKey(
 // GetKey retrieves a key from the Yubikey only (it does not look inside the
 // backup store)
 func (s *YubiKeyStore) GetKey(keyID string) (data.PrivateKey, string, error) {
-	ctx, session, err := SetupHSMEnv(pkcs11Lib)
+	ctx, session, err := SetupHSMEnv(pkcs11Lib, s.libLoader)
 	if err != nil {
 		logrus.Debugf("Failed to initialize PKCS11 environment: %s", err.Error())
 		return nil, "", err
@@ -700,7 +714,7 @@ func (s *YubiKeyStore) GetKey(keyID string) (data.PrivateKey, string, error) {
 // RemoveKey deletes a key from the Yubikey only (it does not remove it from the
 // backup store)
 func (s *YubiKeyStore) RemoveKey(keyID string) error {
-	ctx, session, err := SetupHSMEnv(pkcs11Lib)
+	ctx, session, err := SetupHSMEnv(pkcs11Lib, s.libLoader)
 	if err != nil {
 		logrus.Debugf("Failed to initialize PKCS11 environment: %s", err.Error())
 		return nil
@@ -741,18 +755,20 @@ func (s *YubiKeyStore) ImportKey(pemBytes []byte, keyPath string) error {
 	return s.addKey(privKey.ID(), "root", privKey, false)
 }
 
-func cleanup(ctx *pkcs11.Ctx, session pkcs11.SessionHandle) {
+func cleanup(ctx IPKCS11Ctx, session pkcs11.SessionHandle) {
 	ctx.CloseSession(session)
 	ctx.Finalize()
 	ctx.Destroy()
 }
 
 // SetupHSMEnv is a method that depends on the existences
-func SetupHSMEnv(libraryPath string) (*pkcs11.Ctx, pkcs11.SessionHandle, error) {
+func SetupHSMEnv(libraryPath string, libLoader pkcs11LibLoader) (
+	IPKCS11Ctx, pkcs11.SessionHandle, error) {
+
 	if libraryPath == "" {
 		return nil, 0, errors.New("No library found.")
 	}
-	p := pkcs11.New(libraryPath)
+	p := libLoader(libraryPath)
 
 	if p == nil {
 		return nil, 0, errors.New("Failed to init library")
@@ -786,7 +802,7 @@ func YubikeyAccessible() bool {
 	if pkcs11Lib == "" {
 		return false
 	}
-	ctx, session, err := SetupHSMEnv(pkcs11Lib)
+	ctx, session, err := SetupHSMEnv(pkcs11Lib, defaultLoader)
 	if err != nil {
 		return false
 	}
@@ -794,7 +810,7 @@ func YubikeyAccessible() bool {
 	return true
 }
 
-func login(ctx *pkcs11.Ctx, session pkcs11.SessionHandle, passRetriever passphrase.Retriever, userFlag uint, defaultPassw string) error {
+func login(ctx IPKCS11Ctx, session pkcs11.SessionHandle, passRetriever passphrase.Retriever, userFlag uint, defaultPassw string) error {
 	// try default password
 	err := ctx.Login(session, userFlag, defaultPassw)
 	if err == nil {

From 73a26d59ac6d5e59ac9a4204de0b2b01f8941d55 Mon Sep 17 00:00:00 2001
From: Ying Li <ying.li@docker.com>
Date: Wed, 11 Nov 2015 11:31:35 -0800
Subject: [PATCH 3/7] Inject errors into pkcs11 in order to test that the
 yubikey code cleans up.

Signed-off-by: Ying Li <ying.li@docker.com>
Signed-off-by: David Lawrence <david.lawrence@docker.com>

Signed-off-by: Ying Li <ying.li@docker.com> (github: endophage)
---
 trustmanager/yubikey/yubikeystore.go      |  34 +-
 trustmanager/yubikey/yubikeystore_test.go | 594 +++++++++++++++++++++-
 2 files changed, 603 insertions(+), 25 deletions(-)

diff --git a/trustmanager/yubikey/yubikeystore.go b/trustmanager/yubikey/yubikeystore.go
index fae3be6659..f1528292f3 100644
--- a/trustmanager/yubikey/yubikeystore.go
+++ b/trustmanager/yubikey/yubikeystore.go
@@ -162,10 +162,10 @@ func (y YubiPrivateKey) SignatureAlgorithm() data.SigAlgorithm {
 
 func (y *YubiPrivateKey) Sign(rand io.Reader, msg []byte, opts crypto.SignerOpts) ([]byte, error) {
 	ctx, session, err := SetupHSMEnv(pkcs11Lib, y.libLoader)
+	defer cleanup(ctx, session)
 	if err != nil {
 		return nil, err
 	}
-	defer cleanup(ctx, session)
 
 	v := signed.Verifiers[data.ECDSASignature]
 	for i := 0; i < sigAttempts; i++ {
@@ -294,7 +294,7 @@ func getECDSAKey(ctx IPKCS11Ctx, session pkcs11.SessionHandle, pkcs11KeyID []byt
 		return nil, "", errors.New("no matching keys found inside of yubikey")
 	}
 
-	// Retrieve the public-key material to be able to create a new HSMRSAKey
+	// Retrieve the public-key material to be able to create a new ECSAKey
 	attr, err := ctx.GetAttributeValue(session, obj[0], attrTemplate)
 	if err != nil {
 		logrus.Debugf("Failed to get Attribute for: %v", obj[0])
@@ -464,7 +464,7 @@ func yubiListKeys(ctx IPKCS11Ctx, session pkcs11.SessionHandle) (keys map[string
 			cert *x509.Certificate
 			slot []byte
 		)
-		// Retrieve the public-key material to be able to create a new HSMRSAKey
+		// Retrieve the public-key material to be able to create a new ECDSA
 		attr, err := ctx.GetAttributeValue(session, obj, attrTemplate)
 		if err != nil {
 			logrus.Debugf("Failed to get Attribute for: %v", obj)
@@ -612,11 +612,11 @@ func (s *YubiKeyStore) ListKeys() map[string]string {
 		return buildKeyMap(s.keys)
 	}
 	ctx, session, err := SetupHSMEnv(pkcs11Lib, s.libLoader)
+	defer cleanup(ctx, session)
 	if err != nil {
 		logrus.Debugf("Failed to initialize PKCS11 environment: %s", err.Error())
 		return nil
 	}
-	defer cleanup(ctx, session)
 	keys, err := yubiListKeys(ctx, session)
 	if err != nil {
 		logrus.Debugf("Failed to list key from the yubikey: %s", err.Error())
@@ -640,11 +640,11 @@ func (s *YubiKeyStore) addKey(
 	}
 
 	ctx, session, err := SetupHSMEnv(pkcs11Lib, s.libLoader)
+	defer cleanup(ctx, session)
 	if err != nil {
 		logrus.Debugf("Failed to initialize PKCS11 environment: %s", err.Error())
 		return err
 	}
-	defer cleanup(ctx, session)
 
 	if k, ok := s.keys[keyID]; ok {
 		if k.role == role {
@@ -683,11 +683,11 @@ func (s *YubiKeyStore) addKey(
 // backup store)
 func (s *YubiKeyStore) GetKey(keyID string) (data.PrivateKey, string, error) {
 	ctx, session, err := SetupHSMEnv(pkcs11Lib, s.libLoader)
+	defer cleanup(ctx, session)
 	if err != nil {
 		logrus.Debugf("Failed to initialize PKCS11 environment: %s", err.Error())
 		return nil, "", err
 	}
-	defer cleanup(ctx, session)
 
 	key, ok := s.keys[keyID]
 	if !ok {
@@ -715,11 +715,11 @@ func (s *YubiKeyStore) GetKey(keyID string) (data.PrivateKey, string, error) {
 // backup store)
 func (s *YubiKeyStore) RemoveKey(keyID string) error {
 	ctx, session, err := SetupHSMEnv(pkcs11Lib, s.libLoader)
+	defer cleanup(ctx, session)
 	if err != nil {
 		logrus.Debugf("Failed to initialize PKCS11 environment: %s", err.Error())
 		return nil
 	}
-	defer cleanup(ctx, session)
 	key, ok := s.keys[keyID]
 	if !ok {
 		return errors.New("Key not present in yubikey")
@@ -756,9 +756,17 @@ func (s *YubiKeyStore) ImportKey(pemBytes []byte, keyPath string) error {
 }
 
 func cleanup(ctx IPKCS11Ctx, session pkcs11.SessionHandle) {
-	ctx.CloseSession(session)
-	ctx.Finalize()
-	ctx.Destroy()
+	if ctx != nil {
+		err := ctx.CloseSession(session)
+		if err != nil {
+			logrus.Debugf("Error closing session: %s", err.Error())
+		}
+		err = ctx.Finalize()
+		if err != nil {
+			logrus.Debugf("Error finalizing: %s", err.Error())
+		}
+		ctx.Destroy()
+	}
 }
 
 // SetupHSMEnv is a method that depends on the existences
@@ -775,15 +783,18 @@ func SetupHSMEnv(libraryPath string, libLoader pkcs11LibLoader) (
 	}
 
 	if err := p.Initialize(); err != nil {
+		defer cleanup(p, 0)
 		return nil, 0, fmt.Errorf("Initialize error %s", err.Error())
 	}
 
 	slots, err := p.GetSlotList(true)
 	if err != nil {
+		defer cleanup(p, 0)
 		return nil, 0, fmt.Errorf("Failed to list HSM slots %s", err)
 	}
 	// Check to see if we got any slots from the HSM.
 	if len(slots) < 1 {
+		defer cleanup(p, 0)
 		return nil, 0, fmt.Errorf("No HSM Slots found")
 	}
 
@@ -791,6 +802,7 @@ func SetupHSMEnv(libraryPath string, libLoader pkcs11LibLoader) (
 	// CKF_RW_SESSION: TRUE if the session is read/write; FALSE if the session is read-only
 	session, err := p.OpenSession(slots[0], pkcs11.CKF_SERIAL_SESSION|pkcs11.CKF_RW_SESSION)
 	if err != nil {
+		defer cleanup(p, 0)
 		return nil, 0, fmt.Errorf("Failed to Start Session with HSM %s", err)
 	}
 
@@ -803,10 +815,10 @@ func YubikeyAccessible() bool {
 		return false
 	}
 	ctx, session, err := SetupHSMEnv(pkcs11Lib, defaultLoader)
+	defer cleanup(ctx, session)
 	if err != nil {
 		return false
 	}
-	defer cleanup(ctx, session)
 	return true
 }
 
diff --git a/trustmanager/yubikey/yubikeystore_test.go b/trustmanager/yubikey/yubikeystore_test.go
index a5c10a0375..2451dc3d08 100644
--- a/trustmanager/yubikey/yubikeystore_test.go
+++ b/trustmanager/yubikey/yubikeystore_test.go
@@ -3,14 +3,15 @@
 package yubikey
 
 import (
-	"bytes"
 	"crypto/rand"
+	"fmt"
 	"reflect"
 	"testing"
 
 	"github.com/docker/notary/passphrase"
 	"github.com/docker/notary/trustmanager"
 	"github.com/docker/notary/tuf/data"
+	"github.com/miekg/pkcs11"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -248,7 +249,7 @@ func TestYubiRemoveKey(t *testing.T) {
 }
 
 // ImportKey imports a key as root without adding it to the backup store
-func TestYubiImportKey(t *testing.T) {
+func TestYubiImportNewKey(t *testing.T) {
 	if !YubikeyAccessible() {
 		t.Skip("Must have Yubikey access.")
 	}
@@ -289,6 +290,74 @@ func TestYubiImportKey(t *testing.T) {
 	}
 }
 
+// Importing an existing key succeeds, but doesn't actually add the key, nor
+// does it write it to backup.
+func TestYubiImportExistingKey(t *testing.T) {
+	if !YubikeyAccessible() {
+		t.Skip("Must have Yubikey access.")
+	}
+	clearAllKeys(t)
+
+	SetYubikeyKeyMode(KeymodeNone)
+	defer func() {
+		SetYubikeyKeyMode(KeymodeTouch | KeymodePinOnce)
+	}()
+
+	origStore, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
+	assert.NoError(t, err)
+	key, err := testAddKey(t, origStore)
+
+	backup := trustmanager.NewKeyMemoryStore(ret)
+	newStore, err := NewYubiKeyStore(backup, ret)
+	assert.NoError(t, err)
+
+	// for sanity, ensure that the key is already in the Yubikey
+	k, _, err := newStore.GetKey(key.ID())
+	assert.NoError(t, err)
+	assert.NotNil(t, k)
+
+	// import the key, which should have already been added to the yubikey
+	pemBytes, err := trustmanager.EncryptPrivateKey(key, "passphrase")
+	assert.NoError(t, err)
+	err = newStore.ImportKey(pemBytes, "root")
+	assert.NoError(t, err)
+
+	// key is not in backup store
+	_, _, err = backup.GetKey(key.ID())
+	assert.Error(t, err)
+}
+
+// Importing a key not as root fails, and it is not added to the backup store
+func TestYubiImportNonRootKey(t *testing.T) {
+	if !YubikeyAccessible() {
+		t.Skip("Must have Yubikey access.")
+	}
+	clearAllKeys(t)
+
+	SetYubikeyKeyMode(KeymodeNone)
+	defer func() {
+		SetYubikeyKeyMode(KeymodeTouch | KeymodePinOnce)
+	}()
+
+	backup := trustmanager.NewKeyMemoryStore(ret)
+	origStore, err := NewYubiKeyStore(backup, ret)
+	assert.NoError(t, err)
+
+	// generate key and import it
+	privKey, err := trustmanager.GenerateECDSAKey(rand.Reader)
+	assert.NoError(t, err)
+
+	pemBytes, err := trustmanager.EncryptPrivateKey(privKey, "passphrase")
+	assert.NoError(t, err)
+
+	err = origStore.ImportKey(pemBytes, privKey.ID())
+	assert.Error(t, err)
+
+	// key is not in backup store
+	_, _, err = backup.GetKey(privKey.ID())
+	assert.Error(t, err)
+}
+
 // One cannot export from hardware - it will not export from the backup
 func TestYubiExportKeyFails(t *testing.T) {
 	if !YubikeyAccessible() {
@@ -335,8 +404,9 @@ func TestYubiListAndGetKeysIgnoresBackup(t *testing.T) {
 }
 
 // Get a YubiPrivateKey.  Check that it has the right algorithm, etc, and
-// specifically that you cannot get the private bytes out.
-func TestYubiKey(t *testing.T) {
+// specifically that you cannot get the private bytes out.  Assume we can
+// sign something.
+func TestYubiKeyAndSign(t *testing.T) {
 	if !YubikeyAccessible() {
 		t.Skip("Must have Yubikey access.")
 	}
@@ -360,13 +430,319 @@ func TestYubiKey(t *testing.T) {
 	assert.Equal(t, data.ECDSASignature, yubiPrivateKey.SignatureAlgorithm())
 	assert.Equal(t, ecdsaPrivateKey.Public(), yubiPrivateKey.Public())
 	assert.Nil(t, yubiPrivateKey.Private())
+
+	// The signature should be verified, but the importing the verifiers causes
+	// an import cycle.  A bigger refactor needs to be done to fix it.
+	msg := []byte("Hello there")
+	_, err = yubiPrivateKey.Sign(rand.Reader, msg, nil)
+	assert.NoError(t, err)
 }
 
-// Get a YubiPrivateKey.  Sign something with it.
-func TestYubiSigning(t *testing.T) {
-	// TODO(cyli): the signature should be verified, but the importing the
-	// verifiers causes an import cycle.  A bigger refactor needs to be done
-	// to fix it.
+// ----- Negative tests that use stubbed pkcs11 for error injection -----
+
+type pkcs11Stubbable interface {
+	setLibLoader(pkcs11LibLoader)
+}
+
+var setupErrors = []string{"Initialize", "GetSlotList", "OpenSession"}
+
+// Create a new store, so that we avoid any cache issues, and list keys
+func cleanListKeys(t *testing.T) map[string]string {
+	cleanStore, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
+	assert.NoError(t, err)
+	return cleanStore.ListKeys()
+}
+
+// If an error occurs during login, which only some functions do, the function
+// under test will clean up after itself
+func testYubiFunctionCleansUpOnLoginError(t *testing.T, toStub pkcs11Stubbable,
+	functionUnderTest func() error) {
+
+	toStub.setLibLoader(func(string) IPKCS11Ctx {
+		return NewStubCtx(map[string]bool{"Login": true})
+	})
+
+	err := functionUnderTest()
+	assert.Error(t, err)
+	// a lot of these functions wrap other errors
+	assert.Contains(t, err.Error(), trustmanager.ErrAttemptsExceeded{}.Error())
+
+	// Set Up another time, to ensure we weren't left in a bad state
+	// by the previous runs
+	ctx, session, err := SetupHSMEnv(pkcs11Lib, defaultLoader)
+	assert.NoError(t, err)
+	cleanup(ctx, session)
+}
+
+// If one of the specified pkcs11 functions errors, the function under test
+// will clean up after itself
+func testYubiFunctionCleansUpOnSpecifiedErrors(t *testing.T,
+	toStub pkcs11Stubbable, functionUnderTest func() error,
+	dependentFunctions []string, functionShouldError bool) {
+
+	for _, methodName := range dependentFunctions {
+
+		toStub.setLibLoader(func(string) IPKCS11Ctx {
+			return NewStubCtx(
+				map[string]bool{methodName: true})
+		})
+
+		err := functionUnderTest()
+		if functionShouldError {
+			assert.Error(t, err,
+				fmt.Sprintf("Didn't error when %s errored.", methodName))
+			// a lot of these functions wrap other errors
+			assert.Contains(t, err.Error(), errInjected{methodName}.Error())
+		} else {
+			assert.NoError(t, err)
+		}
+	}
+
+	// Set Up another time, to ensure we weren't left in a bad state
+	// by the previous runs
+	ctx, session, err := SetupHSMEnv(pkcs11Lib, defaultLoader)
+	assert.NoError(t, err)
+	cleanup(ctx, session)
+}
+
+func TestYubiAddKeyCleansUpOnError(t *testing.T) {
+	if !YubikeyAccessible() {
+		t.Skip("Must have Yubikey access.")
+	}
+	clearAllKeys(t)
+
+	SetYubikeyKeyMode(KeymodeNone)
+	defer func() {
+		SetYubikeyKeyMode(KeymodeTouch | KeymodePinOnce)
+	}()
+
+	backup := trustmanager.NewKeyMemoryStore(ret)
+	origStore, err := NewYubiKeyStore(backup, ret)
+	assert.NoError(t, err)
+
+	var _addkey = func() error {
+		_, err := testAddKey(t, origStore)
+		return err
+	}
+
+	testYubiFunctionCleansUpOnLoginError(t, origStore, _addkey)
+	// all the PKCS11 functions AddKey depends on that aren't the login/logout
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, origStore, _addkey,
+		append(
+			setupErrors,
+			"FindObjectsInit",
+			"FindObjects",
+			"CreateObject",
+		), true)
+
+	// given that everything should have errored, there should be no keys on
+	// the yubikey and no keys in backup
+	assert.Len(t, backup.ListKeys(), 0)
+	assert.Len(t, cleanListKeys(t), 0)
+
+	// Logout should not cause a function failure - it s a cleanup failure,
+	// which shouldn't break anything, and it should clean up after itself.
+	// The key should be added to both stores
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, origStore, _addkey,
+		[]string{"Logout"}, false)
+
+	listedKeys := cleanListKeys(t)
+	assert.Len(t, backup.ListKeys(), 1)
+	assert.Len(t, listedKeys, 1)
+
+	// Currently, if GetAttributeValue fails, the function succeeds, because if
+	// we can't get the attribute value of an object, we don't know what slot
+	// it's in, we assume its occupied slot is free (hence this failure will
+	// cause the previous key to be overwritten).  This behavior may need to
+	// be revisited.
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, origStore, _addkey,
+		[]string{"GetAttributeValue"}, false)
+
+	newListedKeys := cleanListKeys(t)
+	// because the original key got overwritten
+	assert.Len(t, backup.ListKeys(), 2)
+	assert.Len(t, newListedKeys, 1)
+	for k := range newListedKeys {
+		_, ok := listedKeys[k]
+		assert.False(t, ok)
+	}
+}
+
+func TestYubiGetKeyCleansUpOnError(t *testing.T) {
+	if !YubikeyAccessible() {
+		t.Skip("Must have Yubikey access.")
+	}
+	clearAllKeys(t)
+
+	SetYubikeyKeyMode(KeymodeNone)
+	defer func() {
+		SetYubikeyKeyMode(KeymodeTouch | KeymodePinOnce)
+	}()
+
+	store, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
+	assert.NoError(t, err)
+	key, err := testAddKey(t, store)
+	assert.NoError(t, err)
+
+	var _getkey = func() error {
+		_, _, err := store.GetKey(key.ID())
+		return err
+	}
+
+	// all the PKCS11 functions GetKey depends on
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, store, _getkey,
+		append(
+			setupErrors,
+			"FindObjectsInit",
+			"FindObjects",
+			"FindObjectsFinal",
+			"GetAttributeValue",
+		), true)
+}
+
+func TestYubiImportKeyCleansUpOnError(t *testing.T) {
+	if !YubikeyAccessible() {
+		t.Skip("Must have Yubikey access.")
+	}
+	clearAllKeys(t)
+
+	SetYubikeyKeyMode(KeymodeNone)
+	defer func() {
+		SetYubikeyKeyMode(KeymodeTouch | KeymodePinOnce)
+	}()
+
+	origStore, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
+	assert.NoError(t, err)
+
+	privKey, err := trustmanager.GenerateECDSAKey(rand.Reader)
+	assert.NoError(t, err)
+
+	pemBytes, err := trustmanager.EncryptPrivateKey(privKey, "passphrase")
+	assert.NoError(t, err)
+
+	var _importkey = func() error { return origStore.ImportKey(pemBytes, "root") }
+
+	testYubiFunctionCleansUpOnLoginError(t, origStore, _importkey)
+	// all the PKCS11 functions ImportKey depends on that aren't the login/logout
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, origStore, _importkey,
+		append(
+			setupErrors,
+			"FindObjectsInit",
+			"FindObjects",
+			"CreateObject",
+		), true)
+
+	// given that everything should have errored, there should be no keys on
+	// the yubikey
+	assert.Len(t, cleanListKeys(t), 0)
+
+	// Logout should not cause a function failure - it s a cleanup failure,
+	// which shouldn't break anything, and it should clean up after itself.
+	// The key should be added to both stores
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, origStore, _importkey,
+		[]string{"Logout"}, false)
+
+	listedKeys := cleanListKeys(t)
+	assert.Len(t, listedKeys, 1)
+
+	// Currently, if GetAttributeValue fails, the function succeeds, because if
+	// we can't get the attribute value of an object, we don't know what slot
+	// it's in, we assume its occupied slot is free (hence this failure will
+	// cause the previous key to be overwritten).  This behavior may need to
+	// be revisited.
+	for k := range listedKeys {
+		err := origStore.RemoveKey(k)
+		assert.NoError(t, err)
+	}
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, origStore, _importkey,
+		[]string{"GetAttributeValue"}, false)
+
+	assert.Len(t, cleanListKeys(t), 1)
+}
+
+func TestYubiRemoveKeyCleansUpOnError(t *testing.T) {
+	if !YubikeyAccessible() {
+		t.Skip("Must have Yubikey access.")
+	}
+	clearAllKeys(t)
+
+	SetYubikeyKeyMode(KeymodeNone)
+	defer func() {
+		SetYubikeyKeyMode(KeymodeTouch | KeymodePinOnce)
+	}()
+
+	origStore, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
+	assert.NoError(t, err)
+	key, err := testAddKey(t, origStore)
+	assert.NoError(t, err)
+
+	var _removekey = func() error { return origStore.RemoveKey(key.ID()) }
+
+	testYubiFunctionCleansUpOnLoginError(t, origStore, _removekey)
+	// RemoveKey just succeeds if we can't set up the yubikey
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, origStore, _removekey, setupErrors, false)
+	// all the PKCS11 functions RemoveKey depends on that aren't the login/logout
+	// or setup/cleanup
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, origStore, _removekey,
+		[]string{
+			"FindObjectsInit",
+			"FindObjects",
+			"FindObjectsFinal",
+			"DestroyObject",
+		}, true)
+
+	// given that everything should have errored, there should still be 1 key
+	// on the yubikey
+	assert.Len(t, cleanListKeys(t), 1)
+
+	// this will not fail, but it should clean up after itself, and the key
+	// should be added to both stores
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, origStore, _removekey,
+		[]string{"Logout"}, false)
+
+	assert.Len(t, cleanListKeys(t), 0)
+}
+
+func TestYubiListKeyCleansUpOnError(t *testing.T) {
+	if !YubikeyAccessible() {
+		t.Skip("Must have Yubikey access.")
+	}
+	clearAllKeys(t)
+
+	SetYubikeyKeyMode(KeymodeNone)
+	defer func() {
+		SetYubikeyKeyMode(KeymodeTouch | KeymodePinOnce)
+	}()
+
+	// Do not call NewYubiKeyStore, because it list keys immediately to
+	// build the cache.
+	store := &YubiKeyStore{
+		passRetriever: ret,
+		keys:          make(map[string]yubiSlot),
+		backupStore:   trustmanager.NewKeyMemoryStore(ret),
+		libLoader:     defaultLoader,
+	}
+
+	var _listkeys = func() error {
+		// ListKeys never fails
+		store.ListKeys()
+		return nil
+	}
+
+	// all the PKCS11 functions ListKey depends on - list keys never errors
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, store, _listkeys,
+		append(
+			setupErrors,
+			"FindObjectsInit",
+			"FindObjects",
+			"FindObjectsFinal",
+			"GetAttributeValue",
+		), false)
+}
+
+// export key fails anyway, don't bother testing
+
+func TestYubiSignCleansUpOnError(t *testing.T) {
 	if !YubikeyAccessible() {
 		t.Skip("Must have Yubikey access.")
 	}
@@ -380,13 +756,203 @@ func TestYubiSigning(t *testing.T) {
 	store, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
 	assert.NoError(t, err)
 
-	ecdsaPrivateKey, err := testAddKey(t, store)
+	key, err := testAddKey(t, store)
 	assert.NoError(t, err)
 
-	yubiPrivateKey, _, err := store.GetKey(ecdsaPrivateKey.ID())
+	privKey, _, err := store.GetKey(key.ID())
 	assert.NoError(t, err)
 
-	msg := []byte("Hello there")
-	_, err = yubiPrivateKey.Sign(bytes.NewBuffer(msg), msg, nil)
-	assert.NoError(t, err)
+	yubiPrivateKey, ok := privKey.(*YubiPrivateKey)
+	assert.True(t, ok)
+
+	var _sign = func() error {
+		_, err = yubiPrivateKey.Sign(rand.Reader, []byte("Hello there"), nil)
+		return err
+	}
+
+	testYubiFunctionCleansUpOnLoginError(t, yubiPrivateKey, _sign)
+	// all the PKCS11 functions SignKey depends on that is not login/logout
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, yubiPrivateKey, _sign,
+		append(
+			setupErrors,
+			"FindObjectsInit",
+			"FindObjects",
+			"FindObjectsFinal",
+			"SignInit",
+			"Sign",
+		), true)
+	// this will not fail, but it should clean up after itself, and the key
+	// should be added to both stores
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, yubiPrivateKey, _sign,
+		[]string{"Logout"}, false)
+}
+
+// -----  Stubbed pkcs11 for testing error conditions ------
+// This is just a passthrough to the underlying pkcs11 library, with optional
+// error injection.  This is to ensure that if errors occur during the process
+// of interacting with the Yubikey, that everything gets cleaned up sanely.
+
+type errInjected struct {
+	methodName string
+}
+
+func (e errInjected) Error() string {
+	return fmt.Sprintf("Injected failure in %s", e.methodName)
+}
+
+const (
+	uninitialized = 0
+	initialized   = 1
+	sessioned     = 2
+	loggedin      = 3
+)
+
+type StubCtx struct {
+	ctx                IPKCS11Ctx
+	functionShouldFail map[string]bool
+}
+
+func NewStubCtx(functionShouldFail map[string]bool) *StubCtx {
+	realCtx := defaultLoader(pkcs11Lib)
+	return &StubCtx{
+		ctx:                realCtx,
+		functionShouldFail: functionShouldFail,
+	}
+}
+
+// Returns an error if we're supposed to error for this method
+func (s *StubCtx) checkErr(methodName string) error {
+	if val, ok := s.functionShouldFail[methodName]; ok && val {
+		return errInjected{methodName: methodName}
+	}
+	return nil
+}
+
+func (s *StubCtx) Destroy() {
+	// can't error
+	s.ctx.Destroy()
+}
+
+func (s *StubCtx) Initialize() error {
+	err := s.checkErr("Initialize")
+	if err != nil {
+		return err
+	}
+	return s.ctx.Initialize()
+}
+
+func (s *StubCtx) Finalize() error {
+	err := s.checkErr("Finalize")
+	if err != nil {
+		return err
+	}
+	return s.ctx.Finalize()
+}
+
+func (s *StubCtx) GetSlotList(tokenPresent bool) ([]uint, error) {
+	err := s.checkErr("GetSlotList")
+	if err != nil {
+		return nil, err
+	}
+	return s.ctx.GetSlotList(tokenPresent)
+}
+
+func (s *StubCtx) OpenSession(slotID uint, flags uint) (pkcs11.SessionHandle, error) {
+	err := s.checkErr("OpenSession")
+	if err != nil {
+		return pkcs11.SessionHandle(0), err
+	}
+	return s.ctx.OpenSession(slotID, flags)
+}
+
+func (s *StubCtx) CloseSession(sh pkcs11.SessionHandle) error {
+	err := s.checkErr("CloseSession")
+	if err != nil {
+		return err
+	}
+	return s.ctx.CloseSession(sh)
+}
+
+func (s *StubCtx) Login(sh pkcs11.SessionHandle, userType uint, pin string) error {
+	err := s.checkErr("Login")
+	if err != nil {
+		return err
+	}
+	return s.ctx.Login(sh, userType, pin)
+}
+
+func (s *StubCtx) Logout(sh pkcs11.SessionHandle) error {
+	err := s.checkErr("Logout")
+	if err != nil {
+		return err
+	}
+	return s.ctx.Logout(sh)
+}
+
+func (s *StubCtx) CreateObject(sh pkcs11.SessionHandle, temp []*pkcs11.Attribute) (
+	pkcs11.ObjectHandle, error) {
+	err := s.checkErr("CreateObject")
+	if err != nil {
+		return pkcs11.ObjectHandle(0), err
+	}
+	return s.ctx.CreateObject(sh, temp)
+}
+
+func (s *StubCtx) DestroyObject(sh pkcs11.SessionHandle, oh pkcs11.ObjectHandle) error {
+	err := s.checkErr("DestroyObject")
+	if err != nil {
+		return err
+	}
+	return s.ctx.DestroyObject(sh, oh)
+}
+
+func (s *StubCtx) GetAttributeValue(sh pkcs11.SessionHandle, o pkcs11.ObjectHandle,
+	a []*pkcs11.Attribute) ([]*pkcs11.Attribute, error) {
+	err := s.checkErr("GetAttributeValue")
+	if err != nil {
+		return nil, err
+	}
+	return s.ctx.GetAttributeValue(sh, o, a)
+}
+
+func (s *StubCtx) FindObjectsInit(sh pkcs11.SessionHandle, temp []*pkcs11.Attribute) error {
+	err := s.checkErr("FindObjectsInit")
+	if err != nil {
+		return err
+	}
+	return s.ctx.FindObjectsInit(sh, temp)
+}
+
+func (s *StubCtx) FindObjects(sh pkcs11.SessionHandle, max int) (
+	[]pkcs11.ObjectHandle, bool, error) {
+	err := s.checkErr("FindObjects")
+	if err != nil {
+		return nil, false, err
+	}
+	return s.ctx.FindObjects(sh, max)
+}
+
+func (s *StubCtx) FindObjectsFinal(sh pkcs11.SessionHandle) error {
+	err := s.checkErr("FindObjectsFinal")
+	if err != nil {
+		return err
+	}
+	return s.ctx.FindObjectsFinal(sh)
+}
+
+func (s *StubCtx) SignInit(sh pkcs11.SessionHandle, m []*pkcs11.Mechanism,
+	o pkcs11.ObjectHandle) error {
+	err := s.checkErr("SignInit")
+	if err != nil {
+		return err
+	}
+	return s.ctx.SignInit(sh, m, o)
+}
+
+func (s *StubCtx) Sign(sh pkcs11.SessionHandle, message []byte) ([]byte, error) {
+	err := s.checkErr("Sign")
+	if err != nil {
+		return nil, err
+	}
+	return s.ctx.Sign(sh, message)
 }

From 10057562d8f11259178169e159fdd4684538019f Mon Sep 17 00:00:00 2001
From: Ying Li <ying.li@docker.com>
Date: Wed, 11 Nov 2015 16:28:23 -0800
Subject: [PATCH 4/7] Add fixes for Sign (do not continue if SignInit fails).

Signed-off-by: Ying Li <ying.li@docker.com>
Signed-off-by: David Lawrence <david.lawrence@docker.com>

Signed-off-by: Ying Li <ying.li@docker.com> (github: endophage)
---
 trustmanager/yubikey/yubikeystore.go      | 7 ++++++-
 trustmanager/yubikey/yubikeystore_test.go | 9 ++++++++-
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/trustmanager/yubikey/yubikeystore.go b/trustmanager/yubikey/yubikeystore.go
index f1528292f3..1f3a31e579 100644
--- a/trustmanager/yubikey/yubikeystore.go
+++ b/trustmanager/yubikey/yubikeystore.go
@@ -354,7 +354,11 @@ func sign(ctx IPKCS11Ctx, session pkcs11.SessionHandle, pkcs11KeyID []byte, pass
 	}
 
 	var sig []byte
-	ctx.SignInit(session, []*pkcs11.Mechanism{pkcs11.NewMechanism(pkcs11.CKM_ECDSA, nil)}, obj[0])
+	err = ctx.SignInit(
+		session, []*pkcs11.Mechanism{pkcs11.NewMechanism(pkcs11.CKM_ECDSA, nil)}, obj[0])
+	if err != nil {
+		return nil, err
+	}
 
 	// Get the SHA256 of the payload
 	digest := sha256.Sum256(payload)
@@ -363,6 +367,7 @@ func sign(ctx IPKCS11Ctx, session pkcs11.SessionHandle, pkcs11KeyID []byte, pass
 		touchToSignUI()
 		defer touchDoneCallback()
 	}
+	// a call to Sign, whether or not Sign fails, will clear the SignInit
 	sig, err = ctx.Sign(session, digest[:])
 	if err != nil {
 		logrus.Debugf("Error while signing: %s", err)
diff --git a/trustmanager/yubikey/yubikeystore_test.go b/trustmanager/yubikey/yubikeystore_test.go
index 2451dc3d08..70034561b0 100644
--- a/trustmanager/yubikey/yubikeystore_test.go
+++ b/trustmanager/yubikey/yubikeystore_test.go
@@ -792,6 +792,10 @@ func TestYubiSignCleansUpOnError(t *testing.T) {
 // error injection.  This is to ensure that if errors occur during the process
 // of interacting with the Yubikey, that everything gets cleaned up sanely.
 
+// Note that this does not actually replicate an actual PKCS11 failure, since
+// who knows what the pkcs11 function call may have done to the key before it
+// errored. This just tests that we handle an error ok.
+
 type errInjected struct {
 	methodName string
 }
@@ -950,9 +954,12 @@ func (s *StubCtx) SignInit(sh pkcs11.SessionHandle, m []*pkcs11.Mechanism,
 }
 
 func (s *StubCtx) Sign(sh pkcs11.SessionHandle, message []byte) ([]byte, error) {
+	// a call to Sign will clear SignInit whether or not it fails, so
+	// replicate that by calling Sign, then optionally returning an error.
+	sig, sigErr := s.ctx.Sign(sh, message)
 	err := s.checkErr("Sign")
 	if err != nil {
 		return nil, err
 	}
-	return s.ctx.Sign(sh, message)
+	return sig, sigErr
 }

From 38a5b5a342cf7963bef1c9beac77179c63917c20 Mon Sep 17 00:00:00 2001
From: Ying Li <ying.li@docker.com>
Date: Wed, 11 Nov 2015 16:37:58 -0800
Subject: [PATCH 5/7] Add FindObjectsFinalize to getNextEmptySlot.

Signed-off-by: Ying Li <ying.li@docker.com>
Signed-off-by: David Lawrence <david.lawrence@docker.com>

Signed-off-by: Ying Li <ying.li@docker.com> (github: endophage)
---
 trustmanager/yubikey/yubikeystore.go      | 4 ++++
 trustmanager/yubikey/yubikeystore_test.go | 2 ++
 2 files changed, 6 insertions(+)

diff --git a/trustmanager/yubikey/yubikeystore.go b/trustmanager/yubikey/yubikeystore.go
index 1f3a31e579..c7c54d03bd 100644
--- a/trustmanager/yubikey/yubikeystore.go
+++ b/trustmanager/yubikey/yubikeystore.go
@@ -545,6 +545,10 @@ func getNextEmptySlot(ctx IPKCS11Ctx, session pkcs11.SessionHandle) ([]byte, err
 		logrus.Debugf("Failed to find: %s %v", err.Error(), b)
 		return nil, err
 	}
+	if err = ctx.FindObjectsFinal(session); err != nil {
+		logrus.Debugf("Failed to finalize: %s\n", err.Error())
+		return nil, err
+	}
 	for _, obj := range objs {
 		// Retrieve the slot ID
 		attr, err := ctx.GetAttributeValue(session, obj, attrTemplate)
diff --git a/trustmanager/yubikey/yubikeystore_test.go b/trustmanager/yubikey/yubikeystore_test.go
index 70034561b0..3cd151737b 100644
--- a/trustmanager/yubikey/yubikeystore_test.go
+++ b/trustmanager/yubikey/yubikeystore_test.go
@@ -532,6 +532,7 @@ func TestYubiAddKeyCleansUpOnError(t *testing.T) {
 			setupErrors,
 			"FindObjectsInit",
 			"FindObjects",
+			"FindObjectsFinal",
 			"CreateObject",
 		), true)
 
@@ -629,6 +630,7 @@ func TestYubiImportKeyCleansUpOnError(t *testing.T) {
 			setupErrors,
 			"FindObjectsInit",
 			"FindObjects",
+			"FindObjectsFinal",
 			"CreateObject",
 		), true)
 

From cee92fa36336071da217e0f224afc8f37af464a9 Mon Sep 17 00:00:00 2001
From: Ying Li <ying.li@docker.com>
Date: Wed, 11 Nov 2015 16:52:15 -0800
Subject: [PATCH 6/7] Undo some changes from a bad stash pop that were
 unintentional.

Signed-off-by: Ying Li <ying.li@docker.com>
Signed-off-by: David Lawrence <david.lawrence@docker.com>

Signed-off-by: Ying Li <ying.li@docker.com> (github: endophage)
---
 trustmanager/yubikey/yubikeystore.go      | 32 ++++-----
 trustmanager/yubikey/yubikeystore_test.go | 86 +++++++++++------------
 2 files changed, 59 insertions(+), 59 deletions(-)

diff --git a/trustmanager/yubikey/yubikeystore.go b/trustmanager/yubikey/yubikeystore.go
index c7c54d03bd..47a97dc9bd 100644
--- a/trustmanager/yubikey/yubikeystore.go
+++ b/trustmanager/yubikey/yubikeystore.go
@@ -162,10 +162,10 @@ func (y YubiPrivateKey) SignatureAlgorithm() data.SigAlgorithm {
 
 func (y *YubiPrivateKey) Sign(rand io.Reader, msg []byte, opts crypto.SignerOpts) ([]byte, error) {
 	ctx, session, err := SetupHSMEnv(pkcs11Lib, y.libLoader)
-	defer cleanup(ctx, session)
 	if err != nil {
 		return nil, err
 	}
+	defer cleanup(ctx, session)
 
 	v := signed.Verifiers[data.ECDSASignature]
 	for i := 0; i < sigAttempts; i++ {
@@ -621,11 +621,12 @@ func (s *YubiKeyStore) ListKeys() map[string]string {
 		return buildKeyMap(s.keys)
 	}
 	ctx, session, err := SetupHSMEnv(pkcs11Lib, s.libLoader)
-	defer cleanup(ctx, session)
 	if err != nil {
 		logrus.Debugf("Failed to initialize PKCS11 environment: %s", err.Error())
 		return nil
 	}
+	defer cleanup(ctx, session)
+
 	keys, err := yubiListKeys(ctx, session)
 	if err != nil {
 		logrus.Debugf("Failed to list key from the yubikey: %s", err.Error())
@@ -649,11 +650,11 @@ func (s *YubiKeyStore) addKey(
 	}
 
 	ctx, session, err := SetupHSMEnv(pkcs11Lib, s.libLoader)
-	defer cleanup(ctx, session)
 	if err != nil {
 		logrus.Debugf("Failed to initialize PKCS11 environment: %s", err.Error())
 		return err
 	}
+	defer cleanup(ctx, session)
 
 	if k, ok := s.keys[keyID]; ok {
 		if k.role == role {
@@ -692,11 +693,11 @@ func (s *YubiKeyStore) addKey(
 // backup store)
 func (s *YubiKeyStore) GetKey(keyID string) (data.PrivateKey, string, error) {
 	ctx, session, err := SetupHSMEnv(pkcs11Lib, s.libLoader)
-	defer cleanup(ctx, session)
 	if err != nil {
 		logrus.Debugf("Failed to initialize PKCS11 environment: %s", err.Error())
 		return nil, "", err
 	}
+	defer cleanup(ctx, session)
 
 	key, ok := s.keys[keyID]
 	if !ok {
@@ -724,11 +725,12 @@ func (s *YubiKeyStore) GetKey(keyID string) (data.PrivateKey, string, error) {
 // backup store)
 func (s *YubiKeyStore) RemoveKey(keyID string) error {
 	ctx, session, err := SetupHSMEnv(pkcs11Lib, s.libLoader)
-	defer cleanup(ctx, session)
 	if err != nil {
 		logrus.Debugf("Failed to initialize PKCS11 environment: %s", err.Error())
 		return nil
 	}
+	defer cleanup(ctx, session)
+
 	key, ok := s.keys[keyID]
 	if !ok {
 		return errors.New("Key not present in yubikey")
@@ -765,17 +767,15 @@ func (s *YubiKeyStore) ImportKey(pemBytes []byte, keyPath string) error {
 }
 
 func cleanup(ctx IPKCS11Ctx, session pkcs11.SessionHandle) {
-	if ctx != nil {
-		err := ctx.CloseSession(session)
-		if err != nil {
-			logrus.Debugf("Error closing session: %s", err.Error())
-		}
-		err = ctx.Finalize()
-		if err != nil {
-			logrus.Debugf("Error finalizing: %s", err.Error())
-		}
-		ctx.Destroy()
+	err := ctx.CloseSession(session)
+	if err != nil {
+		logrus.Debugf("Error closing session: %s", err.Error())
 	}
+	err = ctx.Finalize()
+	if err != nil {
+		logrus.Debugf("Error finalizing: %s", err.Error())
+	}
+	ctx.Destroy()
 }
 
 // SetupHSMEnv is a method that depends on the existences
@@ -824,10 +824,10 @@ func YubikeyAccessible() bool {
 		return false
 	}
 	ctx, session, err := SetupHSMEnv(pkcs11Lib, defaultLoader)
-	defer cleanup(ctx, session)
 	if err != nil {
 		return false
 	}
+	defer cleanup(ctx, session)
 	return true
 }
 
diff --git a/trustmanager/yubikey/yubikeystore_test.go b/trustmanager/yubikey/yubikeystore_test.go
index 3cd151737b..ab53f2eaea 100644
--- a/trustmanager/yubikey/yubikeystore_test.go
+++ b/trustmanager/yubikey/yubikeystore_test.go
@@ -89,9 +89,9 @@ func TestYubiAddKeysAndRetrieve(t *testing.T) {
 
 	// create 4 keys on the original store
 	backup := trustmanager.NewKeyMemoryStore(ret)
-	origStore, err := NewYubiKeyStore(backup, ret)
+	store, err := NewYubiKeyStore(backup, ret)
 	assert.NoError(t, err)
-	keys := addMaxKeys(t, origStore)
+	keys := addMaxKeys(t, store)
 
 	// create a new store, since we want to be sure the original store's cache
 	// is not masking any issues
@@ -101,7 +101,7 @@ func TestYubiAddKeysAndRetrieve(t *testing.T) {
 	// All 4 keys should be in the original store, in the clean store (which
 	// makes sure the keys are actually on the Yubikey and not on the original
 	// store's cache, and on the backup store)
-	for _, store := range []trustmanager.KeyStore{origStore, cleanStore, backup} {
+	for _, store := range []trustmanager.KeyStore{store, cleanStore, backup} {
 		listedKeys := store.ListKeys()
 		assert.Len(t, listedKeys, numSlots)
 		for _, k := range keys {
@@ -129,12 +129,12 @@ func TestYubiAddKeyFailureIfNoMoreSlots(t *testing.T) {
 
 	// create 4 keys on the original store
 	backup := trustmanager.NewKeyMemoryStore(ret)
-	origStore, err := NewYubiKeyStore(backup, ret)
+	store, err := NewYubiKeyStore(backup, ret)
 	assert.NoError(t, err)
-	addMaxKeys(t, origStore)
+	addMaxKeys(t, store)
 
 	// add another key - should fail because there are no more slots
-	badKey, err := testAddKey(t, origStore)
+	badKey, err := testAddKey(t, store)
 	assert.Error(t, err)
 
 	// create a new store, since we want to be sure the original store's cache
@@ -144,7 +144,7 @@ func TestYubiAddKeyFailureIfNoMoreSlots(t *testing.T) {
 
 	// The key should not be in the original store, in the new clean store, or
 	// in teh backup store.
-	for _, store := range []trustmanager.KeyStore{origStore, cleanStore, backup} {
+	for _, store := range []trustmanager.KeyStore{store, cleanStore, backup} {
 		// the key that wasn't created should not appear in ListKeys or GetKey
 		_, _, err := store.GetKey(badKey.ID())
 		assert.Error(t, err)
@@ -169,16 +169,16 @@ func TestYubiAddKeyCanAddToMiddleSlot(t *testing.T) {
 
 	// create 4 keys on the original store
 	backup := trustmanager.NewKeyMemoryStore(ret)
-	origStore, err := NewYubiKeyStore(backup, ret)
+	store, err := NewYubiKeyStore(backup, ret)
 	assert.NoError(t, err)
-	keys := addMaxKeys(t, origStore)
+	keys := addMaxKeys(t, store)
 
 	// delete one of the middle keys, and assert we can still create a new key
 	keyIDToDelete := keys[numSlots/2]
-	err = origStore.RemoveKey(keyIDToDelete)
+	err = store.RemoveKey(keyIDToDelete)
 	assert.NoError(t, err)
 
-	newKey, err := testAddKey(t, origStore)
+	newKey, err := testAddKey(t, store)
 	assert.NoError(t, err)
 
 	// create a new store, since we want to be sure the original store's cache
@@ -189,7 +189,7 @@ func TestYubiAddKeyCanAddToMiddleSlot(t *testing.T) {
 	// The new key should be in the original store, in the new clean store, and
 	// in the backup store.  The old key should not be in the original store,
 	// or the new clean store.
-	for _, store := range []trustmanager.KeyStore{origStore, cleanStore, backup} {
+	for _, store := range []trustmanager.KeyStore{store, cleanStore, backup} {
 		// new key should appear in all stores
 		gottenKey, _, err := store.GetKey(newKey.ID())
 		assert.NoError(t, err)
@@ -222,12 +222,12 @@ func TestYubiRemoveKey(t *testing.T) {
 	}()
 
 	backup := trustmanager.NewKeyMemoryStore(ret)
-	origStore, err := NewYubiKeyStore(backup, ret)
+	store, err := NewYubiKeyStore(backup, ret)
 	assert.NoError(t, err)
 
-	key, err := testAddKey(t, origStore)
+	key, err := testAddKey(t, store)
 	assert.NoError(t, err)
-	err = origStore.RemoveKey(key.ID())
+	err = store.RemoveKey(key.ID())
 	assert.NoError(t, err)
 
 	// key remains in the backup store
@@ -242,7 +242,7 @@ func TestYubiRemoveKey(t *testing.T) {
 	assert.NoError(t, err)
 
 	// key is not in either the original store or the clean store
-	for _, store := range []*YubiKeyStore{origStore, cleanStore} {
+	for _, store := range []*YubiKeyStore{store, cleanStore} {
 		_, _, err := store.GetKey(key.ID())
 		assert.Error(t, err)
 	}
@@ -261,7 +261,7 @@ func TestYubiImportNewKey(t *testing.T) {
 	}()
 
 	backup := trustmanager.NewKeyMemoryStore(ret)
-	origStore, err := NewYubiKeyStore(backup, ret)
+	store, err := NewYubiKeyStore(backup, ret)
 	assert.NoError(t, err)
 
 	// generate key and import it
@@ -271,7 +271,7 @@ func TestYubiImportNewKey(t *testing.T) {
 	pemBytes, err := trustmanager.EncryptPrivateKey(privKey, "passphrase")
 	assert.NoError(t, err)
 
-	err = origStore.ImportKey(pemBytes, "root")
+	err = store.ImportKey(pemBytes, "root")
 	assert.NoError(t, err)
 
 	// key is not in backup store
@@ -282,7 +282,7 @@ func TestYubiImportNewKey(t *testing.T) {
 	// is not masking any issues
 	cleanStore, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
 	assert.NoError(t, err)
-	for _, store := range []*YubiKeyStore{origStore, cleanStore} {
+	for _, store := range []*YubiKeyStore{store, cleanStore} {
 		gottenKey, role, err := store.GetKey(privKey.ID())
 		assert.NoError(t, err)
 		assert.Equal(t, data.CanonicalRootRole, role)
@@ -303,9 +303,9 @@ func TestYubiImportExistingKey(t *testing.T) {
 		SetYubikeyKeyMode(KeymodeTouch | KeymodePinOnce)
 	}()
 
-	origStore, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
+	store, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
 	assert.NoError(t, err)
-	key, err := testAddKey(t, origStore)
+	key, err := testAddKey(t, store)
 
 	backup := trustmanager.NewKeyMemoryStore(ret)
 	newStore, err := NewYubiKeyStore(backup, ret)
@@ -340,7 +340,7 @@ func TestYubiImportNonRootKey(t *testing.T) {
 	}()
 
 	backup := trustmanager.NewKeyMemoryStore(ret)
-	origStore, err := NewYubiKeyStore(backup, ret)
+	store, err := NewYubiKeyStore(backup, ret)
 	assert.NoError(t, err)
 
 	// generate key and import it
@@ -350,7 +350,7 @@ func TestYubiImportNonRootKey(t *testing.T) {
 	pemBytes, err := trustmanager.EncryptPrivateKey(privKey, "passphrase")
 	assert.NoError(t, err)
 
-	err = origStore.ImportKey(pemBytes, privKey.ID())
+	err = store.ImportKey(pemBytes, privKey.ID())
 	assert.Error(t, err)
 
 	// key is not in backup store
@@ -517,17 +517,17 @@ func TestYubiAddKeyCleansUpOnError(t *testing.T) {
 	}()
 
 	backup := trustmanager.NewKeyMemoryStore(ret)
-	origStore, err := NewYubiKeyStore(backup, ret)
+	store, err := NewYubiKeyStore(backup, ret)
 	assert.NoError(t, err)
 
 	var _addkey = func() error {
-		_, err := testAddKey(t, origStore)
+		_, err := testAddKey(t, store)
 		return err
 	}
 
-	testYubiFunctionCleansUpOnLoginError(t, origStore, _addkey)
+	testYubiFunctionCleansUpOnLoginError(t, store, _addkey)
 	// all the PKCS11 functions AddKey depends on that aren't the login/logout
-	testYubiFunctionCleansUpOnSpecifiedErrors(t, origStore, _addkey,
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, store, _addkey,
 		append(
 			setupErrors,
 			"FindObjectsInit",
@@ -544,7 +544,7 @@ func TestYubiAddKeyCleansUpOnError(t *testing.T) {
 	// Logout should not cause a function failure - it s a cleanup failure,
 	// which shouldn't break anything, and it should clean up after itself.
 	// The key should be added to both stores
-	testYubiFunctionCleansUpOnSpecifiedErrors(t, origStore, _addkey,
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, store, _addkey,
 		[]string{"Logout"}, false)
 
 	listedKeys := cleanListKeys(t)
@@ -556,7 +556,7 @@ func TestYubiAddKeyCleansUpOnError(t *testing.T) {
 	// it's in, we assume its occupied slot is free (hence this failure will
 	// cause the previous key to be overwritten).  This behavior may need to
 	// be revisited.
-	testYubiFunctionCleansUpOnSpecifiedErrors(t, origStore, _addkey,
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, store, _addkey,
 		[]string{"GetAttributeValue"}, false)
 
 	newListedKeys := cleanListKeys(t)
@@ -612,7 +612,7 @@ func TestYubiImportKeyCleansUpOnError(t *testing.T) {
 		SetYubikeyKeyMode(KeymodeTouch | KeymodePinOnce)
 	}()
 
-	origStore, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
+	store, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
 	assert.NoError(t, err)
 
 	privKey, err := trustmanager.GenerateECDSAKey(rand.Reader)
@@ -621,11 +621,11 @@ func TestYubiImportKeyCleansUpOnError(t *testing.T) {
 	pemBytes, err := trustmanager.EncryptPrivateKey(privKey, "passphrase")
 	assert.NoError(t, err)
 
-	var _importkey = func() error { return origStore.ImportKey(pemBytes, "root") }
+	var _importkey = func() error { return store.ImportKey(pemBytes, "root") }
 
-	testYubiFunctionCleansUpOnLoginError(t, origStore, _importkey)
+	testYubiFunctionCleansUpOnLoginError(t, store, _importkey)
 	// all the PKCS11 functions ImportKey depends on that aren't the login/logout
-	testYubiFunctionCleansUpOnSpecifiedErrors(t, origStore, _importkey,
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, store, _importkey,
 		append(
 			setupErrors,
 			"FindObjectsInit",
@@ -641,7 +641,7 @@ func TestYubiImportKeyCleansUpOnError(t *testing.T) {
 	// Logout should not cause a function failure - it s a cleanup failure,
 	// which shouldn't break anything, and it should clean up after itself.
 	// The key should be added to both stores
-	testYubiFunctionCleansUpOnSpecifiedErrors(t, origStore, _importkey,
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, store, _importkey,
 		[]string{"Logout"}, false)
 
 	listedKeys := cleanListKeys(t)
@@ -653,10 +653,10 @@ func TestYubiImportKeyCleansUpOnError(t *testing.T) {
 	// cause the previous key to be overwritten).  This behavior may need to
 	// be revisited.
 	for k := range listedKeys {
-		err := origStore.RemoveKey(k)
+		err := store.RemoveKey(k)
 		assert.NoError(t, err)
 	}
-	testYubiFunctionCleansUpOnSpecifiedErrors(t, origStore, _importkey,
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, store, _importkey,
 		[]string{"GetAttributeValue"}, false)
 
 	assert.Len(t, cleanListKeys(t), 1)
@@ -673,19 +673,19 @@ func TestYubiRemoveKeyCleansUpOnError(t *testing.T) {
 		SetYubikeyKeyMode(KeymodeTouch | KeymodePinOnce)
 	}()
 
-	origStore, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
+	store, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
 	assert.NoError(t, err)
-	key, err := testAddKey(t, origStore)
+	key, err := testAddKey(t, store)
 	assert.NoError(t, err)
 
-	var _removekey = func() error { return origStore.RemoveKey(key.ID()) }
+	var _removekey = func() error { return store.RemoveKey(key.ID()) }
 
-	testYubiFunctionCleansUpOnLoginError(t, origStore, _removekey)
+	testYubiFunctionCleansUpOnLoginError(t, store, _removekey)
 	// RemoveKey just succeeds if we can't set up the yubikey
-	testYubiFunctionCleansUpOnSpecifiedErrors(t, origStore, _removekey, setupErrors, false)
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, store, _removekey, setupErrors, false)
 	// all the PKCS11 functions RemoveKey depends on that aren't the login/logout
 	// or setup/cleanup
-	testYubiFunctionCleansUpOnSpecifiedErrors(t, origStore, _removekey,
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, store, _removekey,
 		[]string{
 			"FindObjectsInit",
 			"FindObjects",
@@ -699,7 +699,7 @@ func TestYubiRemoveKeyCleansUpOnError(t *testing.T) {
 
 	// this will not fail, but it should clean up after itself, and the key
 	// should be added to both stores
-	testYubiFunctionCleansUpOnSpecifiedErrors(t, origStore, _removekey,
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, store, _removekey,
 		[]string{"Logout"}, false)
 
 	assert.Len(t, cleanListKeys(t), 0)

From 4b7fefd5ef65f7823bd4484c3fe165d43ef94d20 Mon Sep 17 00:00:00 2001
From: Ying Li <ying.li@docker.com>
Date: Wed, 11 Nov 2015 19:51:57 -0800
Subject: [PATCH 7/7] Do not clean up a session if there is no session.

Signed-off-by: Ying Li <ying.li@docker.com>
Signed-off-by: David Lawrence <david.lawrence@docker.com>

Signed-off-by: Ying Li <ying.li@docker.com> (github: endophage)
---
 Makefile                                  |  2 +-
 trustmanager/yubikey/yubikeystore.go      | 14 +++++++++-----
 trustmanager/yubikey/yubikeystore_test.go |  3 ++-
 3 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/Makefile b/Makefile
index d6b7dc4c10..a45271ddda 100644
--- a/Makefile
+++ b/Makefile
@@ -15,7 +15,7 @@ GO_LDFLAGS=-ldflags "-w $(CTIMEVAR)"
 GO_LDFLAGS_STATIC=-ldflags "-w $(CTIMEVAR) -extldflags -static"
 GOOSES = darwin freebsd linux
 GOARCHS = amd64
-NOTARY_BUILDTAGS ?= "pkcs11"
+NOTARY_BUILDTAGS ?= pkcs11
 GO_EXC = go
 NOTARYDIR := /go/src/github.com/docker/notary
 
diff --git a/trustmanager/yubikey/yubikeystore.go b/trustmanager/yubikey/yubikeystore.go
index 47a97dc9bd..5df3f430ed 100644
--- a/trustmanager/yubikey/yubikeystore.go
+++ b/trustmanager/yubikey/yubikeystore.go
@@ -771,7 +771,11 @@ func cleanup(ctx IPKCS11Ctx, session pkcs11.SessionHandle) {
 	if err != nil {
 		logrus.Debugf("Error closing session: %s", err.Error())
 	}
-	err = ctx.Finalize()
+	finalizeAndDestroy(ctx)
+}
+
+func finalizeAndDestroy(ctx IPKCS11Ctx) {
+	err := ctx.Finalize()
 	if err != nil {
 		logrus.Debugf("Error finalizing: %s", err.Error())
 	}
@@ -792,18 +796,18 @@ func SetupHSMEnv(libraryPath string, libLoader pkcs11LibLoader) (
 	}
 
 	if err := p.Initialize(); err != nil {
-		defer cleanup(p, 0)
+		defer finalizeAndDestroy(p)
 		return nil, 0, fmt.Errorf("Initialize error %s", err.Error())
 	}
 
 	slots, err := p.GetSlotList(true)
 	if err != nil {
-		defer cleanup(p, 0)
+		defer finalizeAndDestroy(p)
 		return nil, 0, fmt.Errorf("Failed to list HSM slots %s", err)
 	}
 	// Check to see if we got any slots from the HSM.
 	if len(slots) < 1 {
-		defer cleanup(p, 0)
+		defer finalizeAndDestroy(p)
 		return nil, 0, fmt.Errorf("No HSM Slots found")
 	}
 
@@ -811,7 +815,7 @@ func SetupHSMEnv(libraryPath string, libLoader pkcs11LibLoader) (
 	// CKF_RW_SESSION: TRUE if the session is read/write; FALSE if the session is read-only
 	session, err := p.OpenSession(slots[0], pkcs11.CKF_SERIAL_SESSION|pkcs11.CKF_RW_SESSION)
 	if err != nil {
-		defer cleanup(p, 0)
+		defer cleanup(p, session)
 		return nil, 0, fmt.Errorf("Failed to Start Session with HSM %s", err)
 	}
 
diff --git a/trustmanager/yubikey/yubikeystore_test.go b/trustmanager/yubikey/yubikeystore_test.go
index ab53f2eaea..672a61eb35 100644
--- a/trustmanager/yubikey/yubikeystore_test.go
+++ b/trustmanager/yubikey/yubikeystore_test.go
@@ -64,7 +64,7 @@ func testAddKey(t *testing.T, store trustmanager.KeyStore) (data.PrivateKey, err
 }
 
 func addMaxKeys(t *testing.T, store trustmanager.KeyStore) []string {
-	keys := make([]string, 0, numSlots)
+	var keys []string
 	// create the maximum number of keys
 	for i := 0; i < numSlots; i++ {
 		privKey, err := testAddKey(t, store)
@@ -378,6 +378,7 @@ func TestYubiExportKeyFails(t *testing.T) {
 
 	_, err = store.ExportKey(key.ID())
 	assert.Error(t, err)
+	assert.Equal(t, "Keys cannot be exported from a Yubikey.", err.Error())
 }
 
 // If there are keys in the backup store but no keys in the Yubikey,