From e7347dd6208b21709f06e3d6e1154648b824b149 Mon Sep 17 00:00:00 2001
From: Christopher Jeffrey <chjjeffrey@gmail.com>
Date: Thu, 21 Jul 2016 09:21:36 -0700
Subject: [PATCH] bip151: fixes and tests.

---
 lib/bcoin/bip151.js | 144 +++++++++++++++++++++++++++++++++++++++++---
 test/bip151-test.js | 136 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 271 insertions(+), 9 deletions(-)
 create mode 100644 test/bip151-test.js

diff --git a/lib/bcoin/bip151.js b/lib/bcoin/bip151.js
index 220124c9..f48ebf71 100644
--- a/lib/bcoin/bip151.js
+++ b/lib/bcoin/bip151.js
@@ -37,6 +37,17 @@ function BIP151(cipher, key) {
   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;
@@ -69,6 +80,13 @@ BIP151.prototype.init = function init(publicKey) {
   this.aead.aad(this.sid);
 };
 
+BIP151.prototype.isReady = function isReady() {
+  return this.initSent
+    && this.ackReceived
+    && this.initReceived
+    && this.ackSent;
+};
+
 BIP151.prototype.rekey = function rekey() {
   assert(this.prk, 'Cannot rekey before initialization.');
   this.k1 = utils.hash256(this.k1);
@@ -135,19 +153,35 @@ BIP151.prototype.toEncinit = function toEncinit(writer) {
   if (!writer)
     p = p.render();
 
+  this.initSent = true;
+
   return p;
 };
 
-BIP151.prototype.fromEncinit = function fromEncinit(data) {
+BIP151.prototype.encinit = function encinit(data) {
   var p = bcoin.reader(data);
   var publicKey = p.readBytes(33);
-  this.cipher = p.readU8();
-  this.init(publicKey);
+
+  // 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.lastRekey = utils.ms();
+  this.initReceived = true;
+
   return this;
 };
 
 BIP151.fromEncinit = function fromEncinit(data) {
-  return new BIP151().fromEncinit(data);
+  return new BIP151().encinit(data);
 };
 
 BIP151.prototype.toEncack = function toEncack(writer) {
@@ -158,24 +192,104 @@ BIP151.prototype.toEncack = function toEncack(writer) {
   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;
+    this.lastRekey = utils.ms();
+    utils.nextTick(function() {
+      self.rekey();
+      self.emit('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.'));
+  }, 1000);
+
+  this.once('handshake', function() {
+    self.complete();
+  });
+};
+
 BIP151.prototype.encack = function encack(data) {
   var p = bcoin.reader(data);
   var publicKey = p.readBytes(33);
 
-  if (utils.isZero(publicKey)) {
+  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;
   }
 
-  this.init(publicKey);
+  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) {
   var chunk, payload, tag, p, cmd, body;
 
+  this.maybeRekey(data);
+
   while (data) {
     if (!this.hasHeader) {
       this.pendingHeaderTotal += data.length;
@@ -236,10 +350,18 @@ BIP151.prototype.feed = function feed(data) {
     }
 
     p = bcoin.reader(payload, true);
-    cmd = p.readVarString('ascii');
-    body = p.readBytes(p.readU32());
 
-    this.emit('packet', cmd, body);
+    while (p.left()) {
+      try {
+        cmd = p.readVarString('ascii');
+        body = p.readBytes(p.readU32());
+      } catch (e) {
+        this.emit('error', e);
+        continue;
+      }
+
+      this.emit('packet', cmd, body);
+    }
   }
 };
 
@@ -260,5 +382,9 @@ BIP151.prototype.frame = function frame(cmd, body) {
   this.finish().copy(packet, 4 + payload.length);
   this.sequence();
 
+  this.maybeRekey(payload);
+
   return packet;
 };
+
+module.exports = BIP151;
diff --git a/test/bip151-test.js b/test/bip151-test.js
new file mode 100644
index 00000000..1734558e
--- /dev/null
+++ b/test/bip151-test.js
@@ -0,0 +1,136 @@
+'use strict';
+
+var bn = require('bn.js');
+var bcoin = require('../').set('main');
+var utils = bcoin.utils;
+var constants = bcoin.protocol.constants;
+var network = bcoin.protocol.network;
+var assert = require('assert');
+
+describe('BIP151', function() {
+  var client = new bcoin.bip151();
+  var server = new bcoin.bip151();
+  var payload = new Buffer('deadbeef', 'hex');
+
+  it('should do encinit', function() {
+    client.encinit(server.toEncinit());
+    server.encinit(client.toEncinit());
+    assert(!client.handshake);
+    assert(!server.handshake);
+  });
+
+  it('should do encack', function() {
+    client.encack(server.toEncack());
+    server.encack(client.toEncack());
+    assert(client.handshake);
+    assert(server.handshake);
+  });
+
+  it('should have completed ECDH handshake', function() {
+    assert(client.isReady());
+    assert(server.isReady());
+    assert(client.handshake);
+    assert(server.handshake);
+  });
+
+  it('should encrypt payload from client to server', function() {
+    var packet = client.frame('fake', payload);
+    var emitted = false;
+    server.once('packet', function(cmd, body) {
+      emitted = true;
+      assert.equal(cmd, 'fake');
+      assert.equal(body.toString('hex'), 'deadbeef');
+    });
+    server.feed(packet);
+    assert(emitted);
+  });
+
+  it('should encrypt payload from server to client', function() {
+    var packet = server.frame('fake', payload);
+    var emitted = false;
+    client.once('packet', function(cmd, body) {
+      emitted = true;
+      assert.equal(cmd, 'fake');
+      assert.equal(body.toString('hex'), 'deadbeef');
+    });
+    client.feed(packet);
+    assert(emitted);
+  });
+
+  it('should encrypt payload from client to server (2)', function() {
+    var packet = client.frame('fake', payload);
+    var emitted = false;
+    server.once('packet', function(cmd, body) {
+      emitted = true;
+      assert.equal(cmd, 'fake');
+      assert.equal(body.toString('hex'), 'deadbeef');
+    });
+    server.feed(packet);
+    assert(emitted);
+  });
+
+  it('should encrypt payload from server to client (2)', function() {
+    var packet = server.frame('fake', payload);
+    var emitted = false;
+    client.once('packet', function(cmd, body) {
+      emitted = true;
+      assert.equal(cmd, 'fake');
+      assert.equal(body.toString('hex'), 'deadbeef');
+    });
+    client.feed(packet);
+    assert(emitted);
+  });
+
+  it('client should rekey', function() {
+    client.rekey();
+    server.encack(client.toRekey());
+  });
+
+  it('should encrypt payload from client to server after rekey', function() {
+    var packet = client.frame('fake', payload);
+    var emitted = false;
+    server.once('packet', function(cmd, body) {
+      emitted = true;
+      assert.equal(cmd, 'fake');
+      assert.equal(body.toString('hex'), 'deadbeef');
+    });
+    server.feed(packet);
+    assert(emitted);
+  });
+
+  it('should encrypt payload from server to client after rekey', function() {
+    var packet = server.frame('fake', payload);
+    var emitted = false;
+    client.once('packet', function(cmd, body) {
+      emitted = true;
+      assert.equal(cmd, 'fake');
+      assert.equal(body.toString('hex'), 'deadbeef');
+    });
+    client.feed(packet);
+    assert(emitted);
+  });
+
+  it('should encrypt payload from client to server after rekey (2)', function() {
+    var packet = client.frame('fake', payload);
+    var emitted = false;
+    server.once('packet', function(cmd, body) {
+      emitted = true;
+      assert.equal(cmd, 'fake');
+      assert.equal(body.toString('hex'), 'deadbeef');
+    });
+    server.feed(packet);
+    assert(emitted);
+  });
+
+  it('should encrypt payload from server to client after rekey (2)', function() {
+    var packet = server.frame('fake', payload);
+    var emitted = false;
+    client.once('packet', function(cmd, body) {
+      emitted = true;
+      assert.equal(cmd, 'fake');
+      assert.equal(body.toString('hex'), 'deadbeef');
+    });
+    client.feed(packet);
+    assert(emitted);
+  });
+});