From 6cf0643d7d1fa6446c0ebf6a6fdc45bcf1885e7e Mon Sep 17 00:00:00 2001
From: Ying Li <ying.li@docker.com>
Date: Wed, 11 Nov 2015 21:13:35 -0800
Subject: [PATCH] Roll back an add key to the yubikey if we can't back it 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      | 31 ++++++++----------
 trustmanager/yubikey/yubikeystore_test.go | 38 +++++++++++++++++++++++
 2 files changed, 51 insertions(+), 18 deletions(-)

diff --git a/trustmanager/yubikey/yubikeystore.go b/trustmanager/yubikey/yubikeystore.go
index 5df3f430ed..1b933ac1ae 100644
--- a/trustmanager/yubikey/yubikeystore.go
+++ b/trustmanager/yubikey/yubikeystore.go
@@ -200,7 +200,6 @@ func addECDSAKey(
 	pkcs11KeyID []byte,
 	passRetriever passphrase.Retriever,
 	role string,
-	backupStore trustmanager.KeyStore,
 ) error {
 	logrus.Debugf("Attempting to add key to yubikey with ID: %s", privKey.ID())
 
@@ -253,13 +252,6 @@ func addECDSAKey(
 		return fmt.Errorf("error importing: %v", err)
 	}
 
-	if backupStore != nil {
-		err = backupStore.AddKey(privKey.ID(), role, privKey)
-		if err != nil {
-			return ErrBackupFailed{err: err.Error()}
-		}
-	}
-
 	return nil
 }
 
@@ -639,11 +631,19 @@ func (s *YubiKeyStore) ListKeys() map[string]string {
 
 // AddKey puts a key inside the Yubikey, as well as writing it to the backup store
 func (s *YubiKeyStore) AddKey(keyID, role string, privKey data.PrivateKey) error {
-	return s.addKey(keyID, role, privKey, true)
+	err := s.addKey(keyID, role, privKey)
+	if err != nil {
+		return err
+	}
+	err = s.backupStore.AddKey(privKey.ID(), role, privKey)
+	if err != nil {
+		defer s.RemoveKey(keyID)
+		return ErrBackupFailed{err: err.Error()}
+	}
+	return nil
 }
 
-func (s *YubiKeyStore) addKey(
-	keyID, role string, privKey data.PrivateKey, backup bool) error {
+func (s *YubiKeyStore) addKey(keyID, role string, privKey data.PrivateKey) error {
 	// We only allow adding root keys for now
 	if role != data.CanonicalRootRole {
 		return fmt.Errorf("yubikey only supports storing root keys, got %s for key: %s", role, keyID)
@@ -670,13 +670,8 @@ func (s *YubiKeyStore) addKey(
 	}
 	logrus.Debugf("Attempting to store key using yubikey slot %v", slot)
 
-	backupStore := s.backupStore
-	if !backup {
-		backupStore = nil
-	}
-
 	err = addECDSAKey(
-		ctx, session, privKey, slot, s.passRetriever, role, backupStore)
+		ctx, session, privKey, slot, s.passRetriever, role)
 	if err == nil {
 		s.keys[privKey.ID()] = yubiSlot{
 			role:   role,
@@ -763,7 +758,7 @@ func (s *YubiKeyStore) ImportKey(pemBytes []byte, keyPath string) error {
 	if keyPath != data.CanonicalRootRole {
 		return fmt.Errorf("yubikey only supports storing root keys")
 	}
-	return s.addKey(privKey.ID(), "root", privKey, false)
+	return s.addKey(privKey.ID(), "root", privKey)
 }
 
 func cleanup(ctx IPKCS11Ctx, session pkcs11.SessionHandle) {
diff --git a/trustmanager/yubikey/yubikeystore_test.go b/trustmanager/yubikey/yubikeystore_test.go
index 7605ce66a0..a38b0c697d 100644
--- a/trustmanager/yubikey/yubikeystore_test.go
+++ b/trustmanager/yubikey/yubikeystore_test.go
@@ -4,6 +4,7 @@ package yubikey
 
 import (
 	"crypto/rand"
+	"errors"
 	"fmt"
 	"reflect"
 	"testing"
@@ -209,6 +210,43 @@ func TestYubiAddKeyCanAddToMiddleSlot(t *testing.T) {
 	}
 }
 
+type nonworkingBackup struct {
+	trustmanager.KeyMemoryStore
+}
+
+// AddKey stores the contents of a PEM-encoded private key as a PEM block
+func (s *nonworkingBackup) AddKey(name, alias string, privKey data.PrivateKey) error {
+	return errors.New("Nope!")
+}
+
+// If, when adding a key to the Yubikey, we can't back up the key, it should
+// be removed from the Yubikey too because otherwise there is no way for
+// the user to later get a backup of the key.
+func TestYubiAddKeyRollsBackIfCannotBackup(t *testing.T) {
+	if !YubikeyAccessible() {
+		t.Skip("Must have Yubikey access.")
+	}
+	clearAllKeys(t)
+
+	SetYubikeyKeyMode(KeymodeNone)
+	defer func() {
+		SetYubikeyKeyMode(KeymodeTouch | KeymodePinOnce)
+	}()
+
+	backup := &nonworkingBackup{
+		KeyMemoryStore: *trustmanager.NewKeyMemoryStore(ret),
+	}
+	store, err := NewYubiKeyStore(backup, ret)
+	assert.NoError(t, err)
+
+	_, err = testAddKey(t, store)
+	assert.Error(t, err)
+	assert.IsType(t, ErrBackupFailed{}, err)
+
+	// there should be no keys on the yubikey
+	assert.Len(t, cleanListKeys(t), 0)
+}
+
 // RemoveKey removes a key from the yubikey, but not from the backup store.
 func TestYubiRemoveKey(t *testing.T) {
 	if !YubikeyAccessible() {