From b4553527084d6a677d6457df41a1ca6807f7d7a4 Mon Sep 17 00:00:00 2001
From: Christopher Jeffrey <chjjeffrey@gmail.com>
Date: Thu, 21 Jul 2016 09:24:09 -0700
Subject: [PATCH] peer/bip151: rewrite bip151. handle bip151 on p2p network.

---
 lib/bcoin/bip151.js          | 363 +++++++++++++++++++----------------
 lib/bcoin/peer.js            |  79 +++++++-
 lib/bcoin/protocol/framer.js |  12 +-
 lib/bcoin/protocol/parser.js |   2 +
 test/bip151-test.js          |  47 +++--
 5 files changed, 318 insertions(+), 185 deletions(-)

diff --git a/lib/bcoin/bip151.js b/lib/bcoin/bip151.js
index 30192b30..bb6b8efc 100644
--- a/lib/bcoin/bip151.js
+++ b/lib/bcoin/bip151.js
@@ -19,9 +19,9 @@ var INFO_KEY1 = new Buffer('BitcoinK1', 'ascii');
 var INFO_KEY2 = new Buffer('BitcoinK2', 'ascii');
 var INFO_SID = new Buffer('BitcoinSessionID', 'ascii');
 
-function BIP151(cipher, key) {
-  if (!(this instanceof BIP151))
-    return new BIP151(cipher, key);
+function BIP151Stream(cipher, key) {
+  if (!(this instanceof BIP151Stream))
+    return new BIP151Stream(cipher, key);
 
   EventEmitter.call(this);
 
@@ -29,25 +29,21 @@ function BIP151(cipher, key) {
   this.privateKey = key || bcoin.ec.generatePrivateKey();
   this.cipher = cipher || 0;
   this.secret = null;
+  this.prk = null;
   this.k1 = null;
   this.k2 = null;
   this.sid = null;
+
+  assert(this.cipher === 0, 'Unknown cipher type.');
+
   this.chacha = new chachapoly.ChaCha20();
   this.aead = new chachapoly.AEAD();
-  this.prk = null;
   this.tag = null;
   this.seq = 0;
-  this.initReceived = false;
-  this.ackReceived = false;
-  this.initSent = false;
-  this.ackSent = false;
+
   this.highWaterMark = 1024 * (1 << 20);
   this.processed = 0;
   this.lastRekey = 0;
-  this.timeout = null;
-  this.callback = null;
-  this.completed = false;
-  this.handshake = false;
 
   this.pendingHeader = [];
   this.pendingHeaderTotal = 0;
@@ -57,9 +53,9 @@ function BIP151(cipher, key) {
   this.waiting = 0;
 }
 
-utils.inherits(BIP151, EventEmitter);
+utils.inherits(BIP151Stream, EventEmitter);
 
-BIP151.prototype.init = function init(publicKey) {
+BIP151Stream.prototype.init = function init(publicKey) {
   var p = bcoin.writer();
 
   this.publicKey = publicKey;
@@ -82,14 +78,17 @@ BIP151.prototype.init = function init(publicKey) {
   this.lastRekey = utils.ms();
 };
 
-BIP151.prototype.isReady = function isReady() {
-  return this.initSent
-    && this.ackReceived
-    && this.initReceived
-    && this.ackSent;
+BIP151Stream.prototype.maybeRekey = function maybeRekey(data) {
+  var self = this;
+  this.processed += data.length;
+  if (this.processed >= this.highWaterMark) {
+    this.processed -= this.highWaterMark;
+    this.rekey();
+    this.emit('rekey');
+  }
 };
 
-BIP151.prototype.rekey = function rekey() {
+BIP151Stream.prototype.rekey = function rekey() {
   assert(this.prk, 'Cannot rekey before initialization.');
 
   this.k1 = utils.hash256(this.k1);
@@ -104,193 +103,54 @@ BIP151.prototype.rekey = function rekey() {
   this.lastRekey = utils.ms();
 };
 
-BIP151.prototype.sequence = function sequence() {
+BIP151Stream.prototype.sequence = function sequence() {
   this.seq++;
   this.chacha.init(this.k1, this.iv());
   this.aead.init(this.k2, this.iv());
   this.aead.aad(this.sid);
 };
 
-BIP151.prototype.iv = function iv() {
+BIP151Stream.prototype.iv = function iv() {
   var p = bcoin.writer();
   p.writeU64(this.seq);
   p.writeU32(0);
   return p.render();
 };
 
-BIP151.prototype.getPublicKey = function getPublicKey() {
+BIP151Stream.prototype.getPublicKey = function getPublicKey() {
   return bcoin.ec.publicKeyCreate(this.privateKey, true);
 };
 
-BIP151.prototype.encryptSize = function encryptSize(size) {
+BIP151Stream.prototype.encryptSize = function encryptSize(size) {
   var data = new Buffer(4);
   data.writeUInt32LE(size, 0, true);
   return this.chacha.encrypt(data);
 };
 
-BIP151.prototype.decryptSize = function decryptSize(data) {
+BIP151Stream.prototype.decryptSize = function decryptSize(data) {
   data = data.slice(0, 4);
   this.chacha.encrypt(data);
   return data.readUInt32LE(0, true);
 };
 
-BIP151.prototype.encrypt = function encrypt(data) {
+BIP151Stream.prototype.encrypt = function encrypt(data) {
   return this.aead.encrypt(data);
 };
 
-BIP151.prototype.decrypt = function decrypt(data) {
+BIP151Stream.prototype.decrypt = function decrypt(data) {
   return this.aead.decrypt(data);
 };
 
-BIP151.prototype.finish = function finish(data) {
+BIP151Stream.prototype.finish = function finish(data) {
   this.tag = this.aead.finish(data);
   return this.tag;
 };
 
-BIP151.prototype.verify = function verify(tag) {
+BIP151Stream.prototype.verify = function verify(tag) {
   return chachapoly.Poly1305.verify(this.tag, tag);
 };
 
-BIP151.prototype.toEncinit = function toEncinit(writer) {
-  var p = bcoin.writer(writer);
-
-  p.writeBytes(this.getPublicKey());
-  p.writeU8(this.cipher);
-
-  if (!writer)
-    p = p.render();
-
-  this.initSent = true;
-
-  return p;
-};
-
-BIP151.prototype.encinit = function encinit(data) {
-  var p = bcoin.reader(data);
-  var publicKey = p.readBytes(33);
-
-  // this.cipher = p.readU8();
-  assert(p.readU8() === this.cipher, 'Wrong cipher type.');
-
-  assert(!this.initReceived, 'Already initialized.');
-
-  if (!this.ackReceived) {
-    this.init(publicKey);
-  } else {
-    assert(utils.equal(publicKey, this.publicKey),
-      'Bad pubkey.');
-  }
-
-  this.initReceived = true;
-
-  return this;
-};
-
-BIP151.fromEncinit = function fromEncinit(data) {
-  return new BIP151().encinit(data);
-};
-
-BIP151.prototype.toEncack = function toEncack(writer) {
-  var p = bcoin.writer(writer);
-
-  p.writeBytes(this.getPublicKey());
-
-  if (!writer)
-    p = p.render();
-
-  if (!this.ackSent) {
-    this.ackSent = true;
-    if (this.isReady()) {
-      this.handshake = true;
-      this.emit('handshake');
-    }
-  }
-
-  return p;
-};
-
-BIP151.prototype.toRekey = function toRekey(writer) {
-  var p = bcoin.writer(writer);
-
-  p.writeBytes(constants.ZERO_KEY);
-
-  if (!writer)
-    p = p.render();
-
-  return p;
-};
-
-BIP151.prototype.maybeRekey = function maybeRekey(data) {
-  var self = this;
-  this.processed += data.length;
-  if (this.processed >= this.highWaterMark) {
-    this.processed -= this.highWaterMark;
-    utils.nextTick(function() {
-      self.emit('rekey');
-      self.rekey();
-    });
-  }
-};
-
-BIP151.prototype.complete = function complete(err) {
-  assert(!this.completed, 'Already completed.');
-  assert(this.callback, 'No completion callback.');
-
-  this.completed = true;
-
-  clearTimeout(this.timeout);
-  this.timeout = null;
-
-  this.callback(err);
-  this.callback = null;
-};
-
-BIP151.prototype.wait = function wait(timeout, callback) {
-  var self = this;
-
-  assert(!this.handshake, 'Cannot wait for init after handshake.');
-
-  this.callback = callback;
-
-  this.timeout = setTimeout(function() {
-    self.complete(new Error('Timed out.'));
-  }, timeout);
-
-  this.once('handshake', function() {
-    self.complete();
-  });
-};
-
-BIP151.prototype.encack = function encack(data) {
-  var p = bcoin.reader(data);
-  var publicKey = p.readBytes(33);
-
-  assert(this.initSent, 'Unsolicited ACK.');
-
-  if (utils.equal(publicKey, constants.ZERO_KEY)) {
-    assert(this.ackReceived, 'No ACK before rekey.');
-    assert(this.handshake, 'No initialization before rekey.');
-    this.rekey();
-    return;
-  }
-
-  assert(!this.ackReceived, 'Already ACKed.');
-  this.ackReceived = true;
-
-  if (!this.initReceived) {
-    this.init(publicKey);
-  } else {
-    assert(utils.equal(publicKey, this.publicKey),
-      'Bad pubkey.');
-  }
-
-  if (this.isReady()) {
-    this.handshake = true;
-    this.emit('handshake');
-  }
-};
-
-BIP151.prototype.feed = function feed(data) {
+BIP151Stream.prototype.feed = function feed(data) {
   var chunk, payload, tag, p, cmd, body;
 
   this.maybeRekey(data);
@@ -371,7 +231,7 @@ BIP151.prototype.feed = function feed(data) {
 };
 
 // TODO: We could batch packets here!
-BIP151.prototype.packet = function packet(cmd, body) {
+BIP151Stream.prototype.packet = function packet(cmd, body) {
   var p = bcoin.writer();
   var payload, packet;
 
@@ -393,4 +253,167 @@ BIP151.prototype.packet = function packet(cmd, body) {
   return packet;
 };
 
+function BIP151(cipher) {
+  if (!(this instanceof BIP151))
+    return new BIP151(cipher);
+
+  EventEmitter.call(this);
+
+  this.input = new BIP151Stream(cipher);
+  this.output = null;
+
+  this.initReceived = false;
+  this.ackReceived = false;
+  this.initSent = false;
+  this.ackSent = false;
+  this.timeout = null;
+  this.callback = null;
+  this.completed = false;
+  this.handshake = false;
+
+  this._init();
+}
+
+utils.inherits(BIP151, EventEmitter);
+
+BIP151.prototype._init = function _init() {
+  var self = this;
+
+  this.input.on('rekey', function() {
+    self.emit('rekey');
+  });
+
+  this.input.on('packet', function(cmd, body) {
+    self.emit('packet', cmd, body);
+  });
+};
+
+BIP151.prototype.isReady = function isReady() {
+  return this.initSent
+    && this.ackReceived
+    && this.initReceived
+    && this.ackSent;
+};
+
+BIP151.prototype.toEncinit = function toEncinit(writer) {
+  var p = bcoin.writer(writer);
+
+  p.writeBytes(this.input.getPublicKey());
+  p.writeU8(this.input.cipher);
+
+  if (!writer)
+    p = p.render();
+
+  this.initSent = true;
+
+  return p;
+};
+
+BIP151.prototype.encack = function encack(data) {
+  var p = bcoin.reader(data);
+  var publicKey = p.readBytes(33);
+
+  assert(this.initSent, 'Unsolicited ACK.');
+
+  if (utils.equal(publicKey, constants.ZERO_KEY)) {
+    assert(this.handshake, 'No initialization before rekey.');
+    this.output.rekey();
+    return;
+  }
+
+  assert(!this.ackReceived, 'Already ACKed.');
+  this.ackReceived = true;
+
+  this.input.init(publicKey);
+
+  if (this.isReady()) {
+    this.handshake = true;
+    this.emit('handshake');
+  }
+};
+
+BIP151.prototype.encinit = function encinit(data) {
+  var p = bcoin.reader(data);
+  var publicKey = p.readBytes(33);
+
+  assert(!this.initReceived, 'Already initialized.');
+
+  this.output = new BIP151Stream(p.readU8());
+  this.output.init(publicKey);
+
+  this.initReceived = true;
+
+  return this;
+};
+
+BIP151.prototype.toEncack = function toEncack(writer) {
+  var p = bcoin.writer(writer);
+
+  assert(this.output, 'Cannot ack before init.');
+
+  p.writeBytes(this.output.getPublicKey());
+
+  if (!writer)
+    p = p.render();
+
+  if (!this.ackSent) {
+    this.ackSent = true;
+    if (this.isReady()) {
+      this.handshake = true;
+      this.emit('handshake');
+    }
+  }
+
+  return p;
+};
+
+BIP151.prototype.toRekey = function toRekey(writer) {
+  var p = bcoin.writer(writer);
+
+  p.writeBytes(constants.ZERO_KEY);
+
+  if (!writer)
+    p = p.render();
+
+  return p;
+};
+
+BIP151.prototype.complete = function complete(err) {
+  assert(!this.completed, 'Already completed.');
+  assert(this.callback, 'No completion callback.');
+
+  this.completed = true;
+
+  clearTimeout(this.timeout);
+  this.timeout = null;
+
+  this.callback(err);
+  this.callback = null;
+};
+
+BIP151.prototype.wait = function wait(timeout, callback) {
+  var self = this;
+
+  assert(!this.handshake, 'Cannot wait for init after handshake.');
+
+  this.callback = callback;
+
+  this.timeout = setTimeout(function() {
+    self.complete(new Error('BIP151 handshake timed out.'));
+  }, timeout);
+
+  this.once('handshake', function() {
+    self.complete();
+  });
+};
+
+
+BIP151.prototype.feed = function feed(data) {
+  return this.input.feed(data);
+};
+
+BIP151.prototype.packet = function packet(cmd, body) {
+  return this.output.packet(cmd, body);
+};
+
 module.exports = BIP151;
diff --git a/lib/bcoin/peer.js b/lib/bcoin/peer.js
index 67174798..ac69539c 100644
--- a/lib/bcoin/peer.js
+++ b/lib/bcoin/peer.js
@@ -18,6 +18,7 @@ var VersionPacket = bcoin.packets.VersionPacket;
 var GetBlocksPacket = bcoin.packets.GetBlocksPacket;
 var RejectPacket = bcoin.packets.RejectPacket;
 var NetworkAddress = bcoin.packets.NetworkAddress;
+var Packet = bcoin.protocol.parser.Packet;
 
 /**
  * Represents a remote peer.
@@ -109,6 +110,10 @@ function Peer(pool, options) {
   this.compactMode = null;
   this.compactBlocks = {};
   this.sentAddr = false;
+  this.bip151 = null;
+
+  if (options.bip151)
+    this.bip151 = new bcoin.bip151();
 
   this.challenge = null;
   this.lastPong = -1;
@@ -215,7 +220,10 @@ Peer.prototype._init = function init() {
   });
 
   this.socket.on('data', function(chunk) {
-    self.parser.feed(chunk);
+    if (self.bip151 && self.bip151.handshake)
+      self.bip151.feed(chunk);
+    else
+      self.parser.feed(chunk);
   });
 
   this.parser.on('packet', function(packet) {
@@ -227,6 +235,16 @@ Peer.prototype._init = function init() {
     self.reject(null, 'malformed', 'error parsing message', 10);
   });
 
+  if (this.bip151) {
+    this.bip151.on('error', function(err) {
+      self.reject(null, 'malformed', 'error parsing message', 10);
+      self._error(err, true);
+    });
+    this.bip151.on('rekey', function() {
+      self.write(self.framer.encack(self.bip151.toRekey()));
+    });
+  }
+
   if (this.connected) {
     utils.nextTick(function() {
       self._onConnect();
@@ -253,6 +271,36 @@ Peer.prototype._onConnect = function _onConnect() {
     this._connectTimeout = null;
   }
 
+  // Send encinit. Wait for handshake to complete.
+  if (this.bip151 && !this.bip151.completed) {
+    this.logger.info('Attempting BIP151 handshake (%s).', this.hostname);
+    this.write(this.framer.encinit(this.bip151.toEncinit()));
+    this.bip151.wait(5000, function(err) {
+      if (err)
+        self.logger.error(err);
+      assert(self.bip151.completed);
+      self._onConnect();
+    });
+    return;
+  }
+
+  if (this.bip151 && this.bip151.handshake) {
+    this.logger.info('BIP151 handshake complete (%s).', this.hostname);
+    this.logger.info('Connection is encrypted (%s).', this.hostname);
+    this.bip151.on('packet', function(cmd, body) {
+      var packet = new Packet(cmd, body.length);
+      try {
+        packet.payload = self.parser.parsePayload(cmd, body);
+      } catch (e) {
+        return self.parser._error(e);
+      }
+      self.parser.emit('packet', packet);
+    });
+    this.framer.packet = function packet(cmd, body) {
+      return self.bip151.packet(cmd, body);
+    };
+  }
+
   this.request('verack', function(err) {
     self._onAck(err);
   });
@@ -766,6 +814,13 @@ Peer.prototype._onPacket = function onPacket(packet) {
   var cmd = packet.cmd;
   var payload = packet.payload;
 
+  if (this.bip151
+      && !this.bip151.completed
+      && cmd !== 'encinit'
+      && cmd !== 'encack') {
+    this.bip151.complete(new Error('Message before handshake.'));
+  }
+
   if (this.lastBlock && cmd !== 'tx')
     this._flushMerkle();
 
@@ -1819,6 +1874,18 @@ Peer.prototype._handleAlert = function _handleAlert(alert) {
  */
 
 Peer.prototype._handleEncinit = function _handleEncinit(payload) {
+  if (!this.bip151)
+    return;
+
+  try {
+    this.bip151.encinit(payload);
+  } catch (e) {
+    this._error(e);
+    return;
+  }
+
+  this.write(this.framer.encack(this.bip151.toEncack()));
+
   this.fire('encinit', payload);
 };
 
@@ -1829,6 +1896,16 @@ Peer.prototype._handleEncinit = function _handleEncinit(payload) {
  */
 
 Peer.prototype._handleEncack = function _handleEncack(payload) {
+  if (!this.bip151)
+    return;
+
+  try {
+    this.bip151.encack(payload);
+  } catch (e) {
+    this._error(e);
+    return;
+  }
+
   this.fire('encack', payload);
 };
 
diff --git a/lib/bcoin/protocol/framer.js b/lib/bcoin/protocol/framer.js
index 2e7dc7ef..c8dccd79 100644
--- a/lib/bcoin/protocol/framer.js
+++ b/lib/bcoin/protocol/framer.js
@@ -936,7 +936,11 @@ Framer.feeFilter = function feeFilter(data, writer) {
  */
 
 Framer.encinit = function encinit(data, writer) {
-  return data.toEncinit(writer);
+  if (writer) {
+    writer.writeBytes(data);
+    return writer;
+  }
+  return data;
 };
 
 /**
@@ -947,7 +951,11 @@ Framer.encinit = function encinit(data, writer) {
  */
 
 Framer.encack = function encack(data, writer) {
-  return data.toEncack(writer);
+  if (writer) {
+    writer.writeBytes(data);
+    return writer;
+  }
+  return data;
 };
 
 /**
diff --git a/lib/bcoin/protocol/parser.js b/lib/bcoin/protocol/parser.js
index 4ebf74eb..af95778f 100644
--- a/lib/bcoin/protocol/parser.js
+++ b/lib/bcoin/protocol/parser.js
@@ -705,6 +705,8 @@ function Packet(cmd, size, checksum) {
   this.payload = null;
 }
 
+Parser.Packet = Packet;
+
 /*
  * Expose
  */
diff --git a/test/bip151-test.js b/test/bip151-test.js
index 83c220b9..4eae0392 100644
--- a/test/bip151-test.js
+++ b/test/bip151-test.js
@@ -10,7 +10,9 @@ var assert = require('assert');
 describe('BIP151', function() {
   var client = new bcoin.bip151();
   var server = new bcoin.bip151();
-  var payload = new Buffer('deadbeef', 'hex');
+  function payload() {
+    return new Buffer('deadbeef', 'hex');
+  }
 
   it('should do encinit', function() {
     client.encinit(server.toEncinit());
@@ -34,7 +36,7 @@ describe('BIP151', function() {
   });
 
   it('should encrypt payload from client to server', function() {
-    var packet = client.packet('fake', payload);
+    var packet = client.packet('fake', payload());
     var emitted = false;
     server.once('packet', function(cmd, body) {
       emitted = true;
@@ -46,7 +48,7 @@ describe('BIP151', function() {
   });
 
   it('should encrypt payload from server to client', function() {
-    var packet = server.packet('fake', payload);
+    var packet = server.packet('fake', payload());
     var emitted = false;
     client.once('packet', function(cmd, body) {
       emitted = true;
@@ -58,7 +60,7 @@ describe('BIP151', function() {
   });
 
   it('should encrypt payload from client to server (2)', function() {
-    var packet = client.packet('fake', payload);
+    var packet = client.packet('fake', payload());
     var emitted = false;
     server.once('packet', function(cmd, body) {
       emitted = true;
@@ -70,7 +72,7 @@ describe('BIP151', function() {
   });
 
   it('should encrypt payload from server to client (2)', function() {
-    var packet = server.packet('fake', payload);
+    var packet = server.packet('fake', payload());
     var emitted = false;
     client.once('packet', function(cmd, body) {
       emitted = true;
@@ -83,7 +85,7 @@ describe('BIP151', function() {
 
   it('client should rekey', function() {
     var rekeyed = false;
-    var bytes = client.processed;
+    var bytes = client.input.processed;
 
     client.once('rekey', function() {
       rekeyed = true;
@@ -99,19 +101,19 @@ describe('BIP151', function() {
     });
 
     // Force a rekey after 1gb processed.
-    client.maybeRekey({ length: 1024 * (1 << 20) });
+    client.input.maybeRekey({ length: 1024 * (1 << 20) });
 
     utils.nextTick(function() {
       assert(rekeyed);
 
       // Reset so as not to mess up
       // the symmetry of client and server.
-      client.processed = bytes + 33 + 31;
+      client.input.processed = bytes + 33 + 31;
     });
   });
 
   it('should encrypt payload from client to server after rekey', function() {
-    var packet = client.packet('fake', payload);
+    var packet = client.packet('fake', payload());
     var emitted = false;
     server.once('packet', function(cmd, body) {
       emitted = true;
@@ -123,7 +125,7 @@ describe('BIP151', function() {
   });
 
   it('should encrypt payload from server to client after rekey', function() {
-    var packet = server.packet('fake', payload);
+    var packet = server.packet('fake', payload());
     var emitted = false;
     client.once('packet', function(cmd, body) {
       emitted = true;
@@ -135,7 +137,7 @@ describe('BIP151', function() {
   });
 
   it('should encrypt payload from client to server after rekey (2)', function() {
-    var packet = client.packet('fake', payload);
+    var packet = client.packet('fake', payload());
     var emitted = false;
     server.once('packet', function(cmd, body) {
       emitted = true;
@@ -147,7 +149,7 @@ describe('BIP151', function() {
   });
 
   it('should encrypt payload from server to client after rekey (2)', function() {
-    var packet = server.packet('fake', payload);
+    var packet = server.packet('fake', payload());
     var emitted = false;
     client.once('packet', function(cmd, body) {
       emitted = true;
@@ -157,4 +159,25 @@ describe('BIP151', function() {
     client.feed(packet);
     assert(emitted);
   });
+
+  it('should encrypt payloads both ways asynchronously', function() {
+    var spacket = server.packet('fake', payload());
+    var cpacket = client.packet('fake', payload());
+    var cemitted = false;
+    var semitted = false;
+    client.once('packet', function(cmd, body) {
+      cemitted = true;
+      assert.equal(cmd, 'fake');
+      assert.equal(body.toString('hex'), 'deadbeef');
+    });
+    server.once('packet', function(cmd, body) {
+      semitted = true;
+      assert.equal(cmd, 'fake');
+      assert.equal(body.toString('hex'), 'deadbeef');
+    });
+    client.feed(spacket);
+    server.feed(cpacket);
+    assert(cemitted);
+    assert(semitted);
+  });
 });