Merge pull request #51 from docker/more-yubikey-tests

more yubikey tests

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

Signed-off-by: Diogo Mónica <diogo.monica@gmail.com> (github: endophage)

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
 

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)
+}

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),
@@ -286,7 +294,7 @@ func getECDSAKey(ctx *pkcs11.Ctx, session pkcs11.SessionHandle, pkcs11KeyID []by
 		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])
@@ -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)
@@ -346,7 +354,11 @@ func sign(ctx *pkcs11.Ctx, session pkcs11.SessionHandle, pkcs11KeyID []byte, pas
 	}
 
 	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)
@@ -355,6 +367,7 @@ func sign(ctx *pkcs11.Ctx, session pkcs11.SessionHandle, pkcs11KeyID []byte, pas
 		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)
@@ -367,7 +380,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 +421,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),
@@ -456,7 +469,7 @@ func yubiListKeys(ctx *pkcs11.Ctx, session pkcs11.SessionHandle) (keys map[strin
 			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)
@@ -501,7 +514,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),
 	}
@@ -532,6 +545,10 @@ func getNextEmptySlot(ctx *pkcs11.Ctx, session pkcs11.SessionHandle) ([]byte, er
 		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)
@@ -571,6 +588,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 +600,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,16 +612,21 @@ 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
 	}
 	defer cleanup(ctx, session)
+
 	keys, err := yubiListKeys(ctx, session)
 	if err != nil {
 		logrus.Debugf("Failed to list key from the yubikey: %s", err.Error())
@@ -625,7 +649,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 +692,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,12 +724,13 @@ 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
 	}
 	defer cleanup(ctx, session)
+
 	key, ok := s.keys[keyID]
 	if !ok {
 		return errors.New("Key not present in yubikey")
@@ -741,33 +766,48 @@ 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) {
-	ctx.CloseSession(session)
-	ctx.Finalize()
+func cleanup(ctx IPKCS11Ctx, session pkcs11.SessionHandle) {
+	err := ctx.CloseSession(session)
+	if err != nil {
+		logrus.Debugf("Error closing session: %s", err.Error())
+	}
+	finalizeAndDestroy(ctx)
+}
+
+func finalizeAndDestroy(ctx IPKCS11Ctx) {
+	err := ctx.Finalize()
+	if err != nil {
+		logrus.Debugf("Error finalizing: %s", err.Error())
+	}
 	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")
 	}
 
 	if err := p.Initialize(); err != nil {
+		defer finalizeAndDestroy(p)
 		return nil, 0, fmt.Errorf("Initialize error %s", err.Error())
 	}
 
 	slots, err := p.GetSlotList(true)
 	if err != nil {
+		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 finalizeAndDestroy(p)
 		return nil, 0, fmt.Errorf("No HSM Slots found")
 	}
 
@@ -775,6 +815,7 @@ func SetupHSMEnv(libraryPath string) (*pkcs11.Ctx, pkcs11.SessionHandle, error)
 	// 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, session)
 		return nil, 0, fmt.Errorf("Failed to Start Session with HSM %s", err)
 	}
 
