From d8fd0386a9ad56895caeca711928e583a4320355 Mon Sep 17 00:00:00 2001
From: Cedric Staub <cs@squareup.com>
Date: Fri, 31 May 2019 14:56:54 -0700
Subject: [PATCH] Pad z value to proper size after P-521 scalar multiplication.

Fixes issue #228. After calling ScalarMult for P-521, the output can
sometimes be 65 bytes long instead of 66 bytes. This happens when the
first bit of the computed value is zero. Calling z.Bytes() on the big
integer will then omit the leading zero, giving us a 65-byte value.

This subsequently causes the shared secret computation to be incorrect
as the input into the KDF function should always be 66 bytes which is
the full length for a P-521 coordinate value.
---
 cipher/ecdh_es.go | 24 ++++++++++++++++++++++--
 jwe_test.go       | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/cipher/ecdh_es.go b/cipher/ecdh_es.go
index c128e327..4994a66c 100644
--- a/cipher/ecdh_es.go
+++ b/cipher/ecdh_es.go
@@ -17,8 +17,10 @@
 package josecipher
 
 import (
+	"bytes"
 	"crypto"
 	"crypto/ecdsa"
+	"crypto/elliptic"
 	"encoding/binary"
 )
 
@@ -44,16 +46,34 @@ func DeriveECDHES(alg string, apuData, apvData []byte, priv *ecdsa.PrivateKey, p
 		panic("public key not on same curve as private key")
 	}
 
-	z, _ := priv.PublicKey.Curve.ScalarMult(pub.X, pub.Y, priv.D.Bytes())
-	reader := NewConcatKDF(crypto.SHA256, z.Bytes(), algID, ptyUInfo, ptyVInfo, supPubInfo, []byte{})
+	z, _ := priv.Curve.ScalarMult(pub.X, pub.Y, priv.D.Bytes())
+	zBytes := z.Bytes()
 
+	octSize := dSize(priv.Curve)
+	if len(zBytes) != octSize {
+		zBytes = append(bytes.Repeat([]byte{0}, octSize-len(zBytes)), zBytes...)
+	}
+
+	reader := NewConcatKDF(crypto.SHA256, zBytes, algID, ptyUInfo, ptyVInfo, supPubInfo, []byte{})
 	key := make([]byte, size)
 
 	// Read on the KDF will never fail
 	_, _ = reader.Read(key)
+
 	return key
 }
 
+// dSize returns the size in octets for a coordinate on a elliptic curve.
+func dSize(curve elliptic.Curve) int {
+	order := curve.Params().P
+	bitLen := order.BitLen()
+	size := bitLen / 8
+	if bitLen%8 != 0 {
+		size++
+	}
+	return size
+}
+
 func lengthPrefixed(data []byte) []byte {
 	out := make([]byte, len(data)+4)
 	binary.BigEndian.PutUint32(out, uint32(len(data)))
diff --git a/jwe_test.go b/jwe_test.go
index c5c5f969..bba8c326 100644
--- a/jwe_test.go
+++ b/jwe_test.go
@@ -550,6 +550,41 @@ func TestSampleJose4jJWEMessagesECDH(t *testing.T) {
 	}
 }
 
