From 4259bca35a5bfedcdf7f79582182b1edeeb203a9 Mon Sep 17 00:00:00 2001
From: Quantu <673216+Quantu@users.noreply.github.com>
Date: Tue, 9 Jul 2024 18:24:16 -0400
Subject: [PATCH] Misc small fixes for PR149 comments that were lost

---
 README.md     | 12 ++++++------
 v2/piv/piv.go | 15 +++++++++++++--
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index d19fb84..d7d1c38 100644
--- a/README.md
+++ b/README.md
@@ -75,7 +75,7 @@ if err != nil {
 The PIV applet has three unique credentials:
 
 * Management key (3DES key) used to generate new keys on the YubiKey.
-  * YubiKey firmware 5.4.0+ adds support for AES128/192/256 keys
+	* YubiKey firmware 5.4.0+ adds support for AES128/192/256 keys
 * PIN (up to 8 digits, usually 6) used to access signing operations.
 * PUK (up to 8 digits) used to unblock the PIN. Usually set once and thrown
   away or managed by an administrator.
@@ -95,8 +95,8 @@ newPUKInt, err := rand.Int(rand.Reader, big.NewInt(100_000_000))
 if err != nil {
 	// ...
 }
-var newKey [24]byte
-if _, err := io.ReadFull(rand.Reader, newKey[:]); err != nil {
+newKey := make([]byte, 24)
+if _, err := io.ReadFull(rand.Reader, newKey); err != nil {
 	// ...
 }
 // Format with leading zeros.
@@ -104,7 +104,7 @@ newPIN := fmt.Sprintf("%06d", newPINInt)
 newPUK := fmt.Sprintf("%08d", newPUKInt)
 
 // Set all values to a new value.
-if err := yk.SetManagementKey(piv.DefaultManagementKey, newKey[:]); err != nil {
+if err := yk.SetManagementKey(piv.DefaultManagementKey, newKey); err != nil {
 	// ...
 }
 if err := yk.SetPUK(piv.DefaultPUK, newPUK); err != nil {
@@ -114,8 +114,8 @@ if err := yk.SetPIN(piv.DefaultPIN, newPIN); err != nil {
 	// ...
 }
 // Store management key on the YubiKey.
-m := piv.Metadata{ManagementKey: &newKey[:]}
-if err := yk.SetMetadata(newKey[:], m); err != nil {
+m := piv.Metadata{ManagementKey: &newKey}
+if err := yk.SetMetadata(newKey, m); err != nil {
 	// ...
 }
 
diff --git a/v2/piv/piv.go b/v2/piv/piv.go
index e3356dd..f0ad16e 100644
--- a/v2/piv/piv.go
+++ b/v2/piv/piv.go
@@ -369,6 +369,13 @@ var (
 	aidYubiKey    = [...]byte{0xa0, 0x00, 0x00, 0x05, 0x27, 0x20, 0x01, 0x01}
 )
 
+var managementKeyLengthMap = map[byte]int{
+	alg3DES:   24,
+	algAES128: 16,
+	algAES192: 24,
+	algAES256: 32,
+}
+
 func ykAuthenticate(tx *scTx, key []byte, rand io.Reader, version *version) error {
 	// https://nvlpubs.nist.gov/nistpubs/SpecialPublications/NIST.SP.800-73-4.pdf#page=92
 	// https://tsapps.nist.gov/publication/get_pdf.cfm?pub_id=918402#page=114
@@ -382,9 +389,10 @@ func ykAuthenticate(tx *scTx, key []byte, rand io.Reader, version *version) erro
 			param2:      keyCardManagement,
 		}
 		resp, err := tx.Transmit(cmd)
-		if err == nil {
-			managementKeyType = resp[2:][0]
+		if err != nil {
+			return fmt.Errorf("determining key management type: %w", err)
 		}
+		managementKeyType = resp[2:][0]
 	}
 
 	// set challengeLength based on managementKeyType
@@ -397,6 +405,9 @@ func ykAuthenticate(tx *scTx, key []byte, rand io.Reader, version *version) erro
 		managementKeyType = alg3DES
 		challengeLength = 8
 	}
+	if len(key) != managementKeyLengthMap[managementKeyType] {
+		return fmt.Errorf("invalid management key length: %d bytes (expected %d)", len(key), managementKeyLengthMap[managementKeyType])
+	}
 
 	// request a witness
 	cmd := apdu{