From f2ee5aa89d2b57042f4af8ea67b33ec1b9a2c06e Mon Sep 17 00:00:00 2001
From: koichik <koichik@improvement.jp>
Date: Fri, 22 Jul 2011 02:53:11 +0900
Subject: [PATCH] Add an optional maxLength argument to Buffer.write

Fixes #243.
Fixes #1361.
---
 lib/buffer.js              | 69 +++++++++++++++++++++++++-------------
 test/simple/test-buffer.js | 47 ++++++++++++++++++++++++++
 2 files changed, 92 insertions(+), 24 deletions(-)

diff --git a/lib/buffer.js b/lib/buffer.js
index b3b0b5446f3..fe367ec3804 100644
--- a/lib/buffer.js
+++ b/lib/buffer.js
@@ -92,15 +92,17 @@ SlowBuffer.prototype.toString = function(encoding, start, end) {
 };
 
 
-SlowBuffer.prototype.hexWrite = function(string, offset) {
+SlowBuffer.prototype.hexWrite = function(string, offset, maxLength) {
   var len = string.length;
   offset = +offset || 0;
+  maxLength = Math.min(+maxLength || Number.MAX_VALUE, this.length - offset);
 
   // must be an even number of digits
   if (len % 2) {
     throw new Error('Invalid hex string');
   }
-  for (var i = 0; i < len / 2; i++) {
+  maxLength = Math.min(maxLength, len / 2);
+  for (var i = 0; i < maxLength; i++) {
     var byte = parseInt(string.substr(i * 2, 2), 16);
     if (isNaN(byte)) throw new Error('Invalid hex string');
     this[offset + i] = byte;
@@ -109,38 +111,48 @@ SlowBuffer.prototype.hexWrite = function(string, offset) {
 };
 
 
-SlowBuffer.prototype.write = function(string, offset, encoding) {
-  // Support both (string, offset, encoding)
-  // and the legacy (string, encoding, offset)
-  if (!isFinite(offset)) {
-    var swap = encoding;
-    encoding = offset;
-    offset = swap;
+SlowBuffer.prototype.write = function(string) {
+  // Support both (string, offset, maxLength, encoding)
+  // and the legacy (string, encoding, offset, maxLength)
+  var offset, maxLength, encoding;
+  if (isFinite(arguments[1])) {
+    offset = arguments[1];
+    if (isFinite(arguments[2])) {
+      maxLength = arguments[2];
+      encoding = arguments[3];
+    } else {
+      encoding = arguments[2];
+    }
+  } else {  // legacy
+    encoding = arguments[1];
+    offset = arguments[2];
+    maxLength = arguments[3];
   }
 
   offset = +offset || 0;
+  maxLength = Math.min(+maxLength || Number.MAX_VALUE, this.length - offset);
   encoding = String(encoding || 'utf8').toLowerCase();
 
   switch (encoding) {
     case 'hex':
-      return this.hexWrite(string, offset);
+      return this.hexWrite(string, offset, maxLength);
 
     case 'utf8':
     case 'utf-8':
-      return this.utf8Write(string, offset);
+      return this.utf8Write(string, offset, maxLength);
 
     case 'ascii':
-      return this.asciiWrite(string, offset);
+      return this.asciiWrite(string, offset, maxLength);
 
     case 'binary':
-      return this.binaryWrite(string, offset);
+      return this.binaryWrite(string, offset, maxLength);
 
     case 'base64':
-      return this.base64Write(string, offset);
+      return this.base64Write(string, offset, maxLength);
 
     case 'ucs2':
     case 'ucs-2':
-      return this.ucs2Write(string, offset);
+      return this.ucs2Write(string, offset, maxLength);
 
     default:
       throw new Error('Unknown encoding');
@@ -271,20 +283,29 @@ Buffer.prototype.set = function set(i, v) {
 };
 
 
-// write(string, offset = 0, encoding = 'utf8')
-Buffer.prototype.write = function(string, offset, encoding) {
-  if (!isFinite(offset)) {
-    var swap = encoding;
-    encoding = offset;
-    offset = swap;
+// write(string, offset = 0, maxLength = this.length-offset, encoding = 'utf8')
+Buffer.prototype.write = function(string) {
+  // Support both (string, offset, maxLength, encoding)
+  // and the legacy (string, encoding, offset, maxLength)
+  var offset, maxLength, encoding;
+  if (isFinite(arguments[1])) {
+    offset = arguments[1];
+    if (isFinite(arguments[2])) {
+      maxLength = arguments[2];
+      encoding = arguments[3];
+    } else {
+      encoding = arguments[2];
+    }
+  } else {  // legacy
+    encoding = arguments[1];
+    offset = arguments[2];
+    maxLength = arguments[3];
   }
 
   offset = +offset || 0;
+  maxLength = Math.min(+maxLength || Number.MAX_VALUE, this.length - offset);
   encoding = String(encoding || 'utf8').toLowerCase();
 
-  // Make sure we are not going to overflow
-  var maxLength = this.length - offset;
-
   var ret;
   switch (encoding) {
     case 'hex':
diff --git a/test/simple/test-buffer.js b/test/simple/test-buffer.js
index 9cd4959aa2c..8c0d4fb0ddf 100644
--- a/test/simple/test-buffer.js
+++ b/test/simple/test-buffer.js
@@ -621,3 +621,50 @@ assert.equal(written, 9);
 written = buf.write('あいう\0'); // 3bytes * 3 + 1byte
 assert.equal(written, 10);
 
+// #243 Test write() with maxLength
+var buf = new Buffer(4);
+buf.fill(0xFF);
+var written = buf.write('abcd', 1, 2, 'utf8');
+console.log(buf);
+assert.equal(written, 2);
+assert.equal(buf[0], 0xFF);
+assert.equal(buf[1], 0x61);
+assert.equal(buf[2], 0x62);
+assert.equal(buf[3], 0xFF);
+
+buf.fill(0xFF);
+written = buf.write('abcd', 1, 4);
+console.log(buf);
+assert.equal(written, 3);
+assert.equal(buf[0], 0xFF);
+assert.equal(buf[1], 0x61);
+assert.equal(buf[2], 0x62);
+assert.equal(buf[3], 0x63);
+
+buf.fill(0xFF);
+written = buf.write('abcd', 'utf8', 1, 2);  // legacy style
+console.log(buf);
+assert.equal(written, 2);
+assert.equal(buf[0], 0xFF);
+assert.equal(buf[1], 0x61);
+assert.equal(buf[2], 0x62);
+assert.equal(buf[3], 0xFF);
+
+buf.fill(0xFF);
+written = buf.write('abcdef', 1, 2, 'hex');
+console.log(buf);
+assert.equal(written, 2);
+assert.equal(buf[0], 0xFF);
+assert.equal(buf[1], 0xAB);
+assert.equal(buf[2], 0xCD);
+assert.equal(buf[3], 0xFF);
+
+buf.fill(0xFF);
+written = buf.write('abcd', 0, 2, 'ucs2');
+console.log(buf);
+assert.equal(written, 2);
+assert.equal(buf[0], 0x61);
+assert.equal(buf[1], 0x00);
+assert.equal(buf[2], 0xFF);
+assert.equal(buf[3], 0xFF);
+