From fcae101c0f89876bd13b3a1b1937e95ebd871ab3 Mon Sep 17 00:00:00 2001
From: Christopher Jeffrey <chjjeffrey@gmail.com>
Date: Sat, 14 Jan 2017 23:19:46 -0800
Subject: [PATCH] bip151: extra assertions. refactor.

---
 lib/net/bip151.js | 60 +++++++++++++++++++++++++++++++----------------
 1 file changed, 40 insertions(+), 20 deletions(-)

diff --git a/lib/net/bip151.js b/lib/net/bip151.js
index 9d21f68e..f69fbade 100644
--- a/lib/net/bip151.js
+++ b/lib/net/bip151.js
@@ -60,7 +60,7 @@ function BIP151Stream(cipher) {
   if (!(this instanceof BIP151Stream))
     return new BIP151Stream(cipher);
 
-  this.cipher = cipher || 0;
+  this.cipher = BIP151.ciphers.CHACHAPOLY;
   this.privateKey = ec.generatePrivateKey();
   this.publicKey = null;
   this.secret = null;
@@ -69,7 +69,10 @@ function BIP151Stream(cipher) {
   this.k2 = null;
   this.sid = null;
 
-  assert(this.cipher === 0, 'Unknown cipher type.');
+  if (cipher != null) {
+    assert(cipher === BIP151.ciphers.CHACHAPOLY, 'Unknown cipher type.');
+    this.cipher = cipher;
+  }
 
   this.chacha = new chachapoly.ChaCha20();
   this.aead = new chachapoly.AEAD();
@@ -115,14 +118,14 @@ BIP151Stream.prototype.init = function init(publicKey) {
 /**
  * Add buffer size to `processed`,
  * check whether we need to rekey.
- * @param {Buffer} data
+ * @param {Buffer} packet
  * @returns {Boolean}
  */
 
-BIP151Stream.prototype.shouldRekey = function shouldRekey(data) {
+BIP151Stream.prototype.shouldRekey = function shouldRekey(packet) {
   var now = util.now();
 
-  this.processed += data.length;
+  this.processed += packet.length;
 
   if (now >= this.lastRekey + 10
       || this.processed >= HIGH_WATERMARK) {
@@ -159,6 +162,9 @@ BIP151Stream.prototype.rekey = function rekey(k1, k2) {
     this.k2 = k2;
   }
 
+  assert(this.k1);
+  assert(this.k2);
+
   // All state is reinitialized
   // aside from the sequence number.
   this.chacha.init(this.k1, this.iv);
@@ -207,14 +213,12 @@ BIP151Stream.prototype.getPublicKey = function getPublicKey() {
 
 /**
  * Encrypt a payload size with k1.
- * @param {Number} size
+ * @param {Buffer} data
  * @returns {Buffer}
  */
 
-BIP151Stream.prototype.encryptSize = function encryptSize(size) {
-  var data = new Buffer(4);
-  data.writeUInt32LE(size, 0, true);
-  return this.chacha.encrypt(data);
+BIP151Stream.prototype.encryptSize = function encryptSize(data) {
+  return this.chacha.encrypt(data.slice(0, 4));
 };
 
 /**
@@ -309,8 +313,6 @@ function BIP151(cipher) {
   this.ackReceived = false;
   this.initSent = false;
   this.ackSent = false;
-  this.timeout = null;
-  this.job = null;
   this.completed = false;
   this.handshake = false;
 
@@ -319,11 +321,23 @@ function BIP151(cipher) {
   this.waiting = 4;
   this.hasSize = false;
 
+  this.timeout = null;
+  this.job = null;
+
   this.bip150 = null;
 }
 
 util.inherits(BIP151, EventEmitter);
 
+/**
+ * Cipher list.
+ * @enum {Number}
+ */
+
+BIP151.ciphers = {
+  CHACHAPOLY: 0
+};
+
 /**
  * Max message size.
  * @const {Number}
@@ -378,6 +392,7 @@ BIP151.prototype.toEncack = function toEncack() {
   this.ackSent = true;
 
   if (this.isReady()) {
+    assert(!this.completed, 'No encack after timeout.');
     this.handshake = true;
     this.emit('handshake');
   }
@@ -405,6 +420,7 @@ BIP151.prototype.toRekey = function toRekey() {
 BIP151.prototype.encinit = function encinit(publicKey, cipher) {
   assert(cipher === this.output.cipher, 'Cipher mismatch.');
   assert(!this.initReceived, 'Already initialized.');
+  assert(!this.completed, 'No encinit after timeout.');
   this.initReceived = true;
   this.output.init(publicKey);
 };
@@ -431,6 +447,7 @@ BIP151.prototype.encack = function encack(publicKey) {
   }
 
   assert(!this.ackReceived, 'Already ACKed.');
+  assert(!this.completed, 'No encack after timeout.');
   this.ackReceived = true;
 
   this.input.init(publicKey);
@@ -532,11 +549,11 @@ BIP151.prototype.destroy = function destroy() {
 /**
  * Add buffer size to `processed`,
  * check whether we need to rekey.
- * @param {Buffer} data
+ * @param {Buffer} packet
  */
 
-BIP151.prototype.maybeRekey = function maybeRekey(data) {
-  if (!this.output.shouldRekey(data))
+BIP151.prototype.maybeRekey = function maybeRekey(packet) {
+  if (!this.output.shouldRekey(packet))
     return;
 
   this.emit('rekey');
@@ -576,22 +593,23 @@ BIP151.prototype.packetSize = function packetSize(cmd, body) {
 BIP151.prototype.packet = function packet(cmd, body) {
   var size = this.packetSize(cmd, body);
   var bw = new StaticWriter(size);
+  var payloadSize = size - 20;
   var packet, payload;
 
-  bw.seek(4);
+  bw.writeU32(payloadSize);
   bw.writeVarString(cmd, 'ascii');
   bw.writeU32(body.length);
   bw.writeBytes(body);
   bw.seek(16);
 
   packet = bw.render();
-  payload = packet.slice(4, -16);
+  payload = packet.slice(4, 4 + payloadSize);
 
   this.maybeRekey(packet);
 
-  this.output.encryptSize(payload.length).copy(packet, 0);
+  this.output.encryptSize(packet);
   this.output.encrypt(payload);
-  this.output.finish().copy(packet, 4 + payload.length);
+  this.output.finish().copy(packet, 4 + payloadSize);
   this.output.sequence();
 
   return packet;
@@ -679,6 +697,9 @@ BIP151.prototype.parse = function parse(data) {
   if (!this.hasSize) {
     size = this.input.decryptSize(data);
 
+    assert(this.waiting === 4);
+    assert(data.length === 4);
+
     // Allow 3 batched packets of max message size (12mb).
     // Not technically standard, but this protects us
     // from buffering tons of data due to either an
@@ -686,7 +707,6 @@ BIP151.prototype.parse = function parse(data) {
     // Note that 6 is the minimum size:
     // varint-cmdlen(1) str-cmd(1) u32-size(4) payload(0)
     if (size < 6 || size > BIP151.MAX_MESSAGE) {
-      this.waiting = 4;
       this.error('Bad packet size: %d.', util.mb(size));
       return;
     }