From d4958eb05fd02f0a597e18ad5ca56b4a051638c6 Mon Sep 17 00:00:00 2001
From: Patrick Nagurny <patrick@bitpay.com>
Date: Thu, 4 Jun 2015 11:09:24 -0400
Subject: [PATCH] support adding custom p2p messages

---
 lib/messages/builder.js |  8 ++++----
 lib/messages/index.js   |  5 +++++
 test/messages/index.js  | 21 +++++++++++++++++++++
 3 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/lib/messages/builder.js b/lib/messages/builder.js
index 068d672..d17f86b 100644
--- a/lib/messages/builder.js
+++ b/lib/messages/builder.js
@@ -63,10 +63,7 @@ function builder(options) {
     commands: {}
   };
 
-  Object.keys(exported.commandsMap).forEach(function(key) {
-
-    var Command = require('./commands/' + key);
-
+  exported.add = function(key, Command) {
     exported.commands[key] = function(obj) {
       return new Command(obj, options);
     };
@@ -78,7 +75,10 @@ function builder(options) {
       message.setPayload(buffer);
       return message;
     };
+  };
 
+  Object.keys(exported.commandsMap).forEach(function(key) {
+    exported.add(key, require('./commands/' + key));
   });
 
   exported.inventoryCommands.forEach(function(command) {
diff --git a/lib/messages/index.js b/lib/messages/index.js
index 7fe58a4..8ce7f01 100644
--- a/lib/messages/index.js
+++ b/lib/messages/index.js
@@ -103,4 +103,9 @@ Messages.prototype._buildFromBuffer = function(command, payload) {
   return this.builder.commands[command].fromBuffer(payload);
 };
 
+Messages.prototype.add = function(key, name, Command) {
+  this.builder.add(key, Command);
+  this[name] = this.builder.commands[key];
+};
+
 module.exports = Messages;
diff --git a/test/messages/index.js b/test/messages/index.js
index b5145e4..33d314d 100644
--- a/test/messages/index.js
+++ b/test/messages/index.js
@@ -192,4 +192,25 @@ describe('Messages', function() {
 
   });
 
+  describe('#add', function() {
+    it('should add a custom message', function() {
+      var network = bitcore.Networks.defaultNetwork;
+      var messages = new Messages({
+        network: network,
+        Block: bitcore.Block,
+        Transaction: bitcore.Transaction
+      });
+
+      var CustomMessage = function(arg, options) {
+        this.arg = arg;
+      };
+
+      messages.add('custom', 'Custom', CustomMessage);
+      should.exist(messages.Custom);
+
+      var message = messages.Custom('hello');
+      message.arg.should.equal('hello');
+    });
+  });
+
 });