fix when source stream emits more than 1 chunk

CHANGELOG.md

@@ -1,5 +1,8 @@
 # truncate-stream changelog
 
-## 1.0.0 (2014/11/12)
+## 1.0.1 (2014/11/12)
+
+  - fix streams that emit more than 1 chunk
 
-  - initial version
+## 1.0.0 (2014/11/12)
+*don't use; broken for streams that emit more than 1 chunk*

package.json

@@ -1,6 +1,6 @@
 {
   "name": "truncate-stream",
-  "version": "1.0.0",
+  "version": "1.0.1",
   "description": "stream which truncates input after N bytes",
   "main": "truncate-stream.js",
   "scripts": {

test.js

@@ -1,17 +1,43 @@
 var expect = require('chai').expect;
 
+var stream = require('stream');
 var RandStream = require('randstream');
 var concat = require('concat-stream');
 
 var TruncateStream = require('./truncate-stream');
 
-describe('truncate-stream', function(done) {
-  it('should truncate the input', function (done) {
+describe('truncate-stream', function() {
+
+  it('should truncate the input when enough is provided in the first chunk', function (done) {
     var source = new RandStream();
     var truncate = new TruncateStream({maxBytes: 1024});
 
     source.pipe(truncate).pipe(concat(function(data) {
-      expect(data.length).to.eql(1024);
+      expect(data).to.have.lengthOf(1024);
+      done();
+    }));
+  });
+
+  it('should truncate the input with multi-chunk streams', function (done) {
+    var source = new RandStream();
+    var truncate = new TruncateStream({maxBytes: 100000});
+
+    source.pipe(truncate).pipe(concat(function(data) {
+      expect(data).to.have.lengthOf(100000);
+      done();
+    }));
+  });
+
+  it('should not truncate a stream that was smaller than the maxBytes', function (done) {
+    var source = new stream.Readable();
+    source._read = function () {
+      this.push(new Buffer(100));
+      this.push(null);
+    };
+    var truncate = new TruncateStream({maxBytes: 1024});
+
+    source.pipe(truncate).pipe(concat(function (data) {
+      expect(data).to.have.lengthOf(100);
       done();
     }));
   });

truncate-stream.js

@@ -8,20 +8,46 @@ function TruncateStream(opts) {
 
   this._truncateState = {
     bytes: 0,
-    maxBytes: opts.maxBytes
+    maxBytes: opts.maxBytes,
+    truncated: false,
+    src: []
   };
+
+  this.on('pipe', function (source) {
+    this._truncateState.src.push(source);
+  });
 }
 
 util.inherits(TruncateStream, stream.Transform);
 module.exports = TruncateStream;
 
 TruncateStream.prototype._transform = function _transform(chunk, enc, callback) {
   var state = this._truncateState;
+  if (state.truncated) return callback();
+
   state.bytes += chunk.length;
 
   // are we there yet?
-  if (state.bytes < state.maxBytes) return this.push(chunk);
+  if (state.bytes < state.maxBytes) {
+    // nope
+    this.push(chunk);
+    return callback();
+  }
+
+  state.truncated = true;
 
+  // push exactly what's left to make up maxBytes
   this.push(chunk.slice(0, chunk.length - (state.bytes - state.maxBytes)));
-  this.push(null);
-};
+
+  // unpipe ourselves and redirect their output to /dev/null
+  for (var i = 0; i < state.src.length; i++) {
+    state.src[i].unpipe(this);
+    state.src[i].on('data', devnull);
+  }
+
+  state.src = null;
+  this.end();
+  callback();
+};
+
+function devnull() {}