From ccd19d86952c3ce85a50432f83ae445b8cc7d700 Mon Sep 17 00:00:00 2001 From: pythongosssss <125205205+pythongosssss@users.noreply.github.com> Date: Sat, 2 May 2026 21:18:45 +0100 Subject: [PATCH] test: add metadata parser coverage (#11307) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Adds tests for metadata parsers ## Changes - **What**: - add test file generation script - identified & fixed bug in webp exif parsing over-reading - identified & fix bug in mp3/ogg parser where it would read from a fixed position instead of relative, causing incorrect reads throwing RangeError - added catch in latent + json parsing to resolve errors ┆Issue is synchronized with this [Notion page](https://www.notion.so/PR-11307-test-add-metadata-parser-coverage-3446d73d36508108ac36dddcec0a54d4) by [Unito](https://www.unito.io) --------- Co-authored-by: GitHub Action --- .../generate-embedded-metadata-test-files.py | 177 ++++++++++++++ src/scripts/metadata/__fixtures__/helpers.ts | 48 ++++ .../metadata/__fixtures__/with_metadata.avif | Bin 0 -> 552 bytes .../metadata/__fixtures__/with_metadata.flac | Bin 0 -> 8651 bytes .../metadata/__fixtures__/with_metadata.mp3 | Bin 0 -> 797 bytes .../metadata/__fixtures__/with_metadata.mp4 | Bin 0 -> 1094 bytes .../metadata/__fixtures__/with_metadata.opus | Bin 0 -> 906 bytes .../metadata/__fixtures__/with_metadata.webm | Bin 0 -> 4043 bytes .../metadata/__fixtures__/with_metadata.webp | Bin 0 -> 266 bytes .../with_metadata_exif_prefix.webp | Bin 0 -> 272 bytes src/scripts/metadata/avif.test.ts | 71 +++++- src/scripts/metadata/avif.ts | 1 + src/scripts/metadata/ebml.test.ts | 49 ++++ src/scripts/metadata/ebml.ts | 1 + src/scripts/metadata/flac.test.ts | 56 +++++ src/scripts/metadata/flac.ts | 2 + src/scripts/metadata/gltf.test.ts | 22 +- src/scripts/metadata/gltf.ts | 1 + src/scripts/metadata/isobmff.test.ts | 52 ++++ src/scripts/metadata/isobmff.ts | 1 + src/scripts/metadata/json.test.ts | 91 +++++++ src/scripts/metadata/json.ts | 31 ++- src/scripts/metadata/mp3.test.ts | 106 +++++++++ src/scripts/metadata/mp3.ts | 21 +- src/scripts/metadata/ogg.test.ts | 74 ++++++ src/scripts/metadata/ogg.ts | 17 +- src/scripts/metadata/png.test.ts | 141 +++++------ src/scripts/metadata/png.ts | 1 + src/scripts/metadata/svg.test.ts | 42 ++++ src/scripts/pnginfo.test.ts | 222 ++++++++++++------ src/scripts/pnginfo.ts | 49 ++-- 31 files changed, 1085 insertions(+), 191 deletions(-) create mode 100644 scripts/generate-embedded-metadata-test-files.py create mode 100644 src/scripts/metadata/__fixtures__/helpers.ts create mode 100644 src/scripts/metadata/__fixtures__/with_metadata.avif create mode 100644 src/scripts/metadata/__fixtures__/with_metadata.flac create mode 100644 src/scripts/metadata/__fixtures__/with_metadata.mp3 create mode 100644 src/scripts/metadata/__fixtures__/with_metadata.mp4 create mode 100644 src/scripts/metadata/__fixtures__/with_metadata.opus create mode 100644 src/scripts/metadata/__fixtures__/with_metadata.webm create mode 100644 src/scripts/metadata/__fixtures__/with_metadata.webp create mode 100644 src/scripts/metadata/__fixtures__/with_metadata_exif_prefix.webp create mode 100644 src/scripts/metadata/ebml.test.ts create mode 100644 src/scripts/metadata/flac.test.ts create mode 100644 src/scripts/metadata/isobmff.test.ts create mode 100644 src/scripts/metadata/json.test.ts create mode 100644 src/scripts/metadata/mp3.test.ts create mode 100644 src/scripts/metadata/ogg.test.ts create mode 100644 src/scripts/metadata/svg.test.ts diff --git a/scripts/generate-embedded-metadata-test-files.py b/scripts/generate-embedded-metadata-test-files.py new file mode 100644 index 0000000000..2b57e296d3 --- /dev/null +++ b/scripts/generate-embedded-metadata-test-files.py @@ -0,0 +1,177 @@ +#!/usr/bin/env python3 +""" +Generate test fixture files for metadata parser tests. + +Each fixture embeds the same workflow and prompt JSON, matching the +format the ComfyUI backend uses to write metadata. + +Prerequisites: + source ~/ComfyUI/.venv/bin/activate + python3 scripts/generate-embedded-metadata-test-files.py + +Output: src/scripts/metadata/__fixtures__/ +""" + +import json +import os +import struct +import subprocess + +import av +from PIL import Image + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +FIXTURES_DIR = os.path.join(REPO_ROOT, 'src', 'scripts', 'metadata', '__fixtures__') + +WORKFLOW = { + 'nodes': [ + { + 'id': 1, + 'type': 'KSampler', + 'pos': [100, 100], + 'size': [200, 200], + } + ] +} +PROMPT = {'1': {'class_type': 'KSampler', 'inputs': {}}} + +WORKFLOW_JSON = json.dumps(WORKFLOW, separators=(',', ':')) +PROMPT_JSON = json.dumps(PROMPT, separators=(',', ':')) + + +def out(name: str) -> str: + return os.path.join(FIXTURES_DIR, name) + + +def report(name: str): + size = os.path.getsize(out(name)) + print(f' {name} ({size} bytes)') + + +def make_1x1_image() -> Image.Image: + return Image.new('RGB', (1, 1), (255, 0, 0)) + + +def build_exif_bytes() -> bytes: + """Build EXIF bytes matching the backend's tag assignments. + + Backend: 0x010F (Make) = "workflow:", 0x0110 (Model) = "prompt:" + """ + img = make_1x1_image() + exif = img.getexif() + exif[0x010F] = f'workflow:{WORKFLOW_JSON}' + exif[0x0110] = f'prompt:{PROMPT_JSON}' + return exif.tobytes() + + +def inject_exif_prefix_in_webp(path: str): + """Prepend Exif\\0\\0 to the EXIF chunk in a WEBP file. + + PIL always strips this prefix, so we re-inject it to test that code path. + """ + data = bytearray(open(path, 'rb').read()) + off = 12 + while off < len(data): + chunk_type = data[off:off + 4] + chunk_len = struct.unpack_from(' + this.onerror?.(new ProgressEvent('error') as ProgressEvent) + ) + } + ) +} + +export function mockFileReaderAbort(method: ReadMethod): void { + vi.spyOn(FileReader.prototype, method).mockImplementation( + function (this: FileReader) { + queueMicrotask(() => + this.onabort?.(new ProgressEvent('abort') as ProgressEvent) + ) + } + ) +} + +export function mockFileReaderResult( + method: ReadMethod, + result: string | ArrayBuffer | null +): void { + vi.spyOn(FileReader.prototype, method).mockImplementation( + function (this: FileReader) { + Object.defineProperty(this, 'result', { + value: result, + configurable: true + }) + queueMicrotask(() => + this.onload?.(new ProgressEvent('load') as ProgressEvent) + ) + } + ) +} diff --git a/src/scripts/metadata/__fixtures__/with_metadata.avif b/src/scripts/metadata/__fixtures__/with_metadata.avif new file mode 100644 index 0000000000000000000000000000000000000000..7db0eb4c8f4a03b4e65de07329988fa1143738d6 GIT binary patch literal 552 zcmZWm%}&BV5S|t!4M8MAqH-{}2M?wPinrbnJpl)UH$u#&KXJ5WH*JL?Ho^E#K7mhQ zIGXqbzJs$xyg13X-+VK(Gt({rwESct`3!v^Qyijig~+3!+NlbE6k6b!>4 zs_>cBdx44y)g2K;QzuBNdKXG73LT&fxyq4gpcOU!wMgz~=yLT>N}HvILpUR4fUfQ3 zpLI>@RIBCwIr%98<|2CV12MOkY${wYX7+f=&}Ft|GI|=1*{na}p$xo;nM?{*V71#O zVPdivpQvKop#Wi$O)`3!NQ5$>eioy{33wde{_liS`Iu0%WtP#MeJJl7LaEs5)EiK! SwBG?1gZ0<*VEsM%Sp5Ky)NAAb literal 0 HcmV?d00001 diff --git a/src/scripts/metadata/__fixtures__/with_metadata.flac b/src/scripts/metadata/__fixtures__/with_metadata.flac new file mode 100644 index 0000000000000000000000000000000000000000..3d5b9d5e35b21d2cb7952c6f6846cb5d0a7c4d6f GIT binary patch literal 8651 zcmYfENpxmlU{Dfb5CT$dVVdNd?M8S1VjZqm6(9Kpa~eTPp%IGc_+6XjG9cvMU=I6o4|LU^E0qLtr!nMnhmU1V%$( zGz3ONU^E0qLtr!nMnhmU1V%$(Gz3ONU^E0qLtr!nMnhmU1V%$(Gz3ONU^E0qLtr!n zMnhmU1V%$(sE5G+A1gT+0{mR!lQT0*^74xk6)H8f!W^HQMSOc8UBR&Zul&1Rtz2Kb zmmO>2lRv!-Ro>xu7UNx1fZfTFFq!s#+;IC$YFVzNE4sRmn=pJ2)}7 zASbm*Nk=I&ub{M~7^tGQww3{^I|!(|JijPAEhoPms6Q`1B^9V5x>_kSMajxgM+wuE zf_#tyLjwaHAdb~hD$c9|+7oRA6aeDb+E}33!NIOTm+=A3_DL*DGc(dN)&nXAIU5T8 z|8_V4W5a3BytI5EhaZTUfGC_{8L8ku$pB1yw?fG?{S z_j_GFkytIp+n!MHB-nB1qe_#Sd5ou&(ei;(*}mn)O=X8pJDlM9_3%6oSC{|7g2u5CMVyZ-;YsH|*eDdQzy8!kmXvKP9x zVM6GU@IPC-B0N(RgeA9zK5FPxW6F&Pd=R3-(>6smGBf_K!<9E)4~365ueKK2;TmAF q;ZkGpg3Ct^TsrK<8Wa5?^pO7l{~v{qtc$$))W^}+6`Y7d@c{tTJ5uTZ literal 0 HcmV?d00001 diff --git a/src/scripts/metadata/__fixtures__/with_metadata.mp4 b/src/scripts/metadata/__fixtures__/with_metadata.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..de42638c14f37e04d5ef6a7b62cfe37e65f30b63 GIT binary patch literal 1094 zcmZuxOKTHR6h4!rQi{|d)q+Hbmv$k^8XR&zp0V9D<0{#0&x(rBnPFVH3n@e1UvWIa3I zmGoWs>Pjxp{_Pz=SD${^mCw#(kYqVXN8>QpSsucc@7+)a9%Gzi<%dHkRjn41ZpqYN zlZnT1CT{5NI0@>h?$ovE2hhK2qw~XHOo)T#Mw!@7sYF#)xK+@erzH?m%;J6tvMdNwz&H20nCQw+<@qnypx29hIFc$QrsEumi;1 zFlRX0KC8J>wox{GzwigZ}X^-=#pI)E0 z2yrK`@WQw<>ogTD(V2=t;A(w>=jn-#_;Bc|7*|0wgdfIZqk)*scyCX@Jm#0h6d@`B zbTk#d6rEN>{GTX}7=WkT0QVX~`c3Mt;IZ$6-eo87(Cc0WFGeybD zP)7+{0F_8E%>D^MB; zKt2P7%J25Pbe8|h{~dqS%(=3!{q6S;S9dR&IO%w%nT)sU{hP1E#%)q1HSgVT?xG6ZM>%hm>=wvx_@9+Rh0|NlPjs& zJXP$ax9wk7ANv-_;KhEy=KQ~SlLb@PmA>M*vU*+jC$6+)_l+I__1jjeUs=`Mp39hM zyLqibjQ!^=^Vx2m3m9$xyncAZuZ#cZtpB1HH!h}Wd@09g;#s>7j#9QU0eKtgQ@&^9Rml0`c?Z?Weo}^Gz6SoFuwZl_ujWbjzMnymoH5b3^ojA>%QE){kDO30mCYZs+wi& d3#$4vzHxDU+~v-|3QXWkSF5$_JvHt-0RRz(Rh0k$ literal 0 HcmV?d00001 diff --git a/src/scripts/metadata/__fixtures__/with_metadata.webm b/src/scripts/metadata/__fixtures__/with_metadata.webm new file mode 100644 index 0000000000000000000000000000000000000000..5f7d87b55798ceaa7354e74951bfb6f1f31b3dba GIT binary patch literal 4043 zcmcIneNay+N zAqk`rS&Shn#$8a-f`Z%Ff<;Hib^-CDRHOl2%Tj9vY`d-2dBvITAN$-y?K=C%{j4AZ%J` z{ULyZO8s#Ejf(KZJKw3VsL>PLpsar4gHU67pqEdAmIb6dV3yR>Rj<#l zucB|SsVMT>&~xaioKy~HBZo(4VVX*Vx~8^@F0S2KQ~yq_PE*^!CSYRW)zqrrsnpbO zhLNgaq+BjB)S%@7@I+B&Q4kU!#L!cdw|Mk%7IM6(V9i~cXhQqL8c|+!6D>7Qd6lVe zPw43q*_15P^>!Q)1}(x}6pLL*O0Y0DmpIHxn}M?NfmcKD$f9SJm|qQ2A!)7UE+;+G zd@ub8%_8QPkb$>Ix2zO>(rNDJgEAv6yTBaEU!+++OJUiYMmV3uSwhlStxex@#=kVL z5ZOA-IcZPgEmgu@ae+PAjhAMK;sCJ1v1K|9kk2)X_vjU6{MF+mqY9k3G}kS=#T^_yoSR3M8K@1TVF?J5x>b*0-qC3a7>eZkg7vNZPxqeg^6pc`JMeA8BE))$5N z%#iN(zZ9iQv&L@@F(lupd8;qf##y!|)ISg4Mq{*{Gt&4h$cFSzv+m*?Hw380R-OCvi$Z^1p%RT>eXlDxaGcP10O9TWF|?~Wb= z>eYeW?VbAufm%LL)7F{m1uA=ZH+CeK;OPk{fA!tyeVTg}(d<1x>^gFN6sY{M`+swj zpR4%w{6U8C3dqg-%D-ZhZY)&Ow?tROr`{JXEpo!+&2G_mlZ(|o&I}T>V+STl3IseD z2}lW63YPt>s<0(AQ|qs=Kj>m!nT1L33-Q-e_c&$P+t zwZE2!v@fEWvKT1qrc%?U1~W0LG^et5`+>v(*6yr z?QXh*k>g<2II4ELdJlKrKm5n!6Mhv$CRB*W{{Aj|-!6OauC_&+qx!I`x3BZHEARhs z_SE+F!T~aed0w8@#|6Cs%ETfEG6d4vnW1 zfKcS+4=GnCY0?;`4_k-J*ZDB9PmcemM>@rmu7bdZ6C3urLVs=GV4HBU|KK>Qr)y># zf4*nt8mq~@K<26sUBbATkz+2Q#km}U7J_83uejsf5nBC3&UXy)1V4GLXNedWiI&Hi?Z7rKNe%5>9Z0(g(r%s;k|K;@hU(Wvh zduPs69^sHLONL8}2X%6TPPs}v6N~#?Vz~iOjAaJRt(mMzgBb775M6d?RPSnv&wz;z z-)Wp{P^k_YR3&dIm#k9;gQ`O_SF9;r(x@&?8|IesriP?5KEo7ln431-JFL6?1k4P$ z%UT9{GuAu1YiC#gQHC34ovJc0TU}eW@vj$ZH~m?szWwIuKmRsz>xB@|8!#O||V*j}o+o&PFRzfHpG_7q5RyDifqw%3X& zwh#C^6?@YxHxw+-V1ZdE{+u^JL!A-^X;n72SmS@x<+0OKoDD(J3RDB!Kd6K_#&me|rJM_040%d!7G4@5=j6`E&;D9USs=8ByW2GDwjK5fb9iAu7nFvwcdNv!s>NBF=`fb zEKYrDU@$(D9q8N_rye*@bj}=eNcIt(?;D2WaGUpDr=ov&^rluf{QsT$>wSX?dEPML z)$sCC{g3C(*J0$N1(Mc*L4~rF-}1C``yfDRy^lcvP`Q${A*!}5?Qy5#BnY%`d*?g# zKXK=8%C%`{s9NYl0Z{YPYDjsOMZxj#s?!AfKA70|Aqf(0Xt&VDUFhMFxbtqOox4Eq z!MO7@3zoV-cRDbF!Xi994outa#IUyqVBZH5`~GDfGcF9aEfcjzpx@WguKv79#=~#z zOu76hYmx^$YptIHyH9sDz%dwR{W`corCQ38rNIJ*1-==qfK&R|I92H93|Y%~crkEX znz<4cp)+JqQj2^9N2lD_tei2-E?LX@8E0z0Kl{@!pFMoz5<(9`!cbT6%SBt$+3_n< z{9Zxg26HRSPB4?z-YW0*3`73-#gv;2CUmSZ9o4eO;WRNtVBO?eyJVPa;FvgpUxkg zuINu&_{vQF>|=`h)bQi4_x!{jQfNQl{Df{P@cI<`2ShTtm%H?p`2i%DziZ_$ykdFq zX;?Wz@0ypPtfqah_cM;hwuP=HMoS7ubaDkuup9Ov3>&m*@N;G{F+VR#L}FGx<9VE| z0ljdK)A3Xf@`%)hF0ch2Ch^;n20!&kr3=wH7AgeW6L^7oDkuW_sT5=h+MK9Pjo z%+0()f1fRnNU$d8?Fqr%qeZ!+@M9;3AQFUtBPif3Ly=MBH^YQUg1V0(yA@S7XY*1* zTEaR|G&dG8US)k7EU=D>|2n$7Va?Y~>wnB3i^xsnk#A5$_Qg-bH*R{^ob=V#AxDp4 zVgAbq%>GJ*E-fcVPtvwDkA-e+xfGTS*)Tlg=wcYUm=K(m#cSE$?frD}&cmP>NK&3Y zAt51t3f3VC2_bHMg75TtfmeO8v1kJKg&@dW}RpUPKaa(admb)b{zc{Pq(w@XxxZ;3RxhJuv zGBJajlKAfkb?*=gaQ8`pSCx{eulZl-LkV}N6nIre72qBgA8!$F%goObV_7XPTn6>} z)9QZv)i7Xrw7Utw{%W65m;bi#BdQR#Fj{iZw$s95~S}>vG?It)>zI^vt zR1i9DHKApV{!eJ&6CU_P20km{14c0~8Ucs$BB#TIRyUrl4f_?Dpor`A{&Ro$$b`<= P;oAPNY2`}~Ow0ZW*Wd1y literal 0 HcmV?d00001 diff --git a/src/scripts/metadata/__fixtures__/with_metadata.webp b/src/scripts/metadata/__fixtures__/with_metadata.webp new file mode 100644 index 0000000000000000000000000000000000000000..ab727515517d2426ba824f386fbdf7b29aafce0b GIT binary patch literal 266 zcmWIYbaP{3WMBw)bqWXzu!!JdU|`??Vh8|=C>Q{l7Z@4lGHNk00HK4L>H?;cU5q~% zm^U!|`}Dg0&kfPD#qYoWzx3d*9{;5W|1V8oU~rA_bXx<|=j+R$1*A9_m>BsPn1HN6 zAXZ}(0JDvO>Q{l7Z@4lGHNk00HK4L>H?;cU5q~% zm^U!|`}Dg0&kfPD#qYoWzx3d*9{;5W|1V8oU~rA_blU{f=US1O#=zj~%b*41a4;}2 z@-r|2S%E;T#wY-08w1%DK+I5{UzDAelV5ICt(2Fal3J`}6 vi.restoreAllMocks()) + +describe('AVIF metadata', () => { + it('extracts workflow and prompt from EXIF data in ISOBMFF boxes', async () => { + const bytes = fs.readFileSync(fixturePath) + const file = new File([bytes], 'test.avif', { type: 'image/avif' }) + + const result = await getFromAvifFile(file) + + expect(JSON.parse(result.workflow)).toEqual(EXPECTED_WORKFLOW) + expect(JSON.parse(result.prompt)).toEqual(EXPECTED_PROMPT) + }) + + it('returns empty for non-AVIF data', async () => { + vi.spyOn(console, 'error').mockImplementation(() => {}) + const file = new File([new Uint8Array(16)], 'fake.avif') + + const result = await getFromAvifFile(file) + + expect(result).toEqual({}) + expect(console.error).toHaveBeenCalledWith('Not a valid AVIF file') + }) + + it('returns empty when AVIF has valid ftyp but corrupt internal boxes', async () => { + vi.spyOn(console, 'error').mockImplementation(() => {}) + + const buf = new Uint8Array(40) + const dv = new DataView(buf.buffer) + dv.setUint32(0, 16) + buf.set(new TextEncoder().encode('ftypavif'), 4) + dv.setUint32(16, 24) + buf.set(new TextEncoder().encode('meta'), 20) + + const file = new File([buf], 'corrupt.avif', { type: 'image/avif' }) + const result = await getFromAvifFile(file) + + expect(result).toEqual({}) + expect(console.error).toHaveBeenCalledWith( + expect.stringContaining('Error parsing AVIF metadata'), + expect.anything() + ) + }) + + describe('FileReader failure modes', () => { + const file = new File([new Uint8Array(16)], 'test.avif') + + it('resolves empty when the FileReader fires error', async () => { + vi.spyOn(console, 'error').mockImplementation(() => {}) + mockFileReaderError('readAsArrayBuffer') + expect(await getFromAvifFile(file)).toEqual({}) + }) + + it('resolves empty when the FileReader fires abort', async () => { + mockFileReaderAbort('readAsArrayBuffer') + expect(await getFromAvifFile(file)).toEqual({}) + }) + }) +}) + const setU32BE = (dv: DataView, off: number, val: number) => dv.setUint32(off, val, false) const setU16BE = (dv: DataView, off: number, val: number) => diff --git a/src/scripts/metadata/avif.ts b/src/scripts/metadata/avif.ts index c0d747d9e2..34664d0978 100644 --- a/src/scripts/metadata/avif.ts +++ b/src/scripts/metadata/avif.ts @@ -407,6 +407,7 @@ export function getFromAvifFile(file: File): Promise> { console.error('FileReader: Error reading AVIF file:', err) resolve({}) } + reader.onabort = () => resolve({}) reader.readAsArrayBuffer(file) }) } diff --git a/src/scripts/metadata/ebml.test.ts b/src/scripts/metadata/ebml.test.ts new file mode 100644 index 0000000000..bb66745c9c --- /dev/null +++ b/src/scripts/metadata/ebml.test.ts @@ -0,0 +1,49 @@ +import fs from 'fs' +import path from 'path' +import { afterEach, describe, expect, it, vi } from 'vitest' + +import { + EXPECTED_PROMPT, + EXPECTED_WORKFLOW, + mockFileReaderAbort, + mockFileReaderError +} from './__fixtures__/helpers' +import { getFromWebmFile } from './ebml' + +const fixturePath = path.resolve(__dirname, '__fixtures__/with_metadata.webm') + +describe('WebM/EBML metadata', () => { + it('extracts workflow and prompt from EBML SimpleTag elements', async () => { + const bytes = fs.readFileSync(fixturePath) + const file = new File([bytes], 'test.webm', { type: 'video/webm' }) + + const result = await getFromWebmFile(file) + + expect(result.workflow).toEqual(EXPECTED_WORKFLOW) + expect(result.prompt).toEqual(EXPECTED_PROMPT) + }) + + it('returns empty for non-WebM data', async () => { + const file = new File([new Uint8Array(16)], 'fake.webm') + + const result = await getFromWebmFile(file) + + expect(result).toEqual({}) + }) + + describe('FileReader failure modes', () => { + afterEach(() => vi.restoreAllMocks()) + + const file = new File([new Uint8Array(16)], 'test.webm') + + it('resolves empty when the FileReader fires error', async () => { + mockFileReaderError('readAsArrayBuffer') + expect(await getFromWebmFile(file)).toEqual({}) + }) + + it('resolves empty when the FileReader fires abort', async () => { + mockFileReaderAbort('readAsArrayBuffer') + expect(await getFromWebmFile(file)).toEqual({}) + }) + }) +}) diff --git a/src/scripts/metadata/ebml.ts b/src/scripts/metadata/ebml.ts index 835c2502d6..2e94c81737 100644 --- a/src/scripts/metadata/ebml.ts +++ b/src/scripts/metadata/ebml.ts @@ -353,6 +353,7 @@ export function getFromWebmFile(file: File): Promise { const reader = new FileReader() reader.onload = (event) => handleFileLoad(event, resolve) reader.onerror = () => resolve({}) + reader.onabort = () => resolve({}) reader.readAsArrayBuffer(file.slice(0, MAX_READ_BYTES)) }) } diff --git a/src/scripts/metadata/flac.test.ts b/src/scripts/metadata/flac.test.ts new file mode 100644 index 0000000000..c1e6bc5366 --- /dev/null +++ b/src/scripts/metadata/flac.test.ts @@ -0,0 +1,56 @@ +import fs from 'fs' +import path from 'path' +import { afterEach, describe, expect, it, vi } from 'vitest' + +import { + EXPECTED_PROMPT, + EXPECTED_WORKFLOW, + mockFileReaderAbort, + mockFileReaderError +} from './__fixtures__/helpers' +import { getFromFlacBuffer, getFromFlacFile } from './flac' + +const fixturePath = path.resolve(__dirname, '__fixtures__/with_metadata.flac') + +afterEach(() => vi.restoreAllMocks()) + +describe('FLAC metadata', () => { + it('extracts workflow and prompt from Vorbis comments', () => { + const bytes = fs.readFileSync(fixturePath) + const buffer = bytes.buffer.slice( + bytes.byteOffset, + bytes.byteOffset + bytes.byteLength + ) + + const result = getFromFlacBuffer(buffer) + + expect(result.workflow).toBe(JSON.stringify(EXPECTED_WORKFLOW)) + expect(result.prompt).toBe(JSON.stringify(EXPECTED_PROMPT)) + }) + + it('returns undefined for non-FLAC data', () => { + const buf = new ArrayBuffer(16) + const result = getFromFlacBuffer(buf) + expect(result).toBeUndefined() + }) + + describe('FileReader failure modes', () => { + const file = new File([new Uint8Array(16)], 'test.flac') + + it('resolves empty when the FileReader fires error', async () => { + mockFileReaderError('readAsArrayBuffer') + + const result = await getFromFlacFile(file) + + expect(result).toEqual({}) + }) + + it('resolves empty when the FileReader fires abort', async () => { + mockFileReaderAbort('readAsArrayBuffer') + + const result = await getFromFlacFile(file) + + expect(result).toEqual({}) + }) + }) +}) diff --git a/src/scripts/metadata/flac.ts b/src/scripts/metadata/flac.ts index 5a3efa6ac6..d2bfee8000 100644 --- a/src/scripts/metadata/flac.ts +++ b/src/scripts/metadata/flac.ts @@ -42,6 +42,8 @@ export function getFromFlacFile(file: File): Promise> { const arrayBuffer = event.target.result as ArrayBuffer r(getFromFlacBuffer(arrayBuffer)) } + reader.onerror = () => r({}) + reader.onabort = () => r({}) reader.readAsArrayBuffer(file) }) } diff --git a/src/scripts/metadata/gltf.test.ts b/src/scripts/metadata/gltf.test.ts index cead271807..230340baf8 100644 --- a/src/scripts/metadata/gltf.test.ts +++ b/src/scripts/metadata/gltf.test.ts @@ -1,7 +1,11 @@ -import { describe, expect, it } from 'vitest' +import { afterEach, describe, expect, it, vi } from 'vitest' import { ASCII, GltfSizeBytes } from '@/types/metadataTypes' +import { + mockFileReaderAbort, + mockFileReaderError +} from './__fixtures__/helpers' import { getGltfBinaryMetadata } from './gltf' describe('GLTF binary metadata parser', () => { @@ -160,4 +164,20 @@ describe('GLTF binary metadata parser', () => { const metadata = await getGltfBinaryMetadata(invalidEmptyFile) expect(metadata).toEqual({}) }) + + describe('FileReader failure modes', () => { + afterEach(() => vi.restoreAllMocks()) + + const file = new File([new Uint8Array(16)], 'test.glb') + + it('resolves empty when the FileReader fires error', async () => { + mockFileReaderError('readAsArrayBuffer') + expect(await getGltfBinaryMetadata(file)).toEqual({}) + }) + + it('resolves empty when the FileReader fires abort', async () => { + mockFileReaderAbort('readAsArrayBuffer') + expect(await getGltfBinaryMetadata(file)).toEqual({}) + }) + }) }) diff --git a/src/scripts/metadata/gltf.ts b/src/scripts/metadata/gltf.ts index 83a5c3c3e7..1da543cfcd 100644 --- a/src/scripts/metadata/gltf.ts +++ b/src/scripts/metadata/gltf.ts @@ -165,6 +165,7 @@ export function getGltfBinaryMetadata(file: File): Promise { } } reader.onerror = () => resolve({}) + reader.onabort = () => resolve({}) reader.readAsArrayBuffer(file.slice(0, bytesToRead)) }) } diff --git a/src/scripts/metadata/isobmff.test.ts b/src/scripts/metadata/isobmff.test.ts new file mode 100644 index 0000000000..7d342cc488 --- /dev/null +++ b/src/scripts/metadata/isobmff.test.ts @@ -0,0 +1,52 @@ +import fs from 'fs' +import path from 'path' +import { afterEach, describe, expect, it, vi } from 'vitest' + +import { + EXPECTED_PROMPT, + EXPECTED_WORKFLOW, + mockFileReaderAbort, + mockFileReaderError +} from './__fixtures__/helpers' +import { getFromIsobmffFile } from './isobmff' + +const fixturePath = path.resolve(__dirname, '__fixtures__/with_metadata.mp4') + +describe('ISOBMFF (MP4) metadata', () => { + it('extracts workflow and prompt from QuickTime keys/ilst boxes', async () => { + const bytes = fs.readFileSync(fixturePath) + const file = new File([bytes], 'test.mp4', { type: 'video/mp4' }) + + const result = await getFromIsobmffFile(file) + + expect(result.workflow).toEqual(EXPECTED_WORKFLOW) + expect(result.prompt).toEqual(EXPECTED_PROMPT) + }) + + it('returns empty for non-ISOBMFF data', async () => { + const file = new File([new Uint8Array(16)], 'fake.mp4', { + type: 'video/mp4' + }) + + const result = await getFromIsobmffFile(file) + + expect(result).toEqual({}) + }) + + describe('FileReader failure modes', () => { + afterEach(() => vi.restoreAllMocks()) + + const file = new File([new Uint8Array(16)], 'test.mp4') + + it('resolves empty when the FileReader fires error', async () => { + vi.spyOn(console, 'error').mockImplementation(() => {}) + mockFileReaderError('readAsArrayBuffer') + expect(await getFromIsobmffFile(file)).toEqual({}) + }) + + it('resolves empty when the FileReader fires abort', async () => { + mockFileReaderAbort('readAsArrayBuffer') + expect(await getFromIsobmffFile(file)).toEqual({}) + }) + }) +}) diff --git a/src/scripts/metadata/isobmff.ts b/src/scripts/metadata/isobmff.ts index e50a101699..95f263e4eb 100644 --- a/src/scripts/metadata/isobmff.ts +++ b/src/scripts/metadata/isobmff.ts @@ -274,6 +274,7 @@ export function getFromIsobmffFile(file: File): Promise { console.error('FileReader: Error reading ISOBMFF file:', err) resolve({}) } + reader.onabort = () => resolve({}) reader.readAsArrayBuffer(file.slice(0, MAX_READ_BYTES)) }) } diff --git a/src/scripts/metadata/json.test.ts b/src/scripts/metadata/json.test.ts new file mode 100644 index 0000000000..b48bf15a44 --- /dev/null +++ b/src/scripts/metadata/json.test.ts @@ -0,0 +1,91 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' + +import { + mockFileReaderAbort, + mockFileReaderError, + mockFileReaderResult +} from './__fixtures__/helpers' +import { getDataFromJSON } from './json' + +function jsonFile(content: object): File { + return new File([JSON.stringify(content)], 'test.json', { + type: 'application/json' + }) +} + +describe('getDataFromJSON', () => { + it('detects API-format workflows by class_type on every value', async () => { + const apiData = { + '1': { class_type: 'KSampler', inputs: {} }, + '2': { class_type: 'EmptyLatentImage', inputs: {} } + } + + const result = await getDataFromJSON(jsonFile(apiData)) + + expect(result).toEqual({ prompt: apiData }) + }) + + it('treats objects without universal class_type as a workflow', async () => { + const workflow = { nodes: [], links: [], version: 1 } + + const result = await getDataFromJSON(jsonFile(workflow)) + + expect(result).toEqual({ workflow }) + }) + + it('extracts templates when the root object has a templates key', async () => { + const templates = [{ name: 'basic' }] + + const result = await getDataFromJSON(jsonFile({ templates })) + + expect(result).toEqual({ templates }) + }) + + it('returns undefined for non-JSON content', async () => { + const file = new File(['not valid json'], 'bad.json', { + type: 'application/json' + }) + + const result = await getDataFromJSON(file) + + expect(result).toBeUndefined() + }) + + describe('FileReader failure modes', () => { + afterEach(() => { + vi.restoreAllMocks() + }) + + it('resolves undefined when the FileReader fires error', async () => { + mockFileReaderError('readAsText') + + const result = await getDataFromJSON(jsonFile({ nodes: [] })) + + expect(result).toBeUndefined() + }) + + it('resolves undefined when the FileReader fires abort', async () => { + mockFileReaderAbort('readAsText') + + const result = await getDataFromJSON(jsonFile({ nodes: [] })) + + expect(result).toBeUndefined() + }) + + it('resolves undefined when reader.result is not a string', async () => { + mockFileReaderResult('readAsText', new ArrayBuffer(8)) + + const result = await getDataFromJSON(jsonFile({ nodes: [] })) + + expect(result).toBeUndefined() + }) + + it('resolves undefined when reader.result is null', async () => { + mockFileReaderResult('readAsText', null) + + const result = await getDataFromJSON(jsonFile({ nodes: [] })) + + expect(result).toBeUndefined() + }) + }) +}) diff --git a/src/scripts/metadata/json.ts b/src/scripts/metadata/json.ts index 873b42117b..ee7fd1ad8d 100644 --- a/src/scripts/metadata/json.ts +++ b/src/scripts/metadata/json.ts @@ -6,21 +6,28 @@ export function getDataFromJSON( return new Promise | undefined>((resolve) => { const reader = new FileReader() reader.onload = async () => { - const readerResult = reader.result as string - const jsonContent = JSON.parse(readerResult) - if (jsonContent?.templates) { - resolve({ templates: jsonContent.templates }) - return + try { + if (typeof reader.result !== 'string') { + resolve(undefined) + return + } + const jsonContent = JSON.parse(reader.result) + if (jsonContent?.templates) { + resolve({ templates: jsonContent.templates }) + return + } + if (isApiJson(jsonContent)) { + resolve({ prompt: jsonContent }) + return + } + resolve({ workflow: jsonContent }) + } catch { + resolve(undefined) } - if (isApiJson(jsonContent)) { - resolve({ prompt: jsonContent }) - return - } - resolve({ workflow: jsonContent }) - return } + reader.onerror = () => resolve(undefined) + reader.onabort = () => resolve(undefined) reader.readAsText(file) - return }) } diff --git a/src/scripts/metadata/mp3.test.ts b/src/scripts/metadata/mp3.test.ts new file mode 100644 index 0000000000..c2fba7dbd5 --- /dev/null +++ b/src/scripts/metadata/mp3.test.ts @@ -0,0 +1,106 @@ +import fs from 'fs' +import path from 'path' +import { afterEach, describe, expect, it, vi } from 'vitest' + +import { + EXPECTED_PROMPT, + EXPECTED_WORKFLOW, + mockFileReaderAbort, + mockFileReaderError +} from './__fixtures__/helpers' +import { getMp3Metadata } from './mp3' + +const fixturePath = path.resolve(__dirname, '__fixtures__/with_metadata.mp3') + +afterEach(() => vi.restoreAllMocks()) + +describe('MP3 metadata', () => { + it('extracts workflow and prompt from ID3 tags', async () => { + const bytes = fs.readFileSync(fixturePath) + const file = new File([bytes], 'test.mp3', { type: 'audio/mpeg' }) + + const result = await getMp3Metadata(file) + + expect(result.workflow).toEqual(EXPECTED_WORKFLOW) + expect(result.prompt).toEqual(EXPECTED_PROMPT) + }) + + it('returns undefined fields when file has no embedded metadata', async () => { + vi.spyOn(console, 'error').mockImplementation(() => {}) + const file = new File([new Uint8Array(16)], 'empty.mp3', { + type: 'audio/mpeg' + }) + + const result = await getMp3Metadata(file) + + expect(result.workflow).toBeUndefined() + expect(result.prompt).toBeUndefined() + expect(console.error).toHaveBeenCalledWith('Invalid file signature.') + }) + + it('does not log an invalid signature for a valid MP3 sync header', async () => { + const errorSpy = vi.spyOn(console, 'error').mockImplementation(() => {}) + const buf = new Uint8Array(16) + buf[0] = 0xff + buf[1] = 0xfb + const file = new File([buf], 'valid.mp3', { type: 'audio/mpeg' }) + + await getMp3Metadata(file) + + expect(errorSpy).not.toHaveBeenCalled() + }) + + it('does not log an invalid signature for a valid ID3v2 header', async () => { + const errorSpy = vi.spyOn(console, 'error').mockImplementation(() => {}) + const buf = new Uint8Array(16) + buf[0] = 0x49 + buf[1] = 0x44 + buf[2] = 0x33 + const file = new File([buf], 'valid-id3.mp3', { type: 'audio/mpeg' }) + + await getMp3Metadata(file) + + expect(errorSpy).not.toHaveBeenCalled() + }) + + it('extracts metadata that spans the 4096-byte page boundary', async () => { + vi.spyOn(console, 'error').mockImplementation(() => {}) + const metadata = + `prompt\0${JSON.stringify(EXPECTED_PROMPT)}\0` + + `workflow\0${JSON.stringify(EXPECTED_WORKFLOW)}\0` + const metadataStart = 4090 + const size = metadataStart + metadata.length + 4 + const buf = new Uint8Array(size) + for (let i = 0; i < metadata.length; i++) { + buf[metadataStart + i] = metadata.charCodeAt(i) + } + buf[size - 2] = 0xff + buf[size - 1] = 0xfb + const file = new File([buf], 'large.mp3', { type: 'audio/mpeg' }) + + const result = await getMp3Metadata(file) + + expect(result.workflow).toEqual(EXPECTED_WORKFLOW) + expect(result.prompt).toEqual(EXPECTED_PROMPT) + }) + + describe('FileReader failure modes', () => { + const file = new File([new Uint8Array(16)], 'test.mp3') + + it('resolves undefined fields when the FileReader fires error', async () => { + mockFileReaderError('readAsArrayBuffer') + + const result = await getMp3Metadata(file) + + expect(result).toEqual({ prompt: undefined, workflow: undefined }) + }) + + it('resolves undefined fields when the FileReader fires abort', async () => { + mockFileReaderAbort('readAsArrayBuffer') + + const result = await getMp3Metadata(file) + + expect(result).toEqual({ prompt: undefined, workflow: undefined }) + }) + }) +}) diff --git a/src/scripts/metadata/mp3.ts b/src/scripts/metadata/mp3.ts index f3e5e23bbd..8ea3cc1f16 100644 --- a/src/scripts/metadata/mp3.ts +++ b/src/scripts/metadata/mp3.ts @@ -1,21 +1,28 @@ export async function getMp3Metadata(file: File) { const reader = new FileReader() - const read_process = new Promise( - (r) => (reader.onload = (event) => r(event?.target?.result)) - ) + const read_process = new Promise((r) => { + reader.onload = (event) => r((event?.target?.result as ArrayBuffer) ?? null) + reader.onerror = () => r(null) + reader.onabort = () => r(null) + }) reader.readAsArrayBuffer(file) - const arrayBuffer = (await read_process) as ArrayBuffer + const arrayBuffer = await read_process + if (!arrayBuffer) return { prompt: undefined, workflow: undefined } //https://stackoverflow.com/questions/7302439/how-can-i-determine-that-a-particular-file-is-in-fact-an-mp3-file#7302482 const sig_bytes = new Uint8Array(arrayBuffer, 0, 3) if ( - (sig_bytes[0] != 0xff && sig_bytes[1] != 0xfb) || - (sig_bytes[0] != 0x49 && sig_bytes[1] != 0x44 && sig_bytes[2] != 0x33) + (sig_bytes[0] != 0xff || sig_bytes[1] != 0xfb) && + (sig_bytes[0] != 0x49 || sig_bytes[1] != 0x44 || sig_bytes[2] != 0x33) ) console.error('Invalid file signature.') let header = '' while (header.length < arrayBuffer.byteLength) { const page = String.fromCharCode( - ...new Uint8Array(arrayBuffer, header.length, header.length + 4096) + ...new Uint8Array( + arrayBuffer, + header.length, + Math.min(4096, arrayBuffer.byteLength - header.length) + ) ) header += page if (page.match('\u00ff\u00fb')) break diff --git a/src/scripts/metadata/ogg.test.ts b/src/scripts/metadata/ogg.test.ts new file mode 100644 index 0000000000..13c3a4bea0 --- /dev/null +++ b/src/scripts/metadata/ogg.test.ts @@ -0,0 +1,74 @@ +import fs from 'fs' +import path from 'path' +import { afterEach, describe, expect, it, vi } from 'vitest' + +import { + EXPECTED_PROMPT, + EXPECTED_WORKFLOW, + mockFileReaderAbort, + mockFileReaderError +} from './__fixtures__/helpers' +import { getOggMetadata } from './ogg' + +const fixturePath = path.resolve(__dirname, '__fixtures__/with_metadata.opus') + +afterEach(() => vi.restoreAllMocks()) + +describe('OGG/Opus metadata', () => { + it('extracts workflow and prompt from an Opus file', async () => { + const bytes = fs.readFileSync(fixturePath) + const file = new File([bytes], 'test.opus', { type: 'audio/ogg' }) + + const result = await getOggMetadata(file) + + expect(result.workflow).toEqual(EXPECTED_WORKFLOW) + expect(result.prompt).toEqual(EXPECTED_PROMPT) + }) + + it('returns undefined fields for non-OGG data', async () => { + vi.spyOn(console, 'error').mockImplementation(() => {}) + const file = new File([new Uint8Array(16)], 'fake.ogg', { + type: 'audio/ogg' + }) + + const result = await getOggMetadata(file) + + expect(result.workflow).toBeUndefined() + expect(result.prompt).toBeUndefined() + expect(console.error).toHaveBeenCalledWith('Invalid file signature.') + }) + + it('handles files larger than 4096 bytes without RangeError', async () => { + const size = 5000 + const buf = new Uint8Array(size) + const oggs = new TextEncoder().encode('OggS\0') + buf.set(oggs, 0) + buf.set(oggs, 4500) + const file = new File([buf], 'large.ogg', { type: 'audio/ogg' }) + + const result = await getOggMetadata(file) + + expect(result.workflow).toBeUndefined() + expect(result.prompt).toBeUndefined() + }) + + describe('FileReader failure modes', () => { + const file = new File([new Uint8Array(16)], 'test.ogg') + + it('resolves undefined fields when the FileReader fires error', async () => { + mockFileReaderError('readAsArrayBuffer') + + const result = await getOggMetadata(file) + + expect(result).toEqual({ prompt: undefined, workflow: undefined }) + }) + + it('resolves undefined fields when the FileReader fires abort', async () => { + mockFileReaderAbort('readAsArrayBuffer') + + const result = await getOggMetadata(file) + + expect(result).toEqual({ prompt: undefined, workflow: undefined }) + }) + }) +}) diff --git a/src/scripts/metadata/ogg.ts b/src/scripts/metadata/ogg.ts index 5dc49c02b3..8adec350a7 100644 --- a/src/scripts/metadata/ogg.ts +++ b/src/scripts/metadata/ogg.ts @@ -1,17 +1,24 @@ export async function getOggMetadata(file: File) { const reader = new FileReader() - const read_process = new Promise( - (r) => (reader.onload = (event) => r(event?.target?.result)) - ) + const read_process = new Promise((r) => { + reader.onload = (event) => r((event?.target?.result as ArrayBuffer) ?? null) + reader.onerror = () => r(null) + reader.onabort = () => r(null) + }) reader.readAsArrayBuffer(file) - const arrayBuffer = (await read_process) as ArrayBuffer + const arrayBuffer = await read_process + if (!arrayBuffer) return { prompt: undefined, workflow: undefined } const signature = String.fromCharCode(...new Uint8Array(arrayBuffer, 0, 4)) if (signature !== 'OggS') console.error('Invalid file signature.') let oggs = 0 let header = '' while (header.length < arrayBuffer.byteLength) { const page = String.fromCharCode( - ...new Uint8Array(arrayBuffer, header.length, header.length + 4096) + ...new Uint8Array( + arrayBuffer, + header.length, + Math.min(4096, arrayBuffer.byteLength - header.length) + ) ) if (page.match('OggS\u0000')) oggs++ header += page diff --git a/src/scripts/metadata/png.test.ts b/src/scripts/metadata/png.test.ts index dcc47bf049..fac4eaeefd 100644 --- a/src/scripts/metadata/png.test.ts +++ b/src/scripts/metadata/png.test.ts @@ -1,11 +1,19 @@ -import { describe, expect, it } from 'vitest' +import { afterEach, describe, expect, it, vi } from 'vitest' -import { getFromPngBuffer } from './png' +import { + mockFileReaderAbort, + mockFileReaderError +} from './__fixtures__/helpers' +import { getFromPngBuffer, getFromPngFile } from './png' + +afterEach(() => vi.restoreAllMocks()) + +const PNG_SIGNATURE = [0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a] function createPngWithChunk( chunkType: string, keyword: string, - content: string, + content: string | Uint8Array, options: { compressionFlag?: number compressionMethod?: number @@ -20,12 +28,11 @@ function createPngWithChunk( translatedKeyword = '' } = options - const signature = new Uint8Array([ - 0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a - ]) + const signature = new Uint8Array(PNG_SIGNATURE) const typeBytes = new TextEncoder().encode(chunkType) const keywordBytes = new TextEncoder().encode(keyword) - const contentBytes = new TextEncoder().encode(content) + const contentBytes = + content instanceof Uint8Array ? content : new TextEncoder().encode(content) let chunkData: Uint8Array if (chunkType === 'iTXt') { @@ -66,12 +73,11 @@ function createPngWithChunk( new DataView(lengthBytes.buffer).setUint32(0, chunkData.length, false) const crc = new Uint8Array(4) - const iendType = new TextEncoder().encode('IEND') const iendLength = new Uint8Array(4) const iendCrc = new Uint8Array(4) - const total = signature.length + 4 + 4 + chunkData.length + 4 + 4 + 4 + 0 + 4 + const total = signature.length + (4 + 4 + chunkData.length + 4) + (4 + 4 + 4) const result = new Uint8Array(total) let offset = 0 @@ -138,6 +144,21 @@ describe('getFromPngBuffer', () => { expect(result['workflow']).toBe(workflow) }) + it('logs warning and skips iTXt chunk with unsupported compression method', async () => { + vi.spyOn(console, 'warn').mockImplementation(() => {}) + const buffer = createPngWithChunk('iTXt', 'workflow', 'data', { + compressionFlag: 1, + compressionMethod: 99 + }) + + const result = await getFromPngBuffer(buffer) + + expect(result['workflow']).toBeUndefined() + expect(console.warn).toHaveBeenCalledWith( + expect.stringContaining('Unsupported compression method 99') + ) + }) + it('parses compressed iTXt chunk', async () => { const workflow = '{"nodes":[{"id":1,"type":"KSampler"}]}' const contentBytes = new TextEncoder().encode(workflow) @@ -163,83 +184,49 @@ describe('getFromPngBuffer', () => { pos += chunk.length } - const buffer = createPngWithCompressedITXt( - 'workflow', - compressedBytes, - '', - '' - ) + const buffer = createPngWithChunk('iTXt', 'workflow', compressedBytes, { + compressionFlag: 1, + compressionMethod: 0 + }) const result = await getFromPngBuffer(buffer) expect(result['workflow']).toBe(workflow) }) }) -function createPngWithCompressedITXt( - keyword: string, - compressedContent: Uint8Array, - languageTag: string, - translatedKeyword: string -): ArrayBuffer { - const signature = new Uint8Array([ - 0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a - ]) - const typeBytes = new TextEncoder().encode('iTXt') - const keywordBytes = new TextEncoder().encode(keyword) - const langBytes = new TextEncoder().encode(languageTag) - const transBytes = new TextEncoder().encode(translatedKeyword) +describe('getFromPngFile', () => { + it('reads metadata from a File object', async () => { + const workflow = '{"nodes":[]}' + const buffer = createPngWithChunk('tEXt', 'workflow', workflow) + const file = new File([buffer], 'test.png', { type: 'image/png' }) - const totalLength = - keywordBytes.length + - 1 + - 2 + - langBytes.length + - 1 + - transBytes.length + - 1 + - compressedContent.length + const result = await getFromPngFile(file) - const chunkData = new Uint8Array(totalLength) - let pos = 0 - chunkData.set(keywordBytes, pos) - pos += keywordBytes.length - chunkData[pos++] = 0 - chunkData[pos++] = 1 - chunkData[pos++] = 0 - chunkData.set(langBytes, pos) - pos += langBytes.length - chunkData[pos++] = 0 - chunkData.set(transBytes, pos) - pos += transBytes.length - chunkData[pos++] = 0 - chunkData.set(compressedContent, pos) + expect(result['workflow']).toBe(workflow) + }) - const lengthBytes = new Uint8Array(4) - new DataView(lengthBytes.buffer).setUint32(0, chunkData.length, false) + it('returns empty for an invalid PNG File', async () => { + vi.spyOn(console, 'error').mockImplementation(() => {}) + const file = new File([new ArrayBuffer(8)], 'bad.png', { + type: 'image/png' + }) - const crc = new Uint8Array(4) - const iendType = new TextEncoder().encode('IEND') - const iendLength = new Uint8Array(4) - const iendCrc = new Uint8Array(4) + const result = await getFromPngFile(file) - const total = signature.length + 4 + 4 + chunkData.length + 4 + 4 + 4 + 0 + 4 - const result = new Uint8Array(total) + expect(result).toEqual({}) + expect(console.error).toHaveBeenCalledWith('Not a valid PNG file') + }) - let offset = 0 - result.set(signature, offset) - offset += signature.length - result.set(lengthBytes, offset) - offset += 4 - result.set(typeBytes, offset) - offset += 4 - result.set(chunkData, offset) - offset += chunkData.length - result.set(crc, offset) - offset += 4 - result.set(iendLength, offset) - offset += 4 - result.set(iendType, offset) - offset += 4 - result.set(iendCrc, offset) + describe('FileReader failure modes', () => { + const file = new File([new Uint8Array(16)], 'test.png') - return result.buffer -} + it('rejects when the FileReader fires error', async () => { + mockFileReaderError('readAsArrayBuffer') + await expect(getFromPngFile(file)).rejects.toBeDefined() + }) + + it('rejects when the FileReader fires abort', async () => { + mockFileReaderAbort('readAsArrayBuffer') + await expect(getFromPngFile(file)).rejects.toThrow('FileReader aborted') + }) + }) +}) diff --git a/src/scripts/metadata/png.ts b/src/scripts/metadata/png.ts index 2d47effcb2..30add0b2e8 100644 --- a/src/scripts/metadata/png.ts +++ b/src/scripts/metadata/png.ts @@ -126,6 +126,7 @@ export async function getFromPngFile( resolve(result) } reader.onerror = () => reject(reader.error) + reader.onabort = () => reject(new Error('FileReader aborted')) reader.readAsArrayBuffer(file) }) } diff --git a/src/scripts/metadata/svg.test.ts b/src/scripts/metadata/svg.test.ts new file mode 100644 index 0000000000..dcc691c8ec --- /dev/null +++ b/src/scripts/metadata/svg.test.ts @@ -0,0 +1,42 @@ +import { describe, expect, it } from 'vitest' + +import { getSvgMetadata } from './svg' + +function svgFile(content: string): File { + return new File([content], 'test.svg', { type: 'image/svg+xml' }) +} + +describe('getSvgMetadata', () => { + it('extracts workflow and prompt from CDATA in ', async () => { + const svg = ` + ${JSON.stringify({ + workflow: { nodes: [] }, + prompt: { '1': {} } + })} + + ` + + const result = await getSvgMetadata(svgFile(svg)) + + expect(result).toEqual({ + workflow: { nodes: [] }, + prompt: { '1': {} } + }) + }) + + it('returns empty when SVG has no metadata element', async () => { + const svg = '' + + const result = await getSvgMetadata(svgFile(svg)) + + expect(result).toEqual({}) + }) + + it('returns empty when CDATA contains invalid JSON', async () => { + const svg = `not valid json` + + const result = await getSvgMetadata(svgFile(svg)) + + expect(result).toEqual({}) + }) +}) diff --git a/src/scripts/pnginfo.test.ts b/src/scripts/pnginfo.test.ts index 21c68a369b..c10aaa4b1a 100644 --- a/src/scripts/pnginfo.test.ts +++ b/src/scripts/pnginfo.test.ts @@ -1,4 +1,6 @@ -import { describe, expect, it, vi } from 'vitest' +import fs from 'fs' +import path from 'path' +import { afterEach, describe, expect, it, vi } from 'vitest' import { getFromAvifFile } from './metadata/avif' import { getFromFlacFile } from './metadata/flac' @@ -21,67 +23,183 @@ vi.mock('./metadata/avif', () => ({ getFromAvifFile: vi.fn() })) -function buildExifPayload(workflowJson: string): Uint8Array { - const fullStr = `workflow:${workflowJson}\0` - const strBytes = new TextEncoder().encode(fullStr) +afterEach(() => vi.restoreAllMocks()) - const headerSize = 22 - const buf = new Uint8Array(headerSize + strBytes.length) +const fixturesDir = path.resolve(__dirname, 'metadata/__fixtures__') + +type AsciiIfdEntry = { tag: number; value: string } + +function encodeAsciiIfd(entries: AsciiIfdEntry[]): Uint8Array { + const tableSize = 10 + 12 * entries.length + const strings = entries.map((e) => new TextEncoder().encode(`${e.value}\0`)) + const totalStringBytes = strings.reduce((sum, s) => sum + s.length, 0) + + const buf = new Uint8Array(tableSize + totalStringBytes) const dv = new DataView(buf.buffer) buf.set([0x49, 0x49], 0) dv.setUint16(2, 0x002a, true) dv.setUint32(4, 8, true) - dv.setUint16(8, 1, true) - dv.setUint16(10, 0, true) - dv.setUint16(12, 2, true) - dv.setUint32(14, strBytes.length, true) - dv.setUint32(18, 22, true) - buf.set(strBytes, 22) + dv.setUint16(8, entries.length, true) + + let stringOffset = tableSize + for (let i = 0; i < entries.length; i++) { + const entryOffset = 10 + i * 12 + dv.setUint16(entryOffset, entries[i].tag, true) + dv.setUint16(entryOffset + 2, 2, true) + dv.setUint32(entryOffset + 4, strings[i].length, true) + dv.setUint32(entryOffset + 8, stringOffset, true) + buf.set(strings[i], stringOffset) + stringOffset += strings[i].length + } return buf } -function buildWebp(precedingChunkLength: number, workflowJson: string): File { - const exifPayload = buildExifPayload(workflowJson) - const precedingPadded = precedingChunkLength + (precedingChunkLength % 2) - const totalSize = 12 + (8 + precedingPadded) + (8 + exifPayload.length) +type WebpChunk = { type: string; payload: Uint8Array } - const buffer = new Uint8Array(totalSize) - const dv = new DataView(buffer.buffer) +function wrapInWebp(chunks: WebpChunk[]): File { + let payloadSize = 0 + for (const c of chunks) { + payloadSize += 8 + c.payload.length + (c.payload.length % 2) + } + const totalSize = 12 + payloadSize + const buf = new Uint8Array(totalSize) + const dv = new DataView(buf.buffer) - buffer.set([0x52, 0x49, 0x46, 0x46], 0) + buf.set([0x52, 0x49, 0x46, 0x46], 0) dv.setUint32(4, totalSize - 8, true) - buffer.set([0x57, 0x45, 0x42, 0x50], 8) + buf.set([0x57, 0x45, 0x42, 0x50], 8) - buffer.set([0x56, 0x50, 0x38, 0x20], 12) - dv.setUint32(16, precedingChunkLength, true) + let offset = 12 + for (const c of chunks) { + for (let i = 0; i < 4; i++) { + buf[offset + i] = c.type.charCodeAt(i) + } + dv.setUint32(offset + 4, c.payload.length, true) + buf.set(c.payload, offset + 8) + offset += 8 + c.payload.length + (c.payload.length % 2) + } - const exifStart = 20 + precedingPadded - buffer.set([0x45, 0x58, 0x49, 0x46], exifStart) - dv.setUint32(exifStart + 4, exifPayload.length, true) - buffer.set(exifPayload, exifStart + 8) + return new File([buf], 'test.webp', { type: 'image/webp' }) +} - return new File([buffer], 'test.webp', { type: 'image/webp' }) +function exifChunk( + entries: AsciiIfdEntry[], + options: { withExifPrefix?: boolean } = {} +): WebpChunk { + const ifd = encodeAsciiIfd(entries) + if (!options.withExifPrefix) { + return { type: 'EXIF', payload: ifd } + } + const prefixed = new Uint8Array(6 + ifd.length) + prefixed.set(new TextEncoder().encode('Exif\0\0'), 0) + prefixed.set(ifd, 6) + return { type: 'EXIF', payload: prefixed } } describe('getWebpMetadata', () => { - it('finds workflow when a preceding chunk has odd length (RIFF padding)', async () => { - const workflow = '{"nodes":[]}' - const file = buildWebp(3, workflow) + it('returns empty when the file is not a valid WEBP', async () => { + vi.spyOn(console, 'error').mockImplementation(() => {}) + const file = new File([new Uint8Array(12)], 'fake.webp') const metadata = await getWebpMetadata(file) - expect(metadata.workflow).toBe(workflow) + expect(metadata).toEqual({}) + expect(console.error).toHaveBeenCalledWith('Not a valid WEBP file') }) - it('finds workflow when preceding chunk has even length (no padding)', async () => { - const workflow = '{"nodes":[1]}' - const file = buildWebp(4, workflow) + it('returns empty when a valid WEBP has no EXIF chunk', async () => { + const file = wrapInWebp([ + { type: 'VP8 ', payload: new Uint8Array([0, 0, 0, 0]) } + ]) const metadata = await getWebpMetadata(file) - expect(metadata.workflow).toBe(workflow) + expect(metadata).toEqual({}) + }) + + it('extracts workflow and prompt from EXIF without prefix', async () => { + const bytes = fs.readFileSync(path.join(fixturesDir, 'with_metadata.webp')) + const file = new File([bytes], 'test.webp', { type: 'image/webp' }) + + const metadata = await getWebpMetadata(file) + + expect(metadata).toEqual({ + workflow: + '{"nodes":[{"id":1,"type":"KSampler","pos":[100,100],"size":[200,200]}]}', + prompt: '{"1":{"class_type":"KSampler","inputs":{}}}' + }) + }) + + it('extracts workflow and prompt from EXIF with Exif\\0\\0 prefix', async () => { + const bytes = fs.readFileSync( + path.join(fixturesDir, 'with_metadata_exif_prefix.webp') + ) + const file = new File([bytes], 'test.webp', { type: 'image/webp' }) + + const metadata = await getWebpMetadata(file) + + expect(metadata).toEqual({ + workflow: + '{"nodes":[{"id":1,"type":"KSampler","pos":[100,100],"size":[200,200]}]}', + prompt: '{"1":{"class_type":"KSampler","inputs":{}}}' + }) + }) + + it('walks past odd-length preceding chunks (RIFF padding)', async () => { + const file = wrapInWebp([ + { type: 'VP8 ', payload: new Uint8Array(3) }, + exifChunk([{ tag: 0, value: 'workflow:{"a":1}' }]) + ]) + + const metadata = await getWebpMetadata(file) + + expect(metadata).toEqual({ workflow: '{"a":1}' }) + }) +}) + +describe('getLatentMetadata', () => { + function buildSafetensors(headerObj: object): File { + const headerBytes = new TextEncoder().encode(JSON.stringify(headerObj)) + const buf = new Uint8Array(8 + headerBytes.length) + const dv = new DataView(buf.buffer) + dv.setUint32(0, headerBytes.length, true) + dv.setUint32(4, 0, true) + buf.set(headerBytes, 8) + return new File([buf], 'test.safetensors') + } + + it('extracts __metadata__ from a safetensors header', async () => { + const workflow = + '{"nodes":[{"id":1,"type":"KSampler","pos":[100,100],"size":[200,200]}]}' + const prompt = '{"1":{"class_type":"KSampler","inputs":{}}}' + const file = buildSafetensors({ + __metadata__: { workflow, prompt }, + 'tensor.weight': { dtype: 'F32', shape: [1], data_offsets: [0, 4] } + }) + + const metadata = await getLatentMetadata(file) + + expect(metadata).toEqual({ workflow, prompt }) + }) + + it('returns undefined when the safetensors header has no __metadata__', async () => { + const file = buildSafetensors({ + 'tensor.weight': { dtype: 'F32', shape: [1], data_offsets: [0, 4] } + }) + + const metadata = await getLatentMetadata(file) + + expect(metadata).toBeUndefined() + }) + + it('returns undefined for a truncated or malformed file', async () => { + const file = new File([new Uint8Array(4)], 'bad.safetensors') + + const metadata = await getLatentMetadata(file) + + expect(metadata).toBeUndefined() }) }) @@ -116,37 +234,3 @@ describe('format-specific metadata wrappers', () => { expect(result).toEqual({ workflow: '{"avif":1}' }) }) }) - -const buildSafetensors = (header: Record): File => { - const headerJson = JSON.stringify(header) - const headerBytes = new TextEncoder().encode(headerJson) - const buf = new ArrayBuffer(8 + headerBytes.length) - const dv = new DataView(buf) - dv.setUint32(0, headerBytes.length, true) - dv.setUint32(4, 0, true) - new Uint8Array(buf, 8).set(headerBytes) - return new File([buf], 'x.safetensors') -} - -describe('getLatentMetadata', () => { - it('returns the __metadata__ object from a safetensors header', async () => { - const file = buildSafetensors({ - __metadata__: { workflow: '{"nodes":[]}', extra: 'value' }, - 'tensor.weight': { dtype: 'F32', shape: [1], data_offsets: [0, 4] } - }) - - const result = await getLatentMetadata(file) - - expect(result).toEqual({ workflow: '{"nodes":[]}', extra: 'value' }) - }) - - it('resolves undefined when header has no __metadata__ entry', async () => { - const file = buildSafetensors({ - 'tensor.weight': { dtype: 'F32', shape: [1], data_offsets: [0, 4] } - }) - - const result = await getLatentMetadata(file) - - expect(result).toBeUndefined() - }) -}) diff --git a/src/scripts/pnginfo.ts b/src/scripts/pnginfo.ts index d0b6ee3e4d..c27ee50528 100644 --- a/src/scripts/pnginfo.ts +++ b/src/scripts/pnginfo.ts @@ -105,14 +105,17 @@ export function getWebpMetadata(file: File) { ...webp.slice(offset, offset + 4) ) if (chunk_type === 'EXIF') { + let exifOffset = offset + 8 + let exifLength = chunk_length if ( - String.fromCharCode(...webp.slice(offset + 8, offset + 8 + 6)) == + String.fromCharCode(...webp.slice(exifOffset, exifOffset + 6)) == 'Exif\0\0' ) { - offset += 6 + exifOffset += 6 + exifLength -= 6 } - let data = parseExifData( - webp.slice(offset + 8, offset + 8 + chunk_length) + const data = parseExifData( + webp.slice(exifOffset, exifOffset + exifLength) ) for (const key in data) { const value = data[Number(key)] @@ -131,30 +134,38 @@ export function getWebpMetadata(file: File) { r(txt_chunks) } - + reader.onerror = () => r({}) + reader.onabort = () => r({}) reader.readAsArrayBuffer(file) }) } -export function getLatentMetadata(file: File): Promise> { +export function getLatentMetadata( + file: File +): Promise | undefined> { return new Promise((r) => { const reader = new FileReader() reader.onload = (event) => { - const safetensorsData = new Uint8Array( - event.target?.result as ArrayBuffer - ) - const dataView = new DataView(safetensorsData.buffer) - let header_size = dataView.getUint32(0, true) - let offset = 8 - let header = JSON.parse( - new TextDecoder().decode( - safetensorsData.slice(offset, offset + header_size) + try { + const safetensorsData = new Uint8Array( + event.target?.result as ArrayBuffer ) - ) - r(header.__metadata__) + const dataView = new DataView(safetensorsData.buffer) + const headerSize = dataView.getUint32(0, true) + const offset = 8 + const header = JSON.parse( + new TextDecoder().decode( + safetensorsData.slice(offset, offset + headerSize) + ) + ) + r(header.__metadata__) + } catch { + r(undefined) + } } - - var slice = file.slice(0, 1024 * 1024 * 4) + reader.onerror = () => r(undefined) + reader.onabort = () => r(undefined) + const slice = file.slice(0, 1024 * 1024 * 4) reader.readAsArrayBuffer(slice) }) }