From 246c248e33af1c9199443a7afed33980d2a7dacd Mon Sep 17 00:00:00 2001
From: Hood Chatham <roberthoodchatham@gmail.com>
Date: Thu, 26 Sep 2024 13:48:42 +0200
Subject: [PATCH] Mount python worker files after taking memory snapshot

This ensures that the contents of worker files cannot be accessed prior to
taking the snapshot and so won't appear in the linear memory.
---
 src/pyodide/internal/python.ts        |  8 ++++++--
 src/pyodide/internal/setupPackages.ts | 15 +++++++++++----
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/src/pyodide/internal/python.ts b/src/pyodide/internal/python.ts
index 37d1372600c..a46f04c7062 100644
--- a/src/pyodide/internal/python.ts
+++ b/src/pyodide/internal/python.ts
@@ -4,7 +4,8 @@ import {
   TRANSITIVE_REQUIREMENTS,
   SITE_PACKAGES,
   adjustSysPath,
-  mountLib,
+  mountSitePackages,
+  mountWorkerFiles,
 } from 'pyodide-internal:setupPackages';
 import { reportError } from 'pyodide-internal:util';
 import {
@@ -201,7 +202,7 @@ async function instantiateEmscriptenModule(
  */
 async function prepareWasmLinearMemory(Module: Module): Promise<void> {
   // Note: if we are restoring from a snapshot, runtime is not initialized yet.
-  mountLib(Module, SITE_PACKAGES.rootInfo);
+  mountSitePackages(Module, SITE_PACKAGES.rootInfo);
   entropyMountFiles(Module);
   if (SHOULD_RESTORE_SNAPSHOT) {
     restoreSnapshot(Module);
@@ -229,6 +230,9 @@ export async function loadPyodide(
     prepareWasmLinearMemory(Module)
   );
   maybeSetupSnapshotUpload(Module);
+  // Mount worker files after doing snapshot upload so we ensure that data from the files is never
+  // present in snapshot memory.
+  mountWorkerFiles(Module);
 
   // Finish setting up Pyodide's ffi so we can use the nice Python interface
   await enterJaegerSpan('finalize_bootstrap', Module.API.finalizeBootstrap);
diff --git a/src/pyodide/internal/setupPackages.ts b/src/pyodide/internal/setupPackages.ts
index fe386f714a4..1b1fa84a4fb 100644
--- a/src/pyodide/internal/setupPackages.ts
+++ b/src/pyodide/internal/setupPackages.ts
@@ -189,19 +189,26 @@ export function getSitePackagesPath(Module: Module): string {
  * details, so even though we want these directories to be on sys.path, we
  * handle that separately in adjustSysPath.
  */
-export function mountLib(Module: Module, info: TarFSInfo): void {
+export function mountSitePackages(Module: Module, info: TarFSInfo): void {
   const tarFS = createTarFS(Module);
-  const mdFS = createMetadataFS(Module);
   const site_packages = getSitePackagesPath(Module);
   Module.FS.mkdirTree(site_packages);
-  Module.FS.mkdirTree('/session/metadata');
   if (!LOAD_WHEELS_FROM_R2 && !LOAD_WHEELS_FROM_ARTIFACT_BUNDLER) {
     // if we are not loading additional wheels, then we're done
     // with site-packages and we can mount it here. Otherwise, we must mount it in
     // loadPackages().
     Module.FS.mount(tarFS, { info }, site_packages);
   }
+}
+
+export function mountWorkerFiles(Module: Module) {
+  Module.FS.mkdirTree('/session/metadata');
+  const mdFS = createMetadataFS(Module);
   Module.FS.mount(mdFS, {}, '/session/metadata');
+  simpleRunPython(
+    Module,
+    `import sys; sys.path.append("/session/metadata"); del sys`
+  );
 }
 
 /**
@@ -212,7 +219,7 @@ export function adjustSysPath(Module: Module): void {
   const site_packages = getSitePackagesPath(Module);
   simpleRunPython(
     Module,
-    `import sys; sys.path.append("/session/metadata"); sys.path.append("${site_packages}"); del sys`
+    `import sys; sys.path.append("${site_packages}"); del sys`
   );
 }