Enforce output file count and size limits in the Pyodide worker

GitOrigin-RevId: 2cc61613381243d810a8cb9e1c2c32fa9e751da7

services/web/frontend/js/features/ide-react/components/editor/python/pyodide-worker-output-limits.ts

@@ -0,0 +1,53 @@
+const BYTES_PER_MB = 1024 * 1024
+
+export const MAX_OUTPUT_FILES = 50
+export const MAX_OUTPUT_TOTAL_BYTES = 100 * BYTES_PER_MB
+export const MAX_OUTPUT_FILE_BYTES = 50 * BYTES_PER_MB
+
+export type OutputLimitViolation = {
+  kind: 'count' | 'total-output-size' | 'single-file-size'
+  message: string
+}
+
+export function checkOutputCount(count: number): OutputLimitViolation | null {
+  if (count > MAX_OUTPUT_FILES) {
+    return {
+      kind: 'count',
+      message: `Output limit exceeded: ${count} files generated (max ${MAX_OUTPUT_FILES})`,
+    }
+  }
+  return null
+}
+
+export function checkOutputLimits(
+  files: { path: string; size: number }[]
+): OutputLimitViolation | null {
+  const countViolation = checkOutputCount(files.length)
+  if (countViolation) {
+    return countViolation
+  }
+
+  let totalBytes = 0
+  for (const file of files) {
+    if (file.size > MAX_OUTPUT_FILE_BYTES) {
+      const fileMB = Math.ceil(file.size / BYTES_PER_MB)
+      const maxMB = MAX_OUTPUT_FILE_BYTES / BYTES_PER_MB
+      return {
+        kind: 'single-file-size',
+        message: `Output limit exceeded: ${file.path} is ${fileMB}MB (max ${maxMB}MB per file)`,
+      }
+    }
+    totalBytes += file.size
+  }
+
+  if (totalBytes > MAX_OUTPUT_TOTAL_BYTES) {
+    const totalMB = Math.ceil(totalBytes / BYTES_PER_MB)
+    const maxMB = MAX_OUTPUT_TOTAL_BYTES / BYTES_PER_MB
+    return {
+      kind: 'total-output-size',
+      message: `Output limit exceeded: ${totalMB}MB total (max ${maxMB}MB)`,
+    }
+  }
+
+  return null
+}

services/web/frontend/js/features/ide-react/components/editor/python/pyodide.worker.ts

@@ -8,6 +8,10 @@ import type {
   PyodideWorkerRequest,
   RunCodeRequest,
 } from './pyodide-worker-messages'
+import {
+  checkOutputCount,
+  checkOutputLimits,
+} from './pyodide-worker-output-limits'
 
 type PyodideFS = PyodideInterface['FS']
 type PyodideModule = typeof import('pyodide')
