Added passthrough stream to track size of bytes downloaded

postmanlabs · Jun 20, 2024 · 5ef3270 · 5ef3270
1 parent 00acf08
commit 5ef3270
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 17 deletions.
diff --git a/lib/helpers.js b/lib/helpers.js
@@ -3,6 +3,7 @@
 var jsonSafeStringify = require('json-stringify-safe')
 var crypto = require('crypto')
 var Buffer = require('safe-buffer').Buffer
+var { Transform } = require('stream')
 
 var defer = typeof setImmediate === 'undefined'
   ? process.nextTick
@@ -58,10 +59,28 @@ function version () {
   }
 }
 
+class SizeTrackerStream extends Transform {
+  constructor (options) {
+    super(options)
+    this.size = 0
+  }
+
+  _transform (chunk, encoding, callback) {
+    this.size += chunk.length
+    this.push(chunk)
+    callback()
+  }
+
+  _flush (callback) {
+    callback()
+  }
+}
+
 exports.safeStringify = safeStringify
 exports.md5 = md5
 exports.isReadStream = isReadStream
 exports.toBase64 = toBase64
 exports.copy = copy
 exports.version = version
 exports.defer = defer
+exports.SizeTrackerStream = SizeTrackerStream
diff --git a/request.js b/request.js
@@ -41,6 +41,7 @@ var toBase64 = helpers.toBase64
 var defer = helpers.defer
 var copy = helpers.copy
 var version = helpers.version
+var SizeTrackerStream = helpers.SizeTrackerStream
 var globalCookieJar = cookies.jar()
 
 var globalPool = {}
@@ -1337,6 +1338,8 @@ Request.prototype.onRequestResponse = function (response) {
     }
 
     var responseContent
