From f576f4ae74af87aa0756702d8e7e3f8a9e77c1fe Mon Sep 17 00:00:00 2001
From: Jeffrey Charles <jeff.charles@shopify.com>
Date: Thu, 5 Dec 2024 09:33:04 -0500
Subject: [PATCH] Remove use of shared mutable static (#854)

---
 crates/cli/src/codegen/mod.rs | 36 +++++++++++++++++------------------
 1 file changed, 17 insertions(+), 19 deletions(-)

diff --git a/crates/cli/src/codegen/mod.rs b/crates/cli/src/codegen/mod.rs
index 11d96398..6fb4781f 100644
--- a/crates/cli/src/codegen/mod.rs
+++ b/crates/cli/src/codegen/mod.rs
@@ -33,12 +33,6 @@
 //! as well as ensuring that any features available in the plugin match the
 //! features requsted by the JavaScript bytecode.
 
-// Wizer doesn't provide a good API to set a custom WASI context so we put our
-// WASI context in a static global variable and instruct the closure for
-// getting the WASI context to use a mutable reference to it. There are never
-// concurrent mutable references to the static WASI context so this is safe.
-#![allow(static_mut_refs)]
-
 mod builder;
 use std::{fs, rc::Rc, sync::OnceLock};
 
@@ -52,14 +46,14 @@ use transform::SourceCodeSection;
 use walrus::{
     DataId, DataKind, ExportItem, FunctionBuilder, FunctionId, LocalId, MemoryId, Module, ValType,
 };
-use wasi_common::{pipe::ReadPipe, sync::WasiCtxBuilder, WasiCtx};
 use wasm_opt::{OptimizationOptions, ShrinkLevel};
+use wasmtime_wasi::{pipe::MemoryInputPipe, WasiCtxBuilder};
 use wizer::{Linker, Wizer};
 
 use crate::{js_config::JsConfig, plugins::Plugin, JS};
 use anyhow::Result;
 
-static mut WASI: OnceLock<WasiCtx> = OnceLock::new();
+static STDIN_PIPE: OnceLock<MemoryInputPipe> = OnceLock::new();
 
 pub(crate) enum CodeGenType {
     /// Static code generation.
@@ -147,21 +141,25 @@ impl Generator {
             CodeGenType::Static => {
                 // Copy config JSON into stdin for `initialize_runtime` function.
                 let runtime_config = self.js_runtime_config.to_json()?;
-                unsafe {
-                    WASI.get_or_init(|| {
-                        WasiCtxBuilder::new()
-                            .inherit_stderr()
-                            .inherit_stdout()
-                            .stdin(Box::new(ReadPipe::from(runtime_config)))
-                            .build()
-                    });
-                };
+                STDIN_PIPE
+                    .set(MemoryInputPipe::new(runtime_config))
+                    .unwrap();
                 let wasm = Wizer::new()
                     .init_func("initialize_runtime")
                     .make_linker(Some(Rc::new(move |engine| {
                         let mut linker = Linker::new(engine);
-                        wasi_common::sync::add_to_linker(&mut linker, |_| unsafe {
-                            WASI.get_mut().unwrap()
+                        wasmtime_wasi::preview1::add_to_linker_sync(&mut linker, move |cx| {
+                            // The underlying buffer backing the pipe is an Arc
+                            // so the cloning should be fast.
+                            let config = STDIN_PIPE.get().unwrap().clone();
+                            cx.wasi_ctx = Some(
+                                WasiCtxBuilder::new()
+                                    .stdin(config)
+                                    .inherit_stdout()
+                                    .inherit_stderr()
+                                    .build_p1(),
+                            );
+                            cx.wasi_ctx.as_mut().unwrap()
                         })?;
                         Ok(linker)
                     })))?