From 575fdbe86e5dd89229707ebec0575ce7d088a4a6 Mon Sep 17 00:00:00 2001
From: Cedric Staub <cs@squareup.com>
Date: Fri, 30 Oct 2015 09:56:38 -0700
Subject: [PATCH] x/crypto/pkcs12: deal with short byte array in PBKDF

A bug in the PBKDF implementation in the pkcs12 package caused certain
inputs to produce the wrong output. In step 6C, we treat Ij as a big.Int
of v bytes and compute Ij + B + 1. The resulting big.Int can have
leading zero bytes, such that Ij.Bytes() produces a byte array that is
smaller than v bytes. In this case we must zero-pad the output to make
Ijb properly sized.

Change-Id: I6e672e6f6bbbd5bea419355d318db289731bb174
Reviewed-on: https://go-review.googlesource.com/16520
Reviewed-by: Adam Langley <agl@golang.org>
---
 pkcs12/pbkdf.go      | 15 +++++++++++++++
 pkcs12/pbkdf_test.go | 14 +++++++++++++-
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/pkcs12/pbkdf.go b/pkcs12/pbkdf.go
index 580cac7..5c419d4 100644
--- a/pkcs12/pbkdf.go
+++ b/pkcs12/pbkdf.go
@@ -103,6 +103,7 @@ func pbkdf(hash func([]byte) []byte, u, v int, salt, password []byte, r int, ID
 
 	//    6.  For i=1, 2, ..., c, do the following:
 	A := make([]byte, c*20)
+	var IjBuf []byte
 	for i := 0; i < c; i++ {
 		//        A.  Set A2=H^r(D||I). (i.e., the r-th hash of D||1,
 		//            H(H(H(... H(D||I))))
@@ -133,9 +134,23 @@ func pbkdf(hash func([]byte) []byte, u, v int, salt, password []byte, r int, ID
 					Ij.Add(Ij, Bbi)
 					Ij.Add(Ij, one)
 					Ijb := Ij.Bytes()
+					// We expect Ijb to be exactly v bytes,
+					// if it is longer or shorter we must
+					// adjust it accordingly.
 					if len(Ijb) > v {
 						Ijb = Ijb[len(Ijb)-v:]
 					}
+					if len(Ijb) < v {
+						if IjBuf == nil {
+							IjBuf = make([]byte, v)
+						}
+						bytesShort := v - len(Ijb)
+						for i := 0; i < bytesShort; i++ {
+							IjBuf[i] = 0
+						}
+						copy(IjBuf[bytesShort:], Ijb)
+						Ijb = IjBuf
+					}
 					copy(I[j*v:(j+1)*v], Ijb)
 				}
 			}
diff --git a/pkcs12/pbkdf_test.go b/pkcs12/pbkdf_test.go
index 1197859..262037d 100644
--- a/pkcs12/pbkdf_test.go
+++ b/pkcs12/pbkdf_test.go
@@ -17,6 +17,18 @@ func TestThatPBKDFWorksCorrectlyForLongKeys(t *testing.T) {
 	key := cipherInfo.deriveKey(salt, password, 2048)
 
 	if expected := []byte("\x7c\xd9\xfd\x3e\x2b\x3b\xe7\x69\x1a\x44\xe3\xbe\xf0\xf9\xea\x0f\xb9\xb8\x97\xd4\xe3\x25\xd9\xd1"); bytes.Compare(key, expected) != 0 {
-		t.Fatalf("expected key '%x', but found '%x'", key, expected)
+		t.Fatalf("expected key '%x', but found '%x'", expected, key)
+	}
+}
+
+func TestThatPBKDFHandlesLeadingZeros(t *testing.T) {
+	// This test triggers a case where I_j (in step 6C) ends up with leading zero
+	// byte, meaning that len(Ijb) < v (leading zeros get stripped by big.Int).
+	// This was previously causing bug whereby certain inputs would break the
+	// derivation and produce the wrong output.
+	key := pbkdf(sha1Sum, 20, 64, []byte("\xf3\x7e\x05\xb5\x18\x32\x4b\x4b"), []byte("\x00\x00"), 2048, 1, 24)
+	expected := []byte("\x00\xf7\x59\xff\x47\xd1\x4d\xd0\x36\x65\xd5\x94\x3c\xb3\xc4\xa3\x9a\x25\x55\xc0\x2a\xed\x66\xe1")
+	if bytes.Compare(key, expected) != 0 {
+		t.Fatalf("expected key '%x', but found '%x'", expected, key)
 	}
 }