From edd003c230fc72be7a9cd6e52a392eca60930eb0 Mon Sep 17 00:00:00 2001
From: Max Schmitt <max@schmitt.mx>
Date: Mon, 15 May 2023 23:32:16 +0200
Subject: [PATCH] chore: validate expected image buffer when comparing images
 (#23030)

Fixes https://github.com/microsoft/playwright/issues/23012
---
 .../playwright-core/src/utils/comparators.ts  | 13 +++++++++++
 .../to-have-screenshot.spec.ts                | 22 +++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/packages/playwright-core/src/utils/comparators.ts b/packages/playwright-core/src/utils/comparators.ts
index 790f58056a..89b5cb2976 100644
--- a/packages/playwright-core/src/utils/comparators.ts
+++ b/packages/playwright-core/src/utils/comparators.ts
@@ -52,6 +52,7 @@ type ImageData = { width: number, height: number, data: Buffer };
 function compareImages(mimeType: string, actualBuffer: Buffer | string, expectedBuffer: Buffer, options: ImageComparatorOptions = {}): ComparatorResult {
   if (!actualBuffer || !(actualBuffer instanceof Buffer))
     return { errorMessage: 'Actual result should be a Buffer.' };
+  validateBuffer(expectedBuffer, mimeType);
 
   let actual: ImageData = mimeType === 'image/png' ? PNG.sync.read(actualBuffer) : jpegjs.decode(actualBuffer, { maxMemoryUsageInMB: JPEG_JS_MAX_BUFFER_SIZE_IN_MB });
   let expected: ImageData = mimeType === 'image/png' ? PNG.sync.read(expectedBuffer) : jpegjs.decode(expectedBuffer, { maxMemoryUsageInMB: JPEG_JS_MAX_BUFFER_SIZE_IN_MB });
@@ -92,6 +93,18 @@ function compareImages(mimeType: string, actualBuffer: Buffer | string, expected
   return null;
 }
 
+function validateBuffer(buffer: Buffer, mimeType: string): void {
+  if (mimeType === 'image/png') {
+    const pngMagicNumber = [137, 80, 78, 71, 13, 10, 26, 10];
+    if (buffer.length < pngMagicNumber.length || !pngMagicNumber.every((byte, index) => buffer[index] === byte))
+      throw new Error('could not decode image as PNG.');
+  } else if (mimeType === 'image/jpeg') {
+    const jpegMagicNumber = [255, 216];
+    if (buffer.length < jpegMagicNumber.length || !jpegMagicNumber.every((byte, index) => buffer[index] === byte))
+      throw new Error('could not decode image as JPEG.');
+  }
+}
+
 function compareText(actual: Buffer | string, expectedBuffer: Buffer): ComparatorResult {
   if (typeof actual !== 'string')
     return { errorMessage: 'Actual result should be a string' };
diff --git a/tests/playwright-test/to-have-screenshot.spec.ts b/tests/playwright-test/to-have-screenshot.spec.ts
index c1a9ab7a2b..a5a46758e7 100644
--- a/tests/playwright-test/to-have-screenshot.spec.ts
+++ b/tests/playwright-test/to-have-screenshot.spec.ts
@@ -1152,6 +1152,28 @@ test('should respect comparator in config', async ({ runInlineTest }) => {
   expect(result.report.suites[0].specs[0].tests[1].status).toBe('unexpected');
 });
 
+test('should throw pretty error if expected PNG file is not a PNG', async ({ runInlineTest }) => {
+  const result = await runInlineTest({
+    ...playwrightConfig({
+      snapshotPathTemplate: '__screenshots__/{testFilePath}/{arg}{ext}',
+    }),
+    '__screenshots__/a.spec.js/snapshot.png': 'not a png',
+    '__screenshots__/a.spec.js/snapshot.jpg': 'not a jpg',
+    'a.spec.js': `
+      const { test, expect } = require('@playwright/test');
+      test('png', async ({ page }) => {
+        await expect(page).toHaveScreenshot('snapshot.png');
+      });
+      test('jpg', async ({ page }) => {
+        expect(await page.screenshot({ type: 'jpeg' })).toMatchSnapshot('snapshot.jpg')
+      });
+    `,
+  });
+  expect(result.exitCode).toBe(1);
+  expect(result.output).toContain('could not decode image as PNG.');
+  expect(result.output).toContain('could not decode image as JPEG.');
+});
+
 function playwrightConfig(obj: any) {
   return {
     'playwright.config.js': `