Merge pull request #72 from pnagurny/feature/custom-messages

Support adding custom p2p messages

docs/messages.md

@@ -88,3 +88,15 @@ Same as `getheaders` and `headers`, but the response comes one block at the time
 ### Transaction
 
 Message that contains a transaction.
+
+## Custom Messages
+
+It is possible to extend the default peer to peer messages and add custom ones. First you will need to create a message which resembles the default messages in `lib/messages/commands`.
+
+Then to add the custom message:
+
+```javascript
+messages.add('custom', 'Custom', CustomMessage);
+
+var customMessage = messages.Custom('argument');
+```

lib/messages/builder.js

@@ -63,10 +63,7 @@ function builder(options) {
     commands: {}
   };
 
-  Object.keys(exported.commandsMap).forEach(function(key) {
+  exported.add = function(key, Command) {
-
-    var Command = require('./commands/' + key);
-
     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) {

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;

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');
+    });
+  });
+
 });