From a2c3e70f091ac73fa740b0d1b123fb9f3b6157d7 Mon Sep 17 00:00:00 2001
From: Christopher Jeffrey <chjjeffrey@gmail.com>
Date: Mon, 4 Sep 2017 13:00:11 -0700
Subject: [PATCH] bip151: use buffer pool.

---
 lib/net/bip151.js | 37 ++++++++++++++++---------------------
 1 file changed, 16 insertions(+), 21 deletions(-)

diff --git a/lib/net/bip151.js b/lib/net/bip151.js
index d01516b4..d6a91ec2 100644
--- a/lib/net/bip151.js
+++ b/lib/net/bip151.js
@@ -47,7 +47,6 @@ const HIGH_WATERMARK = 1024 * (1 << 20);
  * @property {Buffer} publicKey
  * @property {Buffer} privateKey
  * @property {Number} cipher
- * @property {Buffer} prk
  * @property {Buffer} k1
  * @property {Buffer} k2
  * @property {Buffer} sid
@@ -66,8 +65,6 @@ function BIP151Stream(cipher) {
   this.cipher = BIP151.ciphers.CHACHAPOLY;
   this.privateKey = secp256k1.generatePrivateKey();
   this.publicKey = null;
-  this.secret = null;
-  this.prk = null;
   this.k1 = null;
   this.k2 = null;
   this.sid = null;
@@ -95,18 +92,23 @@ function BIP151Stream(cipher) {
  */
 
 BIP151Stream.prototype.init = function init(publicKey) {
-  const bw = new StaticWriter(33);
+  assert(Buffer.isBuffer(publicKey));
 
   this.publicKey = publicKey;
-  this.secret = secp256k1.ecdh(this.publicKey, this.privateKey);
 
-  bw.writeBytes(this.secret);
+  const secret = secp256k1.ecdh(this.publicKey, this.privateKey);
+
+  const bw = StaticWriter.pool(33);
+
+  bw.writeBytes(secret);
   bw.writeU8(this.cipher);
 
-  this.prk = hkdf.extract(bw.render(), HKDF_SALT, 'sha256');
-  this.k1 = hkdf.expand(this.prk, INFO_KEY1, 32, 'sha256');
-  this.k2 = hkdf.expand(this.prk, INFO_KEY2, 32, 'sha256');
-  this.sid = hkdf.expand(this.prk, INFO_SID, 32, 'sha256');
+  const data = bw.render();
+  const prk = hkdf.extract(data, HKDF_SALT, 'sha256');
+
+  this.k1 = hkdf.expand(prk, INFO_KEY1, 32, 'sha256');
+  this.k2 = hkdf.expand(prk, INFO_KEY2, 32, 'sha256');
+  this.sid = hkdf.expand(prk, INFO_SID, 32, 'sha256');
 
   this.seq = 0;
 
@@ -146,18 +148,11 @@ BIP151Stream.prototype.shouldRekey = function shouldRekey(packet) {
  */
 
 BIP151Stream.prototype.rekey = function rekey(k1, k2) {
-  assert(this.prk, 'Cannot rekey before initialization.');
+  assert(this.sid, 'Cannot rekey before initialization.');
 
   if (!k1) {
-    const seed = Buffer.allocUnsafe(64);
-
-    this.sid.copy(seed, 0);
-
-    this.k1.copy(seed, 32);
-    this.k1 = digest.hash256(seed);
-
-    this.k2.copy(seed, 32);
-    this.k2 = digest.hash256(seed);
+    this.k1 = digest.root256(this.sid, this.k1);
+    this.k2 = digest.root256(this.sid, this.k2);
   } else {
     this.k1 = k1;
     this.k2 = k2;
@@ -389,7 +384,7 @@ BIP151.prototype.toEncinit = function toEncinit() {
  */
 
 BIP151.prototype.toEncack = function toEncack() {
-  assert(this.output.prk, 'Cannot ack before init.');
+  assert(this.output.sid, 'Cannot ack before init.');
   assert(!this.ackSent, 'Cannot ack twice.');
   this.ackSent = true;