feat: add memory and resolution limits

patrickhulce · patrickhulce · commit 135705b1510a · 2020-04-23T08:59:31.000-05:00
BREAKING CHANGE: images larger than 100 megapixels or requiring more than 512MB of memory to decode will throw unless `maxMemoryInMB` and `maxResolutionInMP` options are increased
diff --git a/lib/decoder.js b/lib/decoder.js
@@ -359,6 +359,7 @@ var JpegImage = (function jpegImage() {
     var blocksPerLine = component.blocksPerLine;
     var blocksPerColumn = component.blocksPerColumn;
     var samplesPerLine = blocksPerLine << 3;
+    // Only 1 used per invocation of this function and garbage collected after invocation, so no need to account for its memory footprint.
     var R = new Int32Array(64), r = new Uint8Array(64);
 
     // A port of poppler's IDCT method which in turn is taken from:
@@ -521,6 +522,8 @@ var JpegImage = (function jpegImage() {
       }
     }
 
+    requestMemoryAllocation(samplesPerLine * blocksPerColumn * 8);
+
     var i, j;
     for (var blockRow = 0; blockRow < blocksPerColumn; blockRow++) {
       var scanLine = blockRow << 3;
@@ -559,6 +562,7 @@ var JpegImage = (function jpegImage() {
       xhr.send(null);
     },
     parse: function parse(data) {
+      var maxResolutionInPixels = this.opts.maxResolutionInMP * 1000 * 1000;
       var offset = 0, length = data.length;
       function readUint16() {
         var value = (data[offset] << 8) | data[offset + 1];
@@ -590,7 +594,12 @@ var JpegImage = (function jpegImage() {
             var blocksPerColumn = Math.ceil(Math.ceil(frame.scanLines  / 8) * component.v / maxV);
             var blocksPerLineForMcu = mcusPerLine * component.h;
             var blocksPerColumnForMcu = mcusPerColumn * component.v;
+            var blocksToAllocate = blocksPerColumnForMcu * blocksPerLineForMcu;
             var blocks = [];
+
+            // Each block is a Int32Array of length 64 (4 x 64 = 256 bytes)
+            requestMemoryAllocation(blocksToAllocate * 256);
+
             for (var i = 0; i < blocksPerColumnForMcu; i++) {
               var row = [];
               for (var j = 0; j < blocksPerLineForMcu; j++)
@@ -685,6 +694,7 @@ var JpegImage = (function jpegImage() {
             var quantizationTablesEnd = quantizationTablesLength + offset - 2;
             while (offset < quantizationTablesEnd) {
               var quantizationTableSpec = data[offset++];
+              requestMemoryAllocation(64 * 4);
               var tableData = new Int32Array(64);
               if ((quantizationTableSpec >> 4) === 0) { // 8 bit values
                 for (j = 0; j < 64; j++) {
@@ -714,6 +724,13 @@ var JpegImage = (function jpegImage() {
             frame.samplesPerLine = readUint16();
             frame.components = {};
             frame.componentsOrder = [];
+
+            var pixelsInFrame = frame.scanLines * frame.samplesPerLine;
+            if (pixelsInFrame > maxResolutionInPixels) {
+              var exceededAmount = Math.ceil((pixelsInFrame - maxResolutionInPixels) / 1e6);
+              throw new Error(`maxResolutionInMP limit exceeded by ${exceededAmount}MP`);
+            }
+
             var componentsCount = data[offset++], componentId;
             var maxH = 0, maxV = 0;
             for (i = 0; i < componentsCount; i++) {
@@ -739,8 +756,10 @@ var JpegImage = (function jpegImage() {
               var huffmanTableSpec = data[offset++];
               var codeLengths = new Uint8Array(16);
               var codeLengthSum = 0;
-              for (j = 0; j < 16; j++, offset++)
+              for (j = 0; j < 16; j++, offset++) {
                 codeLengthSum += (codeLengths[j] = data[offset]);
+              }
+              requestMemoryAllocation(16 + codeLengthSum);
               var huffmanValues = new Uint8Array(codeLengthSum);
               for (j = 0; j < codeLengthSum; j++, offset++)
                 huffmanValues[j] = data[offset];
@@ -832,6 +851,7 @@ var JpegImage = (function jpegImage() {
       var Y, Cb, Cr, K, C, M, Ye, R, G, B;
       var colorTransform;
       var dataLength = width * height * this.components.length;
+      requestMemoryAllocation(dataLength);
       var data = new Uint8Array(dataLength);
       switch (this.components.length) {
         case 1:
@@ -1009,6 +1029,31 @@ var JpegImage = (function jpegImage() {
     }
   };
 
+
+  // We cap the amount of memory used by jpeg-js to avoid unexpected OOMs from untrusted content.
+  var totalBytesAllocated = 0;
+  var maxMemoryUsageBytes = 0;
+  function requestMemoryAllocation(increaseAmount = 0) {
+    var totalMemoryImpactBytes = totalBytesAllocated + increaseAmount;
+    if (totalMemoryImpactBytes > maxMemoryUsageBytes) {
+      var exceededAmount = Math.ceil((totalMemoryImpactBytes - maxMemoryUsageBytes) / 1024 / 1024);
+      throw new Error(`maxMemoryUsageInMB limit exceeded by at least ${exceededAmount}MB`);
+    }
+
+    totalBytesAllocated = totalMemoryImpactBytes;
+  }
+
+  constructor.resetMaxMemoryUsage = function (maxMemoryUsageBytes_) {
+    totalBytesAllocated = 0;
+    maxMemoryUsageBytes = maxMemoryUsageBytes_;
+  };
+
+  constructor.getBytesAllocated = function () {
+    return totalBytesAllocated;
+  };
+
+  constructor.requestMemoryAllocation = requestMemoryAllocation;
+
   return constructor;
 })();
 
@@ -1026,18 +1071,24 @@ function decode(jpegData, userOpts = {}) {
     colorTransform: undefined,
     formatAsRGBA: true,
     tolerantDecoding: false,
+    maxResolutionInMP: 100, // Don't decode more than 100 megapixels
+    maxMemoryUsageInMB: 512, // Don't decode if memory footprint is more than 512MB
   };
 
   var opts = {...defaultOpts, ...userOpts};
   var arr = new Uint8Array(jpegData);
   var decoder = new JpegImage();
   decoder.opts = opts;
+  // If this constructor ever supports async decoding this will need to be done differently.
+  // Until then, treating as singleton limit is fine.
+  JpegImage.resetMaxMemoryUsage(opts.maxMemoryUsageInMB * 1024 * 1024);
   decoder.parse(arr);
   decoder.colorTransform = opts.colorTransform;
 
   var channels = (opts.formatAsRGBA) ? 4 : 3;
   var bytesNeeded = decoder.width * decoder.height * channels;
   try {
+    JpegImage.requestMemoryAllocation(bytesNeeded);
     var image = {
       width: decoder.width,
       height: decoder.height,
diff --git a/test/fixtures/black-6000x6000.jpg b/test/fixtures/black-6000x6000.jpg
diff --git a/test/index.js b/test/index.js
@@ -7,6 +7,11 @@ function fixture(name) {
   return fs.readFileSync(path.join(__dirname, 'fixtures', name));
 }
 
+const SUPER_LARGE_JPEG_BASE64 =
+  '/9j/wJ39sP//DlKWvX+7xPlXkJa9f7v8DoDVAAD//zb6QAEAI2cBv3P/r4ADpX8Jf14AAAAAgCPE+VeQlr1/uwCAAAAVALNOjAGP2lIS';
+
+const SUPER_LARGE_JPEG_BUFFER = Buffer.from(SUPER_LARGE_JPEG_BASE64, 'base64');
+
 it('should be able to decode a JPEG', function () {
   var jpegData = fixture('grumpycat.jpg');
   var rawImageData = jpeg.decode(jpegData);
@@ -216,3 +221,27 @@ it('should be able to encode/decode image with exif data', function () {
   var loopImageData = jpeg.decode(new Uint8Array(encodedData.data));
   expect(loopImageData.exifBuffer).toEqual(imageData.exifBuffer);
 });
+
+it('should be able to decode large images within memory limits', () => {
+  var jpegData = fixture('black-6000x6000.jpg');
+  var rawImageData = jpeg.decode(jpegData);
+  expect(rawImageData.width).toEqual(6000);
+  expect(rawImageData.height).toEqual(6000);
+}, 30000);
+
+// See https://github.com/eugeneware/jpeg-js/issues/53
+it('should limit resolution exposure', function () {
+  expect(() => jpeg.decode(SUPER_LARGE_JPEG_BUFFER)).toThrow(
+    'maxResolutionInMP limit exceeded by 141MP',
+  );
+});
+
+it('should limit memory exposure', function () {
+  expect(() => jpeg.decode(SUPER_LARGE_JPEG_BUFFER, {maxResolutionInMP: 500})).toThrow(
+    /maxMemoryUsageInMB limit exceeded by at least \d+MB/,
+  );
+
+  // Make sure the limit resets each decode.
+  var jpegData = fixture('grumpycat.jpg');
+  expect(() => jpeg.decode(jpegData)).not.toThrow();
+}, 30000);
