chore: get rid of extra base64 png parsing

Author: KuznetsovRoman

src/image.ts

@@ -3,7 +3,14 @@ import looksSame from "looks-same";
 import { loadEsm } from "./utils/preload-utils";
 import { DiffOptions, ImageSize } from "./types";
 import { convertRgbaToPng } from "./utils/eight-bit-rgba-to-png";
-import { BITS_IN_BYTE, PNG_HEIGHT_OFFSET, PNG_WIDTH_OFFSET, RGBA_CHANNELS } from "./constants/png";
+import {
+    BITS_IN_BYTE,
+    PNG_HEIGHT_OFFSET,
+    PNG_MIN_ASSIST_BYTES,
+    PNG_SIGNATURE,
+    PNG_WIDTH_OFFSET,
+    RGBA_CHANNELS,
+} from "./constants/png";
 
 interface PngImageData {
     data: Buffer;
@@ -49,6 +56,28 @@ const jsquashDecode = (buffer: ArrayBuffer): Promise<ImageData> => {
     ]).then(([mod]) => mod.decode(buffer, { bitDepth: BITS_IN_BYTE }));
 };
 
+export const extractBase64PngSize = (base64EncodedString: string): ImageSize => {
+    // Each 6 bits sequence encoded with 1 base64 char
+    const bytesToBase64CharsRatio = 8 / 6;
+
+    if (base64EncodedString.length <= PNG_MIN_ASSIST_BYTES * bytesToBase64CharsRatio) {
+        throw new Error("Invalid base64 encoded png: too short");
+    }
+
+    const headerBytesToRead = Math.max(PNG_WIDTH_OFFSET, PNG_HEIGHT_OFFSET) + 4;
+    const headerCharsToRead = headerBytesToRead * bytesToBase64CharsRatio;
+    const pngHeader = Buffer.from(base64EncodedString.slice(0, headerCharsToRead), "base64");
+
+    if (!pngHeader.subarray(0, PNG_SIGNATURE.byteLength).equals(PNG_SIGNATURE)) {
+        throw new Error("Invalid base64 encoded png: signature missmatch");
+    }
+
+    return {
+        width: pngHeader.readUInt32BE(PNG_WIDTH_OFFSET),
+        height: pngHeader.readUInt32BE(PNG_HEIGHT_OFFSET),
+    };
+};
+
 export class Image {
     private _imgDataPromise: Promise<Buffer>;
     private _imgData: Buffer | null = null;

src/worker/runner/test-runner/one-time-screenshooter.js

@@ -1,6 +1,6 @@
 "use strict";
 
-const { Image } = require("../../../image");
+const { extractBase64PngSize } = require("../../../image");
 const ScreenShooter = require("../../../browser/screen-shooter");
 const logger = require("../../../utils/logger");
 const { promiseTimeout } = require("../../../utils/promise");
@@ -106,8 +106,7 @@ module.exports = class OneTimeScreenshooter {
 
     async _makeViewportScreenshot() {
         const base64 = await this._browser.publicAPI.takeScreenshot();
-        const image = Image.fromBase64(base64);
-        const size = await image.getSize();
+        const size = extractBase64PngSize(base64);
 
         return { base64, size };
     }

test/src/image.js

@@ -1,6 +1,7 @@
 "use strict";
 
 const proxyquire = require("proxyquire");
+const { extractBase64PngSize } = require("src/image");
 
 describe("Image", () => {
     const sandbox = sinon.createSandbox();
@@ -53,6 +54,36 @@ describe("Image", () => {
 
     afterEach(() => sandbox.restore());
 
+    describe("extractBase64PngSize", () => {
+        it("should throw error on invalid small strings", () => {
+            const fn = () => extractBase64PngSize("foobar");
+
+            assert.throw(fn, "Invalid base64 encoded png: too short");
+        });
+
+        it("should throw error on non-base64 png strings", () => {
+            const fn = () => extractBase64PngSize("foobar".repeat(20));
+
+            assert.throw(fn, "Invalid base64 encoded png: signature missmatch");
+        });
+
+        it("should work with minimal png", () => {
+            const minimalPng =
+                "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVQImWNgYGAAAAAEAAGjChXjAAAAAElFTkSuQmCC";
+            const result = extractBase64PngSize(minimalPng);
+
+            assert.deepEqual(result, { width: 1, height: 1 });
+        });
+
+        it("should extract size", () => {
+            const tenPxSquarePng =
+                "iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs+9AAAAFUlEQVR42mP8z8BQz0AEYBxVSF+FABJADveWkH6oAAAAAElFTkSuQmCC";
+            const result = extractBase64PngSize(tenPxSquarePng);
+
+            assert.deepEqual(result, { width: 10, height: 10 });
+        });
+    });
+
     describe("constructor", () => {
         it("should read width and height from PNG buffer", () => {
             const buffer = createMockPngBuffer(200, 150);