Merge pull request #3804 from heygen-com/fix/prinfra-353-jpeg-input-validation

fix(engine): reject malformed JPEG input before encoding

packages/engine/src/services/__fixtures__/jpeg.ts

@@ -0,0 +1,5 @@
+// A 64x64 JPEG encoded by FFmpeg; contains real DQT, DHT, SOF and scan data.
+export const validJpeg = Buffer.from(
+  "/9j/4AAQSkZJRgABAgAAAQABAAD//gAQTGF2YzYyLjI4LjEwMQD/2wBDAAgEBAQEBAUFBQUFBQYGBgYGBgYGBgYGBgYHBwcICAgHBwcGBgcHCAgICAkJCQgICAgJCQoKCgwMCwsODg4RERT/xABNAAEBAAAAAAAAAAAAAAAAAAAABgEBAQEAAAAAAAAAAAAAAAAAAAYHEAEAAAAAAAAAAAAAAAAAAAAAEQEAAAAAAAAAAAAAAAAAAAAA/8AAEQgAQABAAwEiAAIRAAMRAP/aAAwDAQACEQMRAD8AiwEm38AAAAAAAAAAAAAAAAAAAAAAAAAAAAAB/9k=",
+  "base64",
+);

packages/engine/src/services/streamingEncoder.test.ts

@@ -10,6 +10,8 @@
  * master-display / max-cll and ship as SDR BT.2020 again.
  */
 
+import { validJpeg } from "./__fixtures__/jpeg.js";
+
 import { EventEmitter } from "events";
 import { mkdtempSync } from "fs";
 import { tmpdir } from "os";
@@ -711,6 +713,41 @@ describe("spawnStreamingEncoder lifecycle and cleanup", () => {
     expect(threw).toBe(false);
   });
 