+    var downloadSizeTracker = new SizeTrackerStream()
+
     if ((self.gzip || self.brotli) && !noBody(response.statusCode)) {
       var contentEncoding = response.headers['content-encoding'] || 'identity'
       contentEncoding = contentEncoding.trim().toLowerCase()
@@ -1352,23 +1355,23 @@ Request.prototype.onRequestResponse = function (response) {
 
       if (self.gzip && contentEncoding === 'gzip') {
         responseContent = zlib.createGunzip(zlibOptions)
-        response.pipe(responseContent)
+        response.pipe(downloadSizeTracker).pipe(responseContent)
       } else if (self.gzip && contentEncoding === 'deflate') {
         responseContent = inflate.createInflate(zlibOptions)
-        response.pipe(responseContent)
+        response.pipe(downloadSizeTracker).pipe(responseContent)
       } else if (self.brotli && contentEncoding === 'br') {
         responseContent = brotli.createBrotliDecompress()
-        response.pipe(responseContent)
+        response.pipe(downloadSizeTracker).pipe(responseContent)
       } else {
         // Since previous versions didn't check for Content-Encoding header,
         // ignore any invalid values to preserve backwards-compatibility
         if (contentEncoding !== 'identity') {
           debug('ignoring unrecognized Content-Encoding ' + contentEncoding)
         }
-        responseContent = response
+        responseContent = response.pipe(downloadSizeTracker)
       }
     } else {
-      responseContent = response
+      responseContent = response.pipe(downloadSizeTracker)
     }
 
     if (self.encoding) {
@@ -1398,9 +1401,9 @@ Request.prototype.onRequestResponse = function (response) {
     // results in some other characters.
     // For example: If the server intentionally responds with `ð\x9F\x98\x8A` as status message
     // but if the statusMessageEncoding option is set to `utf8`, then it would get converted to '😊'.
-    var statusMessage = String(responseContent.statusMessage)
+    var statusMessage = String(response.statusMessage)
     if (self.statusMessageEncoding && /[^\w\s-']/.test(statusMessage)) {
-      responseContent.statusMessage = Buffer.from(statusMessage, 'latin1').toString(self.statusMessageEncoding)
+      response.statusMessage = Buffer.from(statusMessage, 'latin1').toString(self.statusMessageEncoding)
     }
 
     if (self._paused) {
@@ -1445,6 +1448,7 @@ Request.prototype.onRequestResponse = function (response) {
       self.emit('data', chunk)
     })
     responseContent.once('end', function (chunk) {
+      self._reqResInfo.response.downloadedBytes = downloadSizeTracker.size
       self.emit('end', chunk)
     })
     responseContent.on('error', function (error) {

diff --git a/tests/test-agent.js b/tests/test-agent.js
@@ -23,16 +23,16 @@ function httpAgent (t, options, req) {
   var r = (req || request)(options, function (_err, res, body) {
     t.ok(r.agent instanceof http.Agent, 'is http.Agent')
     t.equal(r.agent.options.keepAlive, true, 'is keepAlive')
-    t.equal(Object.keys(r.agent.sockets).length, 1, '1 socket name')
-
+    t.equal(Object.keys(r.agent.freeSockets).length, 1, '1 socket name')
     var name = (typeof r.agent.getName === 'function')
       ? r.agent.getName({port: s.port})
       : 'localhost:' + s.port // node 0.10-
-    t.equal(r.agent.sockets[name].length, 1, '1 open socket')
 
-    var socket = r.agent.sockets[name][0]
+    t.equal(r.agent.freeSockets[name].length, 1, '1 open socket')
+
+    var socket = r.agent.freeSockets[name][0]
     socket.on('close', function () {
-      t.equal(Object.keys(r.agent.sockets).length, 0, '0 open sockets')
+      t.equal(Object.keys(r.agent.freeSockets).length, 0, '0 open sockets')
       t.end()
     })
     socket.end()

diff --git a/tests/test-node-debug.js b/tests/test-node-debug.js
@@ -37,7 +37,7 @@ tape('a simple request should not fail with debugging enabled', function (t) {
     t.ok(stderr.length, 'stderr has some messages')
     var url = s.url.replace(/\//g, '\\/')
     var patterns = [
-      /^REQUEST { uri: /,
+      /^REQUEST {\n {2}uri: /,
       new RegExp('^REQUEST make request ' + url + '/\n$'),
       /^REQUEST onRequestResponse /,
       /^REQUEST finish init /,

diff --git a/tests/test-verbose.js b/tests/test-verbose.js
@@ -75,7 +75,7 @@ tape('HTTP: verbose=true', function (t) {
     t.deepEqual(Object.keys(debug[0].session), ['id', 'reused', 'data'])
     t.deepEqual(Object.keys(debug[0].session.data), ['addresses'])
     t.equal(debug[0].session.reused, false)
-    t.deepEqual(Object.keys(debug[0].response), ['statusCode', 'headers', 'httpVersion'])
+    t.deepEqual(Object.keys(debug[0].response), ['statusCode', 'headers', 'httpVersion', 'downloadedBytes'])
 
     t.notEqual(debug[0].response.headers.length, 0)
     t.equal(debug[0].response.headers[0].key, 'Date')
@@ -113,7 +113,7 @@ tape('HTTP: redirect(HTTPS) + verbose=true', function (t) {
     t.deepEqual(Object.keys(debug[1].session.data), ['addresses', 'tls'])
     t.deepEqual(Object.keys(debug[1].session.data.tls), ['reused', 'authorized', 'authorizationError', 'cipher', 'protocol', 'ephemeralKeyInfo', 'peerCertificate'])
     t.equal(debug[1].session.reused, false)
-    t.deepEqual(Object.keys(debug[1].response), ['statusCode', 'headers', 'httpVersion'])
+    t.deepEqual(Object.keys(debug[1].response), ['statusCode', 'headers', 'httpVersion', 'downloadedBytes'])
 
     t.end()
   })
@@ -140,7 +140,7 @@ tape('HTTPS: verbose=true', function (t) {
     t.deepEqual(Object.keys(debug[0].session.data), ['addresses', 'tls'])
     t.deepEqual(Object.keys(debug[0].session.data.tls), ['reused', 'authorized', 'authorizationError', 'cipher', 'protocol', 'ephemeralKeyInfo', 'peerCertificate'])
     t.equal(debug[0].session.reused, false)
-    t.deepEqual(Object.keys(debug[0].response), ['statusCode', 'headers', 'httpVersion'])
+    t.deepEqual(Object.keys(debug[0].response), ['statusCode', 'headers', 'httpVersion', 'downloadedBytes'])
 
     t.end()
   })
@@ -174,7 +174,7 @@ tape('HTTPS: redirect(HTTP) + verbose=true', function (t) {
     t.deepEqual(Object.keys(debug[1].session), ['id', 'reused', 'data'])
     t.deepEqual(Object.keys(debug[1].session.data), ['addresses'])
     t.equal(debug[1].session.reused, false)
-    t.deepEqual(Object.keys(debug[1].response), ['statusCode', 'headers', 'httpVersion'])
+    t.deepEqual(Object.keys(debug[1].response), ['statusCode', 'headers', 'httpVersion', 'downloadedBytes'])
 
     t.end()
   })