Merge pull request #1207 from gemini-testing/TESTPLANE-896.base64_size_parse

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,30 @@ const jsquashDecode = (buffer: ArrayBuffer): Promise<ImageData> => {
     ]).then(([mod]) => mod.decode(buffer, { bitDepth: BITS_IN_BYTE }));
 };
 
+export const extractBase64PngSize = (base64EncodedPng: string): ImageSize => {
+    // Strips data URI prefix if it exists
+    const base64Data = base64EncodedPng.includes(";base64,")
+        ? (base64EncodedPng.split(";base64,").pop() as string)
+        : base64EncodedPng;
+
+    if (base64Data.length <= PNG_MIN_ASSIST_BYTES) {
+        throw new Error("Invalid base64 encoded png: too short");
+    }
+
+    const headerBytesToRead = Math.max(PNG_WIDTH_OFFSET, PNG_HEIGHT_OFFSET) + 4;
+    const headerCharsToRead = Math.ceil(headerBytesToRead / 3) * 4;
+    const pngHeader = Buffer.from(base64Data.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);

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

@@ -11,6 +11,7 @@ describe("worker/runner/test-runner/one-time-screenshooter", () => {
     const sandbox = sinon.createSandbox();
     let OneTimeScreenshooter;
     let logger;
+    let extractBase64PngSize;
 
     const mkBrowser_ = (opts = {}) => {
         const session = mkSessionStub_();
@@ -57,8 +58,10 @@ describe("worker/runner/test-runner/one-time-screenshooter", () => {
         logger = {
             warn: sinon.stub(),
         };
+        extractBase64PngSize = sinon.stub().named("extractBase64PngSize").returns({ width: 100500, height: 500100 });
         OneTimeScreenshooter = proxyquire("src/worker/runner/test-runner/one-time-screenshooter", {
             "../../../utils/logger": logger,
+            "../../../image": { extractBase64PngSize },
         });
 
         sandbox.stub(ScreenShooter.prototype, "capture").resolves(stubImage_());
@@ -71,8 +74,7 @@ describe("worker/runner/test-runner/one-time-screenshooter", () => {
         it('should capture viewport screenshot if option "takeScreenshotOnFailsMode" is not set', async () => {
             const browser = mkBrowser_();
             browser.publicAPI.takeScreenshot.resolves("base64");
-            const imgStub = stubImage_({ width: 100, height: 500 });
-            Image.fromBase64.returns(imgStub);
+            extractBase64PngSize.withArgs("base64").returns({ width: 100, height: 500 });
             const screenshooter = mkScreenshooter_({ browser });
 
             await screenshooter[method](...getArgs());
@@ -86,8 +88,7 @@ describe("worker/runner/test-runner/one-time-screenshooter", () => {
         it('should capture viewport screenshot if option "takeScreenshotOnFailsMode" is set to "viewport"', async () => {
             const browser = mkBrowser_();
             browser.publicAPI.takeScreenshot.resolves("base64");
-            const imgStub = stubImage_({ width: 100, height: 500 });
-            Image.fromBase64.returns(imgStub);
+            extractBase64PngSize.withArgs("base64").returns({ width: 100, height: 500 });
             const config = { takeScreenshotOnFailsMode: "viewport" };
             const screenshooter = mkScreenshooter_({ browser, config });
 
@@ -213,8 +214,8 @@ describe("worker/runner/test-runner/one-time-screenshooter", () => {
         it("should extend passed error with screenshot data", async () => {
             const browser = mkBrowser_();
             browser.publicAPI.takeScreenshot.resolves("base64");
+            extractBase64PngSize.withArgs("base64").returns({ width: 100, height: 200 });
             const screenshooter = mkScreenshooter_({ browser });
-            Image.fromBase64.withArgs("base64").returns(stubImage_({ width: 100, height: 200 }));
 
             const error = await screenshooter.extendWithScreenshot(new Error());
 
@@ -286,7 +287,7 @@ describe("worker/runner/test-runner/one-time-screenshooter", () => {
 
     describe("getScreenshot", () => {
         it("should return captured screenshot", async () => {
-            Image.fromBase64.returns(stubImage_({ width: 100, height: 200 }));
+            extractBase64PngSize.returns({ width: 100, height: 200 });
 
             const screenshooter = mkScreenshooter_({});