+  it("rejects a malformed JPEG before writing it and reports its input frame index", async () => {
+    const { spawn, calls } = createSpawnSpy();
+    vi.resetModules();
+    vi.doMock("child_process", () => ({ spawn }));
+    const { spawnStreamingEncoder } = await import("./streamingEncoder.js");
+    const dir = mkdtempSync(join(tmpdir(), "se-jpeg-"));
+    const encoder = await spawnStreamingEncoder(join(dir, "out.mp4"), baseOptions);
+    const proc = calls[0]!.proc;
+    const write = vi.spyOn(proc.stdin, "write");
+    expect(await encoder.writeFrame(validJpeg)).toBe(true);
+    const malformed = Buffer.from(validJpeg);
+    malformed[malformed.indexOf(Buffer.from([0xff, 0xdb])) + 4] = 0x20;
+    await expect(encoder.writeFrame(malformed)).rejects.toThrow(
+      "Invalid JPEG input at frame 1: invalid JPEG DQT precision 2",
+    );
+    expect(write).toHaveBeenCalledTimes(1);
+    process.nextTick(() => proc.emit("close", 0));
+    await encoder.close();
+  });
+
+  it.each([
+    { ...baseOptions, imageFormat: "png" as const },
+    { ...baseOptions, rawInputFormat: "rgb48le" as const },
+  ])("does not apply JPEG validation to PNG or raw frames", async (options) => {
+    const { spawn, calls } = createSpawnSpy();
+    vi.resetModules();
+    vi.doMock("child_process", () => ({ spawn }));
+    const { spawnStreamingEncoder } = await import("./streamingEncoder.js");
+    const dir = mkdtempSync(join(tmpdir(), "se-non-jpeg-"));
+    const encoder = await spawnStreamingEncoder(join(dir, "out.mp4"), options);
+    expect(await encoder.writeFrame(Buffer.from([0]))).toBe(true);
+    process.nextTick(() => calls[0]!.proc.emit("close", 0));
+    await encoder.close();
+  });
+
   it("writeFrame returns false after ffmpeg has exited", async () => {
     const { spawn, calls } = createSpawnSpy();
     vi.resetModules();
@@ -720,7 +757,7 @@ describe("spawnStreamingEncoder lifecycle and cleanup", () => {
     const dir = mkdtempSync(join(tmpdir(), "se-writefail-"));
     const encoder = await spawnStreamingEncoder(join(dir, "out.mp4"), baseOptions);
 
-    expect(await encoder.writeFrame(Buffer.from([0]))).toBe(true);
+    expect(await encoder.writeFrame(validJpeg)).toBe(true);
 
     const proc = calls[0]!.proc;
     await new Promise<void>((resolve) => {
@@ -730,7 +767,7 @@ describe("spawnStreamingEncoder lifecycle and cleanup", () => {
       });
     });
 
-    expect(await encoder.writeFrame(Buffer.from([0]))).toBe(false);
+    expect(await encoder.writeFrame(validJpeg)).toBe(false);
   });
 
   it("waits for child close when stdin dies first so the interruption reason is observable", async () => {
@@ -744,7 +781,7 @@ describe("spawnStreamingEncoder lifecycle and cleanup", () => {
     const proc = calls[0]!.proc;
     proc.stdin.destroyed = true;
 
-    const writePromise = encoder.writeFrame(Buffer.from([0]));
+    const writePromise = encoder.writeFrame(validJpeg);
     await expect(resolveWithin(writePromise, 10)).resolves.toBe("timeout");
 
     proc.stderr.emit("data", Buffer.from("Exiting normally, received signal 15.\n"));
@@ -766,7 +803,7 @@ describe("spawnStreamingEncoder lifecycle and cleanup", () => {
     const proc = calls[0]!.proc;
     proc.stdin.write = (_chunk: Buffer): boolean => false;
 
-    const writeResult = encoder.writeFrame(Buffer.from([1])) as unknown;
+    const writeResult = encoder.writeFrame(validJpeg) as unknown;
     expect(writeResult).toBeInstanceOf(Promise);
 
     const writePromise = writeResult as Promise<boolean>;
@@ -804,7 +841,7 @@ describe("spawnStreamingEncoder lifecycle and cleanup", () => {
     proc.stdin.write = (_chunk: Buffer): boolean => false;
 
     for (let i = 0; i < 12; i++) {
-      const writePromise = encoder.writeFrame(Buffer.from([i]));
+      const writePromise = encoder.writeFrame(validJpeg);
 
       await Promise.resolve();
       expect(proc.stdin.listenerCount("drain")).toBe(baselineDrainListeners + 1);
@@ -833,7 +870,7 @@ describe("spawnStreamingEncoder lifecycle and cleanup", () => {
     const proc = calls[0]!.proc;
     proc.stdin.write = (_chunk: Buffer): boolean => false;
 
-    const writeResult = encoder.writeFrame(Buffer.from([1])) as unknown;
+    const writeResult = encoder.writeFrame(validJpeg) as unknown;
     expect(writeResult).toBeInstanceOf(Promise);
 
     const writePromise = writeResult as Promise<boolean>;
@@ -872,7 +909,7 @@ describe("spawnStreamingEncoder lifecycle and cleanup", () => {
       return false;
     };
 
-    const writePromise = encoder.writeFrame(Buffer.from([1]));
+    const writePromise = encoder.writeFrame(validJpeg);
 
     await expect(resolveWithin(writePromise)).resolves.toBe(false);
     expect(encoder.getExitStatus()).toBe("error");
@@ -927,7 +964,7 @@ describe("spawnStreamingEncoder lifecycle and cleanup", () => {
       // progressing" capture the encoder must still be alive. The old total-
       // render timeout would have fired SIGTERM at ~1000ms.
       for (let i = 0; i < 9; i++) {
-        await encoder.writeFrame(Buffer.from([i]));
+        await encoder.writeFrame(validJpeg);
         vi.advanceTimersByTime(900);
       }
       expect(proc.kill).not.toHaveBeenCalled();
@@ -964,7 +1001,7 @@ describe("spawnStreamingEncoder lifecycle and cleanup", () => {
       // A buffered write should remain pending and must NOT reset the timer.
       // The 1000ms timer (last reset on spawn) therefore elapses while the
       // caller is correctly back-pressured on the first frame.
-      const writePromise = encoder.writeFrame(Buffer.from([0]));
+      const writePromise = encoder.writeFrame(validJpeg);
       await Promise.resolve();
 
       vi.advanceTimersByTime(1100);

packages/engine/src/services/streamingEncoder.ts

@@ -13,6 +13,7 @@
  *      exposes an async `writeFrame(buffer)` + `close()` API.
  */
 
+import { jpegInputError } from "../utils/jpegInput.js";
 import { spawn, type ChildProcess } from "child_process";
 import { once } from "events";
 import { trackChildProcess } from "../utils/processTracker.js";
@@ -537,6 +538,7 @@ export async function spawnStreamingEncoder(
     }
   };
 
+  let inputFrameIndex = 0;
   const encoder: StreamingEncoder = {
     writeFrame: async (buffer: Buffer): Promise<boolean> => {
       const stdin = ffmpeg.stdin;
@@ -557,7 +559,12 @@ export async function spawnStreamingEncoder(
       // so without this copy the pipe would read partially-overwritten data
       // and flicker.
       const copy = Buffer.from(buffer);
+      if (!options.rawInputFormat && (options.imageFormat ?? "jpeg") === "jpeg") {
+        const error = jpegInputError(copy);
+        if (error) throw new Error(`Invalid JPEG input at frame ${inputFrameIndex}: ${error}`);
+      }
       const accepted = stdin.write(copy);
+      inputFrameIndex++;
       // Reset inactivity timer immediately ONLY on `accepted === true`. `true`
       // means the write went through to the kernel pipe without buffering in
       // Node — proof FFmpeg is actually consuming. `false` means Node's writable

packages/engine/src/utils/jpegInput.test.ts

@@ -0,0 +1,38 @@
+import { describe, expect, it } from "vitest";
+import { validJpeg } from "../services/__fixtures__/jpeg.js";
+import { jpegInputError } from "./jpegInput.js";
+
+function withDqt(descriptor: number, bytes: number): Buffer {
+  return Buffer.concat([
+    Buffer.from([0xff, 0xd8, 0xff, 0xdb, 0, bytes + 3, descriptor]),
+    Buffer.alloc(bytes, 1),
+    Buffer.from([0xff, 0xda, 0, 2]),
+  ]);
+}
+
+describe("JPEG input diagnostics", () => {
+  it("accepts an actual FFmpeg-encoded JPEG", () => {
+    expect(jpegInputError(validJpeg)).toBeUndefined();
+  });
+  it("identifies the reproduced invalid DQT precision", () => {
+    const malformed = Buffer.from(validJpeg);
+    const dqt = malformed.indexOf(Buffer.from([0xff, 0xdb]));
+    expect(dqt).toBeGreaterThan(0);
+    malformed[dqt + 4] = 0x20;
+    expect(jpegInputError(malformed)).toBe("invalid JPEG DQT precision 2");
+  });
+  it.each([0, 1])("accepts DQT precision %i", (precision) => {
+    expect(jpegInputError(withDqt(precision << 4, precision === 0 ? 64 : 128))).toBeUndefined();
+  });
+  it("rejects a truncated DQT table within a complete segment", () => {
+    expect(jpegInputError(withDqt(0, 63))).toBe("truncated JPEG DQT table");
+  });
+  it.each([
+    [Buffer.alloc(0), "empty frame"],
+    [Buffer.from([0]), "missing JPEG SOI marker"],
+    [Buffer.from([0xff, 0xd8, 0xff, 0xdb, 0, 67]), "invalid JPEG segment length"],
+    [Buffer.from([0xff, 0xd8, 0xff]), "truncated JPEG marker"],
+  ])("reports malformed input without reading past its buffer", (buffer, message) => {
+    expect(jpegInputError(buffer)).toBe(message);
+  });
+});

packages/engine/src/utils/jpegInput.ts

@@ -0,0 +1,50 @@
+/** Check JPEG headers before image2pipe; entropy-coded pixels still belong to FFmpeg. */
+export function jpegInputError(buffer: Buffer): string | undefined {
+  if (buffer.length === 0) return "empty frame";
+  if (buffer[0] !== 0xff || buffer[1] !== 0xd8) return "missing JPEG SOI marker";
+  let offset = 2;
+  while (offset < buffer.length) {
+    const segment = readSegment(buffer, offset);
+    if (typeof segment === "string") return segment;
+    const { marker, start, end } = segment;
+    if (marker === 0xdb) {
+      const error = quantizationTableError(buffer, start, end);
+      if (error) return error;
+    }
+    // Do not scan compressed pixel data as markers (byte stuffing/restarts).
+    if (marker === 0xda) return undefined;
+    offset = end;
+  }
+  return "missing JPEG scan";
+}
+
+function readSegment(
+  buffer: Buffer,
+  offset: number,
+): { marker: number; start: number; end: number } | string {
+  if (buffer[offset++] !== 0xff) return "invalid JPEG marker";
+  while (buffer[offset] === 0xff) offset++;
+  const marker = buffer[offset++];
+  if (marker === undefined) return "truncated JPEG marker";
+  if (marker === 0x01) return { marker, start: offset, end: offset };
+  if (offset + 2 > buffer.length) return "truncated JPEG segment length";
+  const length = buffer.readUInt16BE(offset);
+  const end = offset + length;
+  if (length < 2 || end > buffer.length) return "invalid JPEG segment length";
+  return { marker, start: offset + 2, end };
+}
+
+function quantizationTableError(buffer: Buffer, start: number, end: number): string | undefined {
+  if (start === end) return "empty JPEG DQT segment";
+  let offset = start;
+  while (offset < end) {
+    const descriptor = buffer[offset++];
+    if (descriptor === undefined) return "truncated JPEG DQT descriptor";
+    const precision = descriptor >> 4;
+    if (precision > 1) return `invalid JPEG DQT precision ${precision}`;
+    if ((descriptor & 0x0f) > 3) return "invalid JPEG DQT table id";
+    offset += precision === 0 ? 64 : 128;
+    if (offset > end) return "truncated JPEG DQT table";
+  }
+  return undefined;
+}