From 7b9eb8b4f26e771beaa3f2b3b3f64f8c91801319 Mon Sep 17 00:00:00 2001
From: Christopher Jeffrey <chjjeffrey@gmail.com>
Date: Tue, 26 Jul 2016 12:23:40 -0700
Subject: [PATCH] peer: make bip151 handling less hacky.

---
 lib/bcoin/peer.js            | 54 +++++++++++++++++-------------------
 lib/bcoin/protocol/framer.js |  5 ++++
 lib/bcoin/protocol/parser.js | 34 +++++++++++++++++++++--
 3 files changed, 62 insertions(+), 31 deletions(-)

diff --git a/lib/bcoin/peer.js b/lib/bcoin/peer.js
index 7cc4d4b4..379ba46c 100644
--- a/lib/bcoin/peer.js
+++ b/lib/bcoin/peer.js
@@ -18,7 +18,6 @@ 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.
@@ -87,8 +86,6 @@ function Peer(pool, options) {
   this.chain = this.pool.chain;
   this.mempool = this.pool.mempool;
   this.network = this.chain.network;
-  this.parser = new bcoin.protocol.parser({ network: this.network });
-  this.framer = new bcoin.protocol.framer({ network: this.network });
   this.locker = new bcoin.locker(this);
   this.version = null;
   this.destroyed = false;
@@ -115,6 +112,9 @@ function Peer(pool, options) {
   if (this.pool.options.bip151)
     this.bip151 = new bcoin.bip151();
 
+  this.parser = new bcoin.protocol.parser(this);
+  this.framer = new bcoin.protocol.framer(this);
+
   this.challenge = null;
   this.lastPong = -1;
   this.lastPing = -1;
@@ -220,10 +220,7 @@ Peer.prototype._init = function init() {
   });
 
   this.socket.on('data', function(chunk) {
-    if (self.bip151 && self.bip151.handshake)
-      self.bip151.feed(chunk);
-    else
-      self.parser.feed(chunk);
+    self.parser.feed(chunk);
   });
 
   this.parser.on('packet', function(packet) {
@@ -273,33 +270,34 @@ Peer.prototype._onConnect = function _onConnect() {
   }
 
   // Send encinit. Wait for handshake to complete.
-  if (this.bip151 && !this.bip151.completed) {
+  if (this.bip151) {
+    assert(!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) {
+    return this.bip151.wait(5000, function(err) {
       if (err)
-        self.logger.error(err);
-      assert(self.bip151.completed);
-      self._onConnect();
+        self._error(err, true);
+      self._onHandshake();
     });
-    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._onHandshake();
+};
+
+/**
+ * Handle post bip151-handshake.
+ * @private
+ */
+
+Peer.prototype._onHandshake = function _onHandshake() {
+  var self = this;
+
+  if (this.bip151) {
+    assert(this.bip151.completed);
+    if (this.bip151.handshake) {
+      this.logger.info('BIP151 handshake complete (%s).', this.hostname);
+      this.logger.info('Connection is encrypted (%s).', this.hostname);
+    }
   }
 
   this.request('verack', function(err) {
diff --git a/lib/bcoin/protocol/framer.js b/lib/bcoin/protocol/framer.js
index c8dccd79..42cf5422 100644
--- a/lib/bcoin/protocol/framer.js
+++ b/lib/bcoin/protocol/framer.js
@@ -28,7 +28,9 @@ function Framer(options) {
     options = {};
 
   this.options = options;
+
   this.network = bcoin.network.get(options.network);
+  this.bip151 = options.bip151;
 }
 
 /**
@@ -79,6 +81,9 @@ Framer.prototype.packet = function packet(cmd, payload, checksum) {
 
   assert(payload, 'No payload.');
 
+  if (this.bip151 && this.bip151.handshake)
+    return this.bip151.packet(cmd, payload);
+
   header = this.header(cmd, payload, checksum);
 
   return Buffer.concat([header, payload]);
diff --git a/lib/bcoin/protocol/parser.js b/lib/bcoin/protocol/parser.js
index af95778f..d7f335e4 100644
--- a/lib/bcoin/protocol/parser.js
+++ b/lib/bcoin/protocol/parser.js
@@ -32,15 +32,42 @@ function Parser(options) {
 
   EventEmitter.call(this);
 
+  this.network = bcoin.network.get(options.network);
+  this.bip151 = options.bip151;
+
   this.pending = [];
   this.pendingTotal = 0;
   this.waiting = 24;
   this.packet = null;
-  this.network = bcoin.network.get(options.network);
+
+  this._init();
 }
 
 utils.inherits(Parser, EventEmitter);
 
+/**
+ * Initialize. Bind to events.
+ * @private
+ * @param {String} str
+ */
+
+Parser.prototype._init = function _init(str) {
+  var self = this;
+
+  if (!this.bip151)
+    return;
+
+  this.bip151.on('packet', function(cmd, body) {
+    var packet = new Packet(cmd, body.length);
+    try {
+      packet.payload = self.parsePayload(cmd, body);
+    } catch (e) {
+      return self._error(e);
+    }
+    self.emit('packet', packet);
+  });
+};
+
 /**
  * Emit an error.
  * @private
@@ -59,6 +86,9 @@ Parser.prototype._error = function _error(str) {
 Parser.prototype.feed = function feed(data) {
   var chunk, off, len;
 
+  if (this.bip151 && this.bip151.handshake)
+    return this.bip151.feed(data);
+
   this.pendingTotal += data.length;
   this.pending.push(data);
 
@@ -705,8 +735,6 @@ function Packet(cmd, size, checksum) {
   this.payload = null;
 }
 
-Parser.Packet = Packet;
-
 /*
  * Expose
  */