bip151: improve rekey test.

lib/bcoin/bip151.js

@@ -78,6 +78,8 @@ BIP151.prototype.init = function init(publicKey) {
   this.chacha.init(this.k1, this.iv());
   this.aead.init(this.k2, this.iv());
   this.aead.aad(this.sid);
+
+  this.lastRekey = utils.ms();
 };
 
 BIP151.prototype.isReady = function isReady() {
@@ -89,12 +91,17 @@ BIP151.prototype.isReady = function isReady() {
 
 BIP151.prototype.rekey = function rekey() {
   assert(this.prk, 'Cannot rekey before initialization.');
+
   this.k1 = utils.hash256(this.k1);
   this.k2 = utils.hash256(this.k2);
+
   this.seq = 0;
+
   this.chacha.init(this.k1, this.iv());
   this.aead.init(this.k2, this.iv());
   this.aead.aad(this.sid);
+
+  this.lastRekey = utils.ms();
 };
 
 BIP151.prototype.sequence = function sequence() {
@@ -174,7 +181,6 @@ BIP151.prototype.encinit = function encinit(data) {
       'Bad pubkey.');
   }
 
-  this.lastRekey = utils.ms();
   this.initReceived = true;
 
   return this;
@@ -219,10 +225,9 @@ BIP151.prototype.maybeRekey = function maybeRekey(data) {
   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');
+      self.rekey();
     });
   }
 };
@@ -249,7 +254,7 @@ BIP151.prototype.wait = function wait(timeout, callback) {
 
   this.timeout = setTimeout(function() {
     self.complete(new Error('Timed out.'));
-  }, 1000);
+  }, timeout);
 
   this.once('handshake', function() {
     self.complete();
@@ -357,7 +362,7 @@ BIP151.prototype.feed = function feed(data) {
         body = p.readBytes(p.readU32());
       } catch (e) {
         this.emit('error', e);
-        continue;
+        break;
       }
 
       this.emit('packet', cmd, body);
@@ -365,6 +370,7 @@ BIP151.prototype.feed = function feed(data) {
   }
 };
 
+// TODO: We could batch packets here!
 BIP151.prototype.frame = function frame(cmd, body) {
   var p = bcoin.writer();
   var payload, packet;

test/bip151-test.js

@@ -82,8 +82,32 @@ describe('BIP151', function() {
   });
 
   it('client should rekey', function() {
-    client.rekey();
-    server.encack(client.toRekey());
+    var rekeyed = false;
+    var bytes = client.processed;
+
+    client.once('rekey', function() {
+      rekeyed = true;
+      var packet = client.frame('encack', client.toRekey());
+      var emitted = false;
+      server.once('packet', function(cmd, body) {
+        emitted = true;
+        assert.equal(cmd, 'encack');
+        server.encack(body);
+      });
+      server.feed(packet);
+      assert(emitted);
+    });
+
+    // Force a rekey after 1gb processed.
+    client.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;
+    });
   });
 
   it('should encrypt payload from client to server after rekey', function() {