From bbbf9d9fe908682e51d1b03cbb270e1f6e1be525 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lu=C3=ADs=20Duarte?= <lduarte@cloudflare.com>
Date: Mon, 23 Dec 2024 15:46:45 +0000
Subject: [PATCH] feat: make worker-build support custom JS shims

This allows for users to provide their own JS shim and allowing them to
customize behaviour: panic handling, wasm coredumps, and so on...

It defaults to the embedded shim in the binary to avoid breaking
changes,
---
 worker-build/src/main.rs | 25 ++++++++++++++++++++++---
 1 file changed, 22 insertions(+), 3 deletions(-)

diff --git a/worker-build/src/main.rs b/worker-build/src/main.rs
index 2de00f28..31d2eb09 100644
--- a/worker-build/src/main.rs
+++ b/worker-build/src/main.rs
@@ -2,7 +2,7 @@
 
 use std::{
     env::{self, VarError},
-    fs::{self, File},
+    fs::{self, read_to_string, File},
     io::{Read, Write},
     path::{Path, PathBuf},
     process::{Command, Stdio},
@@ -38,10 +38,29 @@ pub fn main() -> Result<()> {
     create_worker_dir()?;
     copy_generated_code_to_worker_dir()?;
 
+    let custom_shim = match env::var("CUSTOM_SHIM") {
+        Ok(path) => {
+            let path = Path::new(&path).to_owned();
+
+            if !path.exists() {
+                None
+            } else {
+                println!("Using custom shim from {}", path.display());
+                Some(path)
+            }
+        }
+        Err(_) => None,
+    };
+
+    let shim_template = match custom_shim {
+        Some(path) => read_to_string(path)?,
+        None => SHIM_TEMPLATE.to_owned(),
+    };
+
     let shim = if env::var("RUN_TO_COMPLETION").is_ok() {
-        SHIM_TEMPLATE.replace("$WAIT_UNTIL_RESPONSE", "this.ctx.waitUntil(response);")
+        shim_template.replace("$WAIT_UNTIL_RESPONSE", "this.ctx.waitUntil(response);")
     } else {
-        SHIM_TEMPLATE.replace("$WAIT_UNTIL_RESPONSE", "")
+        shim_template.replace("$WAIT_UNTIL_RESPONSE", "")
     };
 
     write_string_to_file(worker_path("shim.js"), shim)?;