bake: unregister Bake::SourceProvider from the source map table when it is destroyed

src/jsc/virtual_machine_exports.rs

@@ -257,8 +257,8 @@ pub fn get_verbose_fetch_value() -> i32 {
     }
 }
 
-// `Bun__addBakeSourceProviderSourceMap` / `Bun__addDevServerSourceProvider` /
-// `Bun__removeDevServerSourceProvider` live in
+// `Bun__{add,remove}BakeSourceProviderSourceMap` /
+// `Bun__{add,remove}DevServerSourceProvider` live in
 // `bun_runtime::bake::source_provider_exports` (their callers are bake's C++
 // source providers; LAYERING).
 

src/runtime/bake/BakeSourceProvider.h

@@ -9,6 +9,7 @@ namespace Bake {
 class SourceProvider;
 
 extern "C" void Bun__addBakeSourceProviderSourceMap(void* bun_vm, SourceProvider* opaque_source_provider, const BunString* specifier);
+extern "C" void Bun__removeBakeSourceProviderSourceMap(void* bun_vm, SourceProvider* opaque_source_provider, const BunString* specifier);
 
 class SourceProvider final : public JSC::StringSourceProvider {
 public:
@@ -50,6 +51,15 @@ class SourceProvider final : public JSC::StringSourceProvider {
     {
     }
 
+    // Takes the Rust VirtualMachine, not the Zig::GlobalObject: this runs
+    // from JSC's sweep, possibly after the global object cell itself has been
+    // swept (see DevServerSourceProvider).
+    ~SourceProvider()
+    {
+        auto specifier = Bun::toString(sourceURL());
+        Bun__removeBakeSourceProviderSourceMap(m_bunVM, this, &specifier);
+    }
+
     void* m_bunVM;
 };
 

src/runtime/bake/source_provider_exports.rs

@@ -1,9 +1,9 @@
 //! Rust side of `BakeSourceProvider.h` / `DevServerSourceProvider.h`: the
 //! FFI bindings and `SourceProvider` impls for bake's two C++ source
-//! providers, plus the host exports that register them with the VM's
-//! `SavedSourceMap` so stack remapping can resolve dev-server /
-//! bake-production output. `bun_sourcemap` sees these only as erased
-//! `AnySourceProvider` handles.
+//! providers, plus the host exports that register them with (and, from their
+//! destructors, remove them from) the VM's `SavedSourceMap` so stack remapping
+//! can resolve dev-server / bake-production output. `bun_sourcemap` sees these
+//! only as erased `AnySourceProvider` handles.
 //!
 //! `#[unsafe(no_mangle)] extern "C"` thunks are emitted by
 //! `src/codegen/generate-host-exports.ts` from the `// HOST_EXPORT(Sym, c)`
@@ -18,8 +18,8 @@ use bun_sourcemap::parsed_source_map::AnySourceProvider;
 use bun_sourcemap::{SourceContentPtr, SourceProvider};
 
 bun_opaque::opaque_ffi! {
-    /// Opaque handle to the C++ `Bake::SourceProvider` for production-build
-    /// sources (`BakeSourceProvider.cpp`).
+    /// Opaque handle to the C++ `Bake::SourceProvider` (`BakeSourceProvider.h`),
+    /// registered from its `create()` and unregistered from its destructor.
     pub(crate) struct BakeSourceProvider;
     /// Opaque handle to the C++ `Bake::DevServerSourceProvider`
     /// (`DevServerSourceProvider.cpp`).
@@ -146,6 +146,17 @@ pub(crate) fn add_bake_source_provider_source_map(
     );
 }
 
+// HOST_EXPORT(Bun__removeBakeSourceProviderSourceMap, c)
+pub(crate) fn remove_bake_source_provider_source_map(
+    vm: &mut VirtualMachine,
+    opaque_source_provider: *mut c_void,
+    specifier: &BunString,
+) {
+    let slice = specifier.to_utf8();
+    vm.source_mappings
+        .remove_source_provider(opaque_source_provider, slice.slice());
+}
+
 // HOST_EXPORT(Bun__addDevServerSourceProvider, c)
 pub(crate) fn add_dev_server_source_provider(
     vm: &mut VirtualMachine,

test/bake/deinitialization.test.ts

@@ -1,5 +1,5 @@
 import { expect, test } from "bun:test";
-import { bunEnv, bunExe, tempDir } from "harness";
+import { bunEnv, bunExe, isASAN, tempDir } from "harness";
 import path from "node:path";
 
 test("dev server deinitializes itself", () => {
@@ -50,3 +50,104 @@ test("dev server is deinitialized before its arena when listen fails", async ()
   expect(stdout).toBe('{"code":"EADDRINUSE","deinits":1}\n');
   expect(exitCode).toBe(0);
 });
+
+// Every framework dev server evaluates the server runtime through its own
+// Bake::SourceProvider, registered in the VM's source map table under the one
+// key "bake://server-runtime.js", so the table points at whichever server
+// started last. Once that server is torn down and its provider freed, the
+// next stack trace through a surviving server's runtime looks the map up via
+// the table entry, so the provider must have removed itself by then. Malloc=1
+// puts JSC's allocations under the system malloc so ASAN sees the read of the
+// freed provider instead of bmalloc quietly handing back the stale memory.
+test.skipIf(!isASAN)(
+  "tearing down one dev server does not leave its server runtime source provider registered for another",
+  async () => {
+    using dir = tempDir("bake-deinit-source-provider", {
+      "server.ts": `
+        export function render(req, meta) {
+          return meta.pageModule.default(req, meta);
+        }
+        export function registerClientReference(value, file, uid) {
+          return { value, file, uid };
+        }
+      `,
+      "routes/index.ts": `
+        export default function () {
+          return new Response(new Error("probe").stack);
+        }
+      `,
+      "fixture.ts": `
+        import { getDevServerDeinitCount } from "bun:internal-for-testing";
+        import path from "node:path";
+
+        const serverEntryPoint = path.join(import.meta.dir, "server.ts");
+        const framework = {
+          fileSystemRouterTypes: [{ root: "routes", style: "nextjs-pages", serverEntryPoint }],
+          serverComponents: {
+            separateSSRGraph: false,
+            serverRuntimeImportSource: serverEntryPoint,
+            serverRegisterClientReferenceExport: "registerClientReference",
+          },
+        };
+        const start = () => Bun.serve({ port: 0, app: { framework } });
+
+        async function collect(isDone) {
+          for (let i = 0; i < 200 && !isDone(); i++) {
+            Bun.gc(true);
+            await new Promise(resolve => setTimeout(resolve, 1));
+          }
+          if (!isDone()) throw new Error("dev server was not deinitialized");
+        }
+
+        const survivor = start();
+
+        // Started second, so its provider replaces the survivor's table entry.
+        // Stopped inside its own function so nothing on the stack keeps it alive.
+        const deinitsBefore = getDevServerDeinitCount();
+        const survivorDebugHooks = globalThis.DEBUG;
+        (function startAndStop() {
+          start().stop(true);
+        })();
+        // Debug builds of the server runtime publish globalThis.DEBUG, whose
+        // ASSERT is a function of whichever runtime ran last and so would keep
+        // the second server's provider alive (release builds compile it out).
+        // Put the survivor's back; it is the one that keeps handling requests.
+        globalThis.DEBUG = survivorDebugHooks;
+        await collect(() => getDevServerDeinitCount() > deinitsBefore);
+        // Deinit dropped the dev server's references to the runtime's
+        // functions; a few more collections sweep them and free the provider.
+        for (let i = 0; i < 5; i++) {
+          Bun.gc(true);
+          await new Promise(resolve => setTimeout(resolve, 1));
+        }
+
+        const response = await fetch(survivor.url);
+        const stack = await response.text();
+        survivor.stop(true);
+        console.log("status:", response.status);
+        console.log("runtime frame:", stack.includes("bake://server-runtime.js"));
+      `,
+    });
+
+    await using proc = Bun.spawn({
+      cmd: [bunExe(), "fixture.ts"],
+      cwd: String(dir),
+      env: {
+        ...bunEnv,
+        Malloc: "1",
+        // detect_leaks=0: Malloc=1 exposes JSC's process-lifetime startup
+        // allocations to LSAN; this test is about the use-after-free.
+        ASAN_OPTIONS: [bunEnv.ASAN_OPTIONS, "detect_leaks=0"].filter(Boolean).join(":"),
+      },
+      stdout: "pipe",
+      stderr: "pipe",
+    });
+    const [stdout, stderr, exitCode] = await Promise.all([proc.stdout.text(), proc.stderr.text(), proc.exited]);
+
+    expect(stderr).not.toContain("AddressSanitizer");
+    // The runtime frame is what sends the stack trace through the table entry.
+    expect(stdout).toEndWith("status: 200\nruntime frame: true\n");
+    expect(exitCode).toBe(0);
+  },
+  60_000,
+);