From b1038742086ea8969b12a905941c926955d9e0e1 Mon Sep 17 00:00:00 2001
From: Filip Skokan <panva.ip@gmail.com>
Date: Tue, 3 Dec 2019 15:48:07 +0100
Subject: [PATCH] refactor: improve extra KW header parameter handling

---
 lib/jwa/ecdh/dir.js |  5 +++--
 lib/jwa/ecdh/kw.js  | 10 ++++++----
 lib/jwa/pbes2.js    |  8 +++++---
 3 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/lib/jwa/ecdh/dir.js b/lib/jwa/ecdh/dir.js
index 1a16019b05..7e6023f198 100644
--- a/lib/jwa/ecdh/dir.js
+++ b/lib/jwa/ecdh/dir.js
@@ -17,8 +17,9 @@ const wrapKey = (key, payload, { enc }) => {
   }
 }
 
-const unwrapKey = (key, payload, { apu, apv, epk, enc }) => {
-  return derive(enc, KEYLENGTHS.get(enc), key, epk, { apu, apv })
+const unwrapKey = (key, payload, header) => {
+  const { enc, epk } = header
+  return derive(enc, KEYLENGTHS.get(enc), key, epk, header)
 }
 
 module.exports = (JWA, JWK) => {
diff --git a/lib/jwa/ecdh/kw.js b/lib/jwa/ecdh/kw.js
index ebf8af38ac..f4a525f1d8 100644
--- a/lib/jwa/ecdh/kw.js
+++ b/lib/jwa/ecdh/kw.js
@@ -12,15 +12,17 @@ const wrapKey = (wrap, derive, key, payload) => {
   const derivedKey = derive(epk, key, payload)
 
   const result = wrap({ [KEYOBJECT]: derivedKey }, payload)
-  result.header = { epk: { kty: key.kty, crv: key.crv, x: epk.x, y: epk.y } }
+  result.header = result.header || {}
+  Object.assign(result.header, { epk: { kty: key.kty, crv: key.crv, x: epk.x, y: epk.y } })
 
   return result
 }
 
-const unwrapKey = (unwrap, derive, key, payload, { apu, apv, epk }) => {
-  const derivedKey = derive(key, epk, { apu, apv })
+const unwrapKey = (unwrap, derive, key, payload, header) => {
+  const { epk } = header
+  const derivedKey = derive(key, epk, header)
 
-  return unwrap({ [KEYOBJECT]: derivedKey }, payload)
+  return unwrap({ [KEYOBJECT]: derivedKey }, payload, header)
 }
 
 module.exports = (JWA, JWK) => {
diff --git a/lib/jwa/pbes2.js b/lib/jwa/pbes2.js
index 62b4381607..bd1700fa12 100644
--- a/lib/jwa/pbes2.js
+++ b/lib/jwa/pbes2.js
@@ -27,15 +27,17 @@ const wrapKey = (keylen, sha, concat, wrap, { [KEYOBJECT]: keyObject }, payload)
   const derivedKey = pbkdf2(keyObject.export(), salt, p2c, keylen, sha)
 
   const result = wrap({ [KEYOBJECT]: derivedKey }, payload)
-  result.header = { p2c, p2s: base64url.encodeBuffer(p2s) }
+  result.header = result.header || {}
+  Object.assign(result.header, { p2c, p2s: base64url.encodeBuffer(p2s) })
 
   return result
 }
 
-const unwrapKey = (keylen, sha, concat, unwrap, { [KEYOBJECT]: keyObject }, payload, { p2c, p2s }) => {
+const unwrapKey = (keylen, sha, concat, unwrap, { [KEYOBJECT]: keyObject }, payload, header) => {
+  const { p2s, p2c } = header
   const salt = concat(p2s)
   const derivedKey = pbkdf2(keyObject.export(), salt, p2c, keylen, sha)
-  return unwrap({ [KEYOBJECT]: derivedKey }, payload)
+  return unwrap({ [KEYOBJECT]: derivedKey }, payload, header)
 }
 
 module.exports = (JWA, JWK) => {