From bf7b63063bf06423aea823170721a3ac838b4562 Mon Sep 17 00:00:00 2001
From: Eric Chiang <ericchiang@google.com>
Date: Sat, 25 Apr 2020 19:59:49 -0700
Subject: [PATCH] piv: fix pin prompt (again)

---
 piv/key.go      | 34 ++++++++++++++--------------------
 piv/key_test.go |  4 ++--
 piv/piv.go      |  6 ++++++
 piv/piv_test.go | 14 ++++++++++++++
 4 files changed, 36 insertions(+), 22 deletions(-)

diff --git a/piv/key.go b/piv/key.go
index 6779ee7..032b631 100644
--- a/piv/key.go
+++ b/piv/key.go
@@ -534,27 +534,21 @@ func isAuthErr(err error) bool {
 }
 
 func (k KeyAuth) do(yk *YubiKey, f func(tx *scTx) ([]byte, error)) ([]byte, error) {
-	data, err := f(yk.tx)
-	if err == nil {
-		return data, nil
-	}
-	if !isAuthErr(err) {
-		return nil, err
-	}
-
-	pin := k.PIN
-	if pin == "" && k.PINPrompt != nil {
-		p, err := k.PINPrompt()
-		if err != nil {
-			return nil, fmt.Errorf("pin prompt: %v", err)
+	if ykLoginNeeded(yk.tx) {
+		pin := k.PIN
+		if pin == "" && k.PINPrompt != nil {
+			p, err := k.PINPrompt()
+			if err != nil {
+				return nil, fmt.Errorf("pin prompt: %v", err)
+			}
+			pin = p
 		}
-		pin = p
-	}
-	if pin == "" {
-		return nil, err
-	}
-	if err := ykLogin(yk.tx, pin); err != nil {
-		return nil, err
+		if pin != "" {
+			if err := ykLogin(yk.tx, pin); err != nil {
+				return nil, err
+			}
+		}
+		// If PIN isn't provided, assume this is for a PINPolicyNever key.
 	}
 	return f(yk.tx)
 }
diff --git a/piv/key_test.go b/piv/key_test.go
index b546ee9..3aa3240 100644
--- a/piv/key_test.go
+++ b/piv/key_test.go
@@ -524,7 +524,7 @@ func TestYubiKeyPrivateKeyPINError(t *testing.T) {
 	key := Key{
 		Algorithm:   alg,
 		TouchPolicy: TouchPolicyNever,
-		PINPolicy:   PINPolicyNever,
+		PINPolicy:   PINPolicyAlways,
 	}
 	pub, err := yk.GenerateKey(DefaultManagementKey, slot, key)
 	if err != nil {
@@ -548,7 +548,7 @@ func TestYubiKeyPrivateKeyPINError(t *testing.T) {
 
 	b := sha256.Sum256([]byte("hello"))
 	hash := b[:]
-	if _, err := signer.Sign(rand.Reader, hash, crypto.SHA256); err != nil {
+	if _, err := signer.Sign(rand.Reader, hash, crypto.SHA256); err == nil {
 		t.Errorf("expected sign to fail with pin prompt that returned error")
 	}
 }
diff --git a/piv/piv.go b/piv/piv.go
index 99c3a6a..5bb70b8 100644
--- a/piv/piv.go
+++ b/piv/piv.go
@@ -225,6 +225,12 @@ func ykLogin(tx *scTx, pin string) error {
 	return nil
 }
 
+func ykLoginNeeded(tx *scTx) bool {
+	cmd := apdu{instruction: insVerify, param2: 0x80}
+	_, err := tx.Transmit(cmd)
+	return err != nil
+}
+
 // Retries returns the number of attempts remaining to enter the correct PIN.
 func (yk *YubiKey) Retries() (int, error) {
 	return ykPINRetries(yk.tx)
diff --git a/piv/piv_test.go b/piv/piv_test.go
index 20bac0c..f9ab19e 100644
--- a/piv/piv_test.go
+++ b/piv/piv_test.go
@@ -89,6 +89,20 @@ func TestYubiKeySerial(t *testing.T) {
 	}
 }
 
+func TestYubiKeyLoginNeeded(t *testing.T) {
+	yk, close := newTestYubiKey(t)
+	defer close()
+	if !ykLoginNeeded(yk.tx) {
+		t.Errorf("expected login needed")
+	}
+	if err := ykLogin(yk.tx, DefaultPIN); err != nil {
+		t.Fatalf("login: %v", err)
+	}
+	if ykLoginNeeded(yk.tx) {
+		t.Errorf("expected no login needed")
+	}
+}
+
 func TestYubiKeyPINRetries(t *testing.T) {
 	yk, close := newTestYubiKey(t)
 	defer close()