Skip to content

Commit

Permalink
Mount python worker files after taking memory snapshot
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
hoodmane committed Sep 27, 2024
1 parent ba8a592 commit 246c248
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 6 deletions.
8 changes: 6 additions & 2 deletions src/pyodide/internal/python.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,8 @@ import {
TRANSITIVE_REQUIREMENTS,
SITE_PACKAGES,
adjustSysPath,
mountLib,
mountSitePackages,
mountWorkerFiles,
} from 'pyodide-internal:setupPackages';
import { reportError } from 'pyodide-internal:util';
import {
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand Down
15 changes: 11 additions & 4 deletions src/pyodide/internal/setupPackages.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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`
);
}

/**
Expand All @@ -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`
);
}

Expand Down

0 comments on commit 246c248

Please sign in to comment.