+func TestPrecomputedECDHMessagesFromJose4j(t *testing.T) {
+	data := []struct{ key, message string }{
+		{
+			`{"kty":"EC","x":"fXx-DfOsmecjKh3VrLZFsF98Z1nutsL4UdFTdgA8S7Y","y":"LGzyJY99aqKk52UIExcNFSTs0S7HnNzQ-DRWBTHDad4","crv":"P-256","d":"OeVCWbXuFuJ9U16q7bhLNoKPLLnK-yTx95grzfvQ2l4"}`,
+			`eyJlbmMiOiJBMjU2Q0JDLUhTNTEyIiwiYWxnIjoiRUNESC1FUyIsImVwayI6eyJrdHkiOiJFQyIsIngiOiJ3ZlRHNVFHZkItNHUxanVUUEN1aTNESXhFTV82ZUs5ZEk5TXNZckpxWDRnIiwieSI6Ik8yanlRbHQ2TXFGTGtqMWFCWW1aNXZJWHFVRHh6Ulk3dER0WmdZUUVNa0kiLCJjcnYiOiJQLTI1NiJ9fQ..mk4wQzGSSeZ8uSgEYTIetA.fCw3-TosL4p0D5fEXw0bEA.9mPsdmGTVoVexXqEOdN5VUKk-ZNtfOtUfbdjVHoko_o`,
+		},
+		{
+			`{"kty":"EC","x":"nBr92fh2JsEjIF1LR5PKICBeHNIBe0xb7nlBrrU3WoWgfJYfXve1jxC-5VT5EPLt","y":"sUAxL3L5lJdzFUSR9EHLniuBhEbvXfPa_3OiR6Du0_GOlFXXIi4UmbNpk10_Thfq","crv":"P-384","d":"0f0NnWg__Qgqjj3fl2gAlsID4Ni41FR88cmZPVgb6ch-ZShuVJRjoxymCuzVP7Gi"}`,
+			`eyJlbmMiOiJBMTkyQ0JDLUhTMzg0IiwiYWxnIjoiRUNESC1FUyIsImVwayI6eyJrdHkiOiJFQyIsIngiOiJsX3hXdzIyb1NfOWZGbV96amNzYkstd3R3d0RHSlRQLUxnNFVBWDI3WWF1b1YwNml2emwtcm1ra2h6ci11SDBmIiwieSI6IloyYmVnbzBqeE9nY0YtNVp4SFNBOU5jZDVCOW8wUE1pSVlRbm9sWkNQTHA3YndPd1RLUEZaaFZVUlFPSjdoeUciLCJjcnYiOiJQLTM4NCJ9fQ..jSWP7pfa4KcpqKWZ1x8awg.osb-5641Ej1Uon_f3U8bNw.KUQWwb35Gxq3YQ34_AVkebugx4rxq1lO`,
+		},
+		{
+			`{"kty":"EC","x":"AH3rqSYjKue50ThW0qq_qQ76cNtqWrc7hU6kZR6akxy8iTf8ugcpqnbgbi98AgSwIqgJZDBMCk-8eoiGaf3R_kDD","y":"AeafPdJjHLf6pK5V7iyMsL3-6MShpHS6jXQ8m-Bcbp06yxAMn6TJbdkacvj45dy_pdh1s6XZwoxRxNETg_gj-hq9","crv":"P-521","d":"AB2tm9vgGe2BaxZmJQ016GY-U7NV_EWhrPsLDC5l9tAM9DGEwI2cT2HcO20Z6CQndw0ZhqLZ6MEvS8siL-SCxIl2"}`,
+			`eyJlbmMiOiJBMjU2Q0JDLUhTNTEyIiwiYWxnIjoiRUNESC1FUyIsImVwayI6eyJrdHkiOiJFQyIsIngiOiJBQ1RLMlVPSjJ6SVk3U1U4T0xkaG1QQmE4ZUVpd2JrX09UMXE0MHBsRlRwQmJKUXg3YWdqWG9LYml2NS1OTXB6eXZySm1rblM3SjNRUWlUeFgwWmtjemhEIiwieSI6IkFXeTZCR1dkZld2ekVNeGIxQklCQnZmRDJ4bEh6Rjk2YzVVRVQ4SFBUS0RSeUJyMnQ4T2dTX1J2MnNoUmxGbXlqUWpyX25uQk94akcxVTZNWDNlZ2VETzciLCJjcnYiOiJQLTUyMSJ9fQ..EWqSGntxbO_Y_6JRjFkCgg.DGjDNjAYdsnYTpUFJi1gEI4YtNd7gBPMjD3CDH047RAwZKTme6Ah_ztzxSfVg5kG.yGm5jn2LtbFXaK_yf0b0932sI2O77j2gwmL1Y09YC_Y`,
+		},
+	}
+
+	for i, vector := range data {
+		var jwk JSONWebKey
+		err := jwk.UnmarshalJSON([]byte(vector.key))
+		if err != nil {
+			t.Fatal(i, err)
+		}
+
+		parsed, err := ParseEncrypted(vector.message)
+		if err != nil {
+			t.Fatal(i, err)
+		}
+
+		_, err = parsed.Decrypt(jwk)
+		if err != nil {
+			t.Fatal(i, err)
+		}
+	}
+}
+
 func TestSampleAESCBCHMACMessagesFromNodeJose(t *testing.T) {
 	samples := []struct {
 		key        []byte