From 28b05e3cc154d0895de21b807caa22fdfc5cf80c Mon Sep 17 00:00:00 2001
From: Braydon Fuller <braydon@bitpay.com>
Date: Wed, 1 Apr 2015 11:55:00 -0400
Subject: [PATCH] Added preconditions to ping message

---
 lib/messages/commands/ping.js   |  7 +++++++
 lib/messages/commands/pong.js   |  5 +----
 test/messages/commands/index.js | 22 ++++++++++++++++++++++
 3 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/lib/messages/commands/ping.js b/lib/messages/commands/ping.js
index b87e432..8fa5703 100644
--- a/lib/messages/commands/ping.js
+++ b/lib/messages/commands/ping.js
@@ -4,6 +4,9 @@ var Message = require('../message');
 var inherits = require('util').inherits;
 var bitcore = require('bitcore');
 var utils = require('../utils');
+var $ = bitcore.util.preconditions;
+var _ = bitcore.deps._;
+var BufferUtil = bitcore.util.buffer;
 var BufferReader = bitcore.encoding.BufferReader;
 
 /**
@@ -18,6 +21,10 @@ function PingMessage(arg, options) {
   Message.call(this, arg, options);
   this.command = 'ping';
   this.magicNumber = options.magicNumber;
+  $.checkArgument(
+    _.isUndefined(arg) || (BufferUtil.isBuffer(arg) && arg.length === 8),
+    'First argument is expected to be an 8 byte buffer'
+  );
   this.nonce = arg || utils.getNonce();
 }
 inherits(PingMessage, Message);
diff --git a/lib/messages/commands/pong.js b/lib/messages/commands/pong.js
index f5ff9d6..5b248f5 100644
--- a/lib/messages/commands/pong.js
+++ b/lib/messages/commands/pong.js
@@ -25,10 +25,7 @@ function PongMessage(arg, options) {
     _.isUndefined(arg) || (BufferUtil.isBuffer(arg) && arg.length === 8),
     'First argument is expected to be an 8 byte buffer'
   );
-  this.nonce = arg;
-  if (!this.nonce) {
-    this.nonce = utils.getNonce();
-  }
+  this.nonce = arg || utils.getNonce();
 }
 inherits(PongMessage, Message);
 
diff --git a/test/messages/commands/index.js b/test/messages/commands/index.js
index cbb276f..98edd79 100644
--- a/test/messages/commands/index.js
+++ b/test/messages/commands/index.js
@@ -80,6 +80,28 @@ describe('Command Messages', function() {
 
   });
 
+  describe('Ping', function() {
+
+    it('should error if nonce is not a buffer', function() {
+      (function() {
+        var message = messages.Ping('not a buffer');
+      }).should.throw('First argument is expected to be an 8 byte buffer');
+    });
+
+    it('should error if nonce buffer has invalid length', function() {
+      (function() {
+        var message = messages.Ping(new Buffer(Array(9)));
+      }).should.throw('First argument is expected to be an 8 byte buffer');
+    });
+
+    it('should set a nonce if not included', function() {
+      var message = messages.Ping();
+      should.exist(message.nonce);
+      message.nonce.length.should.equal(8);
+    });
+
+  });
+
   describe('FilterLoad', function() {
 
     it('should return a null payload', function() {