PEE/src/processor.worker.ts

/// <reference lib="ES2021" />
/// <reference lib="webworker" />

import * as platform from './platform';
import pngv3 from "./pngv3";
//import webm from "./webm";
//import gif from "./gif";
import jpg from "./jpg";
import thirdeye from "./thirdeye";
import pomf from "./pomf";

console.log("Worker started");

const pendinggens: Record<number, ((res: IteratorResult<any>) => void)[]> = {};
const pendingcmds: Record<number, ((res: any) => void)[]> = {};

const proxyAsyncGen = (id: number) => {
    return {
        next(arg) {
            postMessage({
                type: 'ag',
                id
            });
            return new Promise<IteratorResult<any>>(res => {
                if (!pendinggens[id])
                    pendinggens[id] = [];
                pendinggens[id].push(res);
            });
        },

        return(v) {
            postMessage({
                type: 'ag',
                id
            });
            return new Promise<IteratorResult<any>>(res => {
                if (!pendinggens[id])
                    pendinggens[id] = [];
                pendinggens[id].push(res);
            });
        },

        throw() {
            postMessage({
                type: 'ag',
                id
            });
            return new Promise<IteratorResult<any>>(res => {
                if (!pendinggens[id])
                    pendinggens[id] = [];
                pendinggens[id].push(res);
            });
        },

        [Symbol.asyncIterator]() {
            return this as AsyncGenerator;
        }
    } as AsyncGenerator;
};

const deserializeMessage = (m: any) => {
    if (typeof m == "object" && m.type == 'AsyncGenerator') {
        return proxyAsyncGen(m.id);
    }
    if (typeof m == "object") {
        for (const p in m) {
            m[p] = deserializeMessage(m[p]);
        }
    }
    return m;
};

type WorkerEmbeddedFileWithPreview = {
    page?: { title: string, url: string }; // can be a booru page
    source?: string; // can be like a twitter post this was posted in originally
    thumbnail: string | Uint8Array;
    filename: string;
    data: WorkerEmbeddedFileWithoutPreview['data'] | { url: string, headers?: Record<string, string> };
};

type WorkerEmbeddedFileWithoutPreview = {
    page: undefined;
    source: undefined;
    thumbnail?: string;
    filename: string;
    data: string | Uint8Array;
};

export type WorkerEmbeddedFile = WorkerEmbeddedFileWithPreview | WorkerEmbeddedFileWithoutPreview;

export interface ImageProcessor {
    skip?: true;
    match(fn: string): boolean;
    has_embed(b: Uint8Array, fn?: string, prevurl?: string): boolean | string | undefined | Promise<boolean | string | undefined>;
    extract(b: Uint8Array, fn?: string): WorkerEmbeddedFile[] | Promise<WorkerEmbeddedFile[]>;
    inject?(b: File, c: string[]): Buffer | Promise<Buffer>;
}

const processors: ImageProcessor[] =
    [thirdeye, pomf, pngv3, jpg]; //, webm, gif

const processImage = async (srcs: AsyncGenerator<string, void, void>, fn: string, hex: string, prevurl: string) => {
    const ret = await Promise.all(processors.filter(e => e.match(fn)).map(async proc => {
        if (proc.skip) {
            // skip file downloading, file is referenced from the filename
            // basically does things like filtering out blacklisted tags
            const md5 = Buffer.from(hex, 'base64');
            if (await proc.has_embed(md5, fn, prevurl) === true) {
                return [await proc.extract(md5, fn), true] as [WorkerEmbeddedFile[], boolean];
            }
            return;
        }
        let succ = false;
        let cumul: Buffer;
        do {
            try {
                const n = await srcs.next();
                if (n.done)
                    return; // no more links to try
                const iter = platform.streamRemote(n.value);
                if (!iter)
                    return;
                cumul = Buffer.alloc(0);
                let found: boolean | undefined;
                let chunk: ReadableStreamDefaultReadResult<Buffer> = { done: true };
                do {
                    const { value, done } = await iter.next(typeof found === "boolean");
                    if (done) {
                        chunk = { done: true } as ReadableStreamDefaultReadDoneResult;
                    } else {
                        chunk = { done: false, value } as ReadableStreamDefaultReadValueResult<Buffer>;
                        cumul = Buffer.concat([cumul, value!]);
                        const v = await proc.has_embed(cumul);
                        if (typeof v == "string") {
                            return [await proc.extract(cumul, v), false] as [WorkerEmbeddedFile[], boolean];
                        }
                        found = v;
                    }
                } while (found !== false && !chunk.done /* Because we only embed links now, it's safe to assume we get everything we need in the first chunk */);
                succ = true;
                await iter.next(true);
                if (found !== true) {
                    //console.log(`Gave up on ${src} after downloading ${cumul.byteLength} bytes...`);
                    return;
                }
                return [await proc.extract(cumul), false] as [WorkerEmbeddedFile[], boolean];
            } catch(err) {
                console.log(err);
                // ignore error and retry with another link
            }
        } while (!succ);
    }));
    return ret.filter(e => e).map(e => e!);
};

(async () => {
    onmessage = async (msg: MessageEvent<any>) => {
        const des = deserializeMessage(msg.data);
        switch (des.type) {
            case 'ipc': {
                if (platform.port1.onmessage)
                    platform.port1.onmessage(new MessageEvent("message", { data: des.res }));
                break;
            }
            case 'cmd': {
                switch (des.fun) {
                    case 'processImage': {
                        //console.log('Received process image command', des);
                        const res = await processImage(des.args[0], des.args[1], des.args[2], des.args[3]);
                        //console.log('Finished process image command', des);
                        const tr: Transferable[] = [];
                        for (const ef of res) {
                            for (const e of ef[0]) {
                                if (Buffer.isBuffer(e.thumbnail) || e.thumbnail instanceof Uint8Array)
                                    tr.push(e.thumbnail.buffer);
                                if (Buffer.isBuffer(e.data) || e.data instanceof Uint8Array)
                                    tr.push(e.data.buffer);
                            }
                        }
                        //console.log('Sent reply', res, des);
                        postMessage({
                            type: 'reply',
                            id: des.id,
                            res
                        }, [...new Set(tr)]);
                    }
                }
                break;
            }
            case 'ag': {
                const cb = pendinggens[des.id].shift();
                if (cb) {
                    cb(des.res);
                }
                break;
            }
        }
    };
})();