@@ -87,11 +91,11 @@ async function handleInit(msg: InitRequest) {
 async function handleRunCode(msg: RunCodeRequest) {
   const { fileId, executionId } = msg
 
-  if (!pyodideModule || !pyodideIndexUrl) {
+  const postFailure = (stream: 'stderr' | 'info', line: string) => {
     self.postMessage({
       type: 'output-line',
-      stream: 'stderr',
-      line: 'Pyodide is not initialized',
+      stream,
+      line,
       fileId,
       executionId,
     })
@@ -103,6 +107,10 @@ async function handleRunCode(msg: RunCodeRequest) {
       outputs: [],
       outputFiles: [],
     })
+  }
+
+  if (!pyodideModule || !pyodideIndexUrl) {
+    postFailure('stderr', 'Pyodide is not initialized')
     return
   }
 
@@ -179,27 +187,39 @@ async function handleRunCode(msg: RunCodeRequest) {
     const errorMessage =
       runError instanceof Error ? runError.message : String(runError)
 
-    self.postMessage({
-      type: 'output-line',
-      stream: 'stderr',
-      line: errorMessage,
-      fileId,
-      executionId,
-    })
-    self.postMessage({
-      type: 'run-code-result',
-      fileId,
-      executionId,
-      success: false,
-      outputs: paths,
-      outputFiles: [],
-    })
+    postFailure('stderr', errorMessage)
+    return
+  }
+
+  const countViolation = checkOutputCount(paths.length)
+  if (countViolation) {
+    postFailure('info', countViolation.message)
+    return
+  }
+
+  const filesWithSizes: { path: string; size: number }[] = []
+  for (const writtenPath of paths) {
+    try {
+      filesWithSizes.push({
+        path: writtenPath,
+        size: fs.stat(writtenPath).size,
+      })
+    } catch {
+      // A script can write a file and later delete or rename it before the run
+      // finishes; fs.stat would then throw and we'd never post a
+      // run-code-result, leaving the UI stuck. Skip paths we can't stat.
+    }
+  }
+
+  const sizeViolation = checkOutputLimits(filesWithSizes)
+  if (sizeViolation) {
+    postFailure('info', sizeViolation.message)
     return
   }
 
   const outputFiles: OutputFileData[] = []
   const transferables: Transferable[] = []
-  for (const writtenPath of paths) {
+  for (const { path: writtenPath } of filesWithSizes) {
     const content = fs.readFile(writtenPath)
     const relativePath = writtenPath.slice(PROJECT_FS_PREFIX.length)
     outputFiles.push({ relativePath, content })
@@ -218,7 +238,7 @@ async function handleRunCode(msg: RunCodeRequest) {
       fileId,
       executionId,
       success: true,
-      outputs: paths,
+      outputs: filesWithSizes.map(f => f.path),
       outputFiles,
     },
     transferables

services/web/test/frontend/features/ide-react/unit/editor/pyodide-worker-output-limits.spec.ts

@@ -0,0 +1,115 @@
+import { expect } from 'chai'
+import {
+  MAX_OUTPUT_FILES,
+  MAX_OUTPUT_FILE_BYTES,
+  MAX_OUTPUT_TOTAL_BYTES,
+  checkOutputLimits,
+} from '@/features/ide-react/components/editor/python/pyodide-worker-output-limits'
+
+const BYTES_PER_MB = 1024 * 1024
+
+function makeFiles(
+  count: number,
+  sizePerFile: number
+): { path: string; size: number }[] {
+  return Array.from({ length: count }, (_, i) => ({
+    path: `/project/file${i}.bin`,
+    size: sizePerFile,
+  }))
+}
+
+describe('checkOutputLimits', function () {
+  it('returns null when both limits are within bounds', function () {
+    const files = makeFiles(10, 1024)
+    expect(checkOutputLimits(files)).to.equal(null)
+  })
+
+  it('returns null for an empty file list', function () {
+    expect(checkOutputLimits([])).to.equal(null)
+  })
+
+  it('returns null at exactly the file count limit', function () {
+    const files = makeFiles(MAX_OUTPUT_FILES, 1024)
+    expect(checkOutputLimits(files)).to.equal(null)
+  })
+
+  it('returns a count violation when the file count exceeds the limit', function () {
+    const files = makeFiles(MAX_OUTPUT_FILES + 1, 1)
+    const violation = checkOutputLimits(files)
+    expect(violation).to.deep.equal({
+      kind: 'count',
+      message: `Output limit exceeded: ${MAX_OUTPUT_FILES + 1} files generated (max ${MAX_OUTPUT_FILES})`,
+    })
+  })
+
+  it('reports the actual file count in the count violation message', function () {
+    const files = makeFiles(73, 1)
+    const violation = checkOutputLimits(files)
+    expect(violation).to.not.equal(null)
+    expect(violation!.kind).to.equal('count')
+    expect(violation!.message).to.equal(
+      'Output limit exceeded: 73 files generated (max 50)'
+    )
+  })
+
+  it('returns null at exactly the per-file size limit', function () {
+    const files = [{ path: '/project/big.bin', size: MAX_OUTPUT_FILE_BYTES }]
+    expect(checkOutputLimits(files)).to.equal(null)
+  })
+
+  it('returns a single-file-size violation when one file exceeds the per-file limit', function () {
+    const files = [
+      { path: '/project/big.bin', size: MAX_OUTPUT_FILE_BYTES + 1 },
+    ]
+    const violation = checkOutputLimits(files)
+    expect(violation).to.not.equal(null)
+    expect(violation!.kind).to.equal('single-file-size')
+    expect(violation!.message).to.equal(
+      'Output limit exceeded: /project/big.bin is 51MB (max 50MB per file)'
+    )
+  })
+
+  it('reports the offending file path and rounded size in the single-file-size message', function () {
+    const files = [
+      { path: '/project/small.bin', size: 1 * BYTES_PER_MB },
+      { path: '/project/huge.bin', size: 80 * BYTES_PER_MB },
+    ]
+    const violation = checkOutputLimits(files)
+    expect(violation).to.deep.equal({
+      kind: 'single-file-size',
+      message:
+        'Output limit exceeded: /project/huge.bin is 80MB (max 50MB per file)',
+    })
+  })
+
+  it('returns a total-output-size violation when the summed size exceeds the total limit', function () {
+    const files = [
+      { path: '/project/a.bin', size: 40 * BYTES_PER_MB },
+      { path: '/project/b.bin', size: 40 * BYTES_PER_MB },
+      { path: '/project/c.bin', size: 40 * BYTES_PER_MB },
+    ]
+    const violation = checkOutputLimits(files)
+    expect(violation).to.not.equal(null)
+    expect(violation!.kind).to.equal('total-output-size')
+    expect(violation!.message).to.equal(
+      'Output limit exceeded: 120MB total (max 100MB)'
+    )
+  })
+
+  it('returns null at exactly the total size limit', function () {
+    const files = [
+      { path: '/project/a.bin', size: 50 * BYTES_PER_MB },
+      { path: '/project/b.bin', size: 50 * BYTES_PER_MB },
+    ]
+    expect(checkOutputLimits(files)).to.equal(null)
+    const total = files.reduce((acc, f) => acc + f.size, 0)
+    expect(total).to.equal(MAX_OUTPUT_TOTAL_BYTES)
+  })
+
+  it('reports the count violation when both limits are exceeded', function () {
+    const files = makeFiles(MAX_OUTPUT_FILES + 10, MAX_OUTPUT_TOTAL_BYTES)
+    const violation = checkOutputLimits(files)
+    expect(violation).to.not.equal(null)
+    expect(violation!.kind).to.equal('count')
+  })
+})