@@ -786,7 +827,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 +835,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 {

trustmanager/yubikey/yubikeystore_test.go

@@ -4,20 +4,22 @@ package yubikey
 
 import (
 	"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"
 )
 
+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 +55,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 +63,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)
-	}()
-
-	keys := make([]string, 0, numSlots)
-
+func addMaxKeys(t *testing.T, store trustmanager.KeyStore) []string {
+	var keys []string
 	// 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)
+	store, err := NewYubiKeyStore(backup, ret)
+	assert.NoError(t, err)
+	keys := addMaxKeys(t, store)
+
+	// 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{store, 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)
+	store, err := NewYubiKeyStore(backup, ret)
+	assert.NoError(t, err)
+	addMaxKeys(t, store)
+
+	// add another key - should fail because there are no more slots
+	badKey, err := testAddKey(t, store)
+	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{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)
+		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)
+	store, err := NewYubiKeyStore(backup, ret)
+	assert.NoError(t, err)
+	keys := addMaxKeys(t, store)
+
+	// delete one of the middle keys, and assert we can still create a new key
+	keyIDToDelete := keys[numSlots/2]
+	err = store.RemoveKey(keyIDToDelete)
+	assert.NoError(t, err)
+
+	newKey, err := testAddKey(t, store)
+	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{store, 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)
+	store, err := NewYubiKeyStore(backup, ret)
+	assert.NoError(t, err)
+
+	key, err := testAddKey(t, store)
+	assert.NoError(t, err)
+	err = store.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{store, cleanStore} {
+		_, _, err := store.GetKey(key.ID())
+		assert.Error(t, err)
+	}
+}
+
+// ImportKey imports a key as root without adding it to the backup store
+func TestYubiImportNewKey(t *testing.T) {
+	if !YubikeyAccessible() {
+		t.Skip("Must have Yubikey access.")
+	}
+	clearAllKeys(t)
+
+	SetYubikeyKeyMode(KeymodeNone)
+	defer func() {
+		SetYubikeyKeyMode(KeymodeTouch | KeymodePinOnce)
+	}()
+
+	backup := trustmanager.NewKeyMemoryStore(ret)
+	store, err := NewYubiKeyStore(backup, ret)
+	assert.NoError(t, err)
+
 	// generate key and import it
 	privKey, err := trustmanager.GenerateECDSAKey(rand.Reader)
 	assert.NoError(t, err)
@@ -137,11 +278,691 @@ func TestImportKey(t *testing.T) {
 	_, _, 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)
-	assert.Equal(t, data.CanonicalRootRole, role)
-	assert.Equal(t, privKey.Public(), gottenKey.Public())
+	for _, store := range []*YubiKeyStore{store, 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())
+	}
+}
+
+// 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)
+	}()
+
+	store, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
+	assert.NoError(t, err)
+	key, err := testAddKey(t, store)
+
+	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)
+	store, 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 = store.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() {
+		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)
+	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,
+// 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.  Assume we can
+// sign something.
+func TestYubiKeyAndSign(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())
+
+	// 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)
+}
+
+// ----- 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)
+	store, err := NewYubiKeyStore(backup, ret)
+	assert.NoError(t, err)
+
+	var _addkey = func() error {
+		_, err := testAddKey(t, store)
+		return err
+	}
+
+	testYubiFunctionCleansUpOnLoginError(t, store, _addkey)
+	// all the PKCS11 functions AddKey depends on that aren't the login/logout
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, store, _addkey,
+		append(
+			setupErrors,
+			"FindObjectsInit",
+			"FindObjects",
+			"FindObjectsFinal",
+			"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, store, _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, store, _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)
+	}()
+
+	store, 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 store.ImportKey(pemBytes, "root") }
+
+	testYubiFunctionCleansUpOnLoginError(t, store, _importkey)
+	// all the PKCS11 functions ImportKey depends on that aren't the login/logout
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, store, _importkey,
+		append(
+			setupErrors,
+			"FindObjectsInit",
+			"FindObjects",
+			"FindObjectsFinal",
+			"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, store, _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 := store.RemoveKey(k)
+		assert.NoError(t, err)
+	}
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, store, _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)
+	}()
+
+	store, err := NewYubiKeyStore(trustmanager.NewKeyMemoryStore(ret), ret)
+	assert.NoError(t, err)
+	key, err := testAddKey(t, store)
+	assert.NoError(t, err)
+
+	var _removekey = func() error { return store.RemoveKey(key.ID()) }
+
+	testYubiFunctionCleansUpOnLoginError(t, store, _removekey)
+	// RemoveKey just succeeds if we can't set up the yubikey
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, store, _removekey, setupErrors, false)
+	// all the PKCS11 functions RemoveKey depends on that aren't the login/logout
+	// or setup/cleanup
+	testYubiFunctionCleansUpOnSpecifiedErrors(t, store, _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, store, _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.")
+	}
+	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)
+
+	privKey, _, err := store.GetKey(key.ID())
+	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.
+
+// 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
+}
+
+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) {
+	// 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 sig, sigErr
 }