From 33fa56c12b80d091a45fdffac0791c46f760e2e2 Mon Sep 17 00:00:00 2001
From: Tom Lienard <tom.lienrd@gmail.com>
Date: Mon, 15 May 2023 09:33:35 +0200
Subject: [PATCH] feat(runtime): allow `export as` for handler function (#868)

---
 .changeset/nasty-ads-buy.md           |  8 ++++++++
 crates/runtime/tests/runtime.rs       | 19 +++++++++++++++++++
 crates/runtime_isolate/src/lib.rs     | 24 ++++++++++++++++++++----
 crates/runtime_isolate/src/options.rs |  6 ++----
 packages/js-runtime/src/index.ts      |  4 ++--
 5 files changed, 51 insertions(+), 10 deletions(-)
 create mode 100644 .changeset/nasty-ads-buy.md

diff --git a/.changeset/nasty-ads-buy.md b/.changeset/nasty-ads-buy.md
new file mode 100644
index 000000000..be166c79f
--- /dev/null
+++ b/.changeset/nasty-ads-buy.md
@@ -0,0 +1,8 @@
+---
+'@lagon/cli': patch
+'@lagon/runtime': patch
+'@lagon/serverless': patch
+'@lagon/js-runtime': patch
+---
+
+Allow `export as` for handler function
diff --git a/crates/runtime/tests/runtime.rs b/crates/runtime/tests/runtime.rs
index 280494247..6b2fb32cd 100644
--- a/crates/runtime/tests/runtime.rs
+++ b/crates/runtime/tests/runtime.rs
@@ -21,6 +21,25 @@ async fn execute_function() {
     );
 }
 
+#[tokio::test]
+async fn execute_function_export_as() {
+    utils::setup();
+    let (send, receiver) = utils::create_isolate(IsolateOptions::new(
+        "function hello() {
+    return new Response('Hello world');
+}
+
+export { hello as handler }"
+            .into(),
+    ));
+    send(Request::default());
+
+    assert_eq!(
+        receiver.recv_async().await.unwrap().as_response(),
+        Response::from("Hello world")
+    );
+}
+
 #[tokio::test]
 async fn execute_function_twice() {
     utils::setup();
diff --git a/crates/runtime_isolate/src/lib.rs b/crates/runtime_isolate/src/lib.rs
index e2ab732f9..14c07cbc1 100644
--- a/crates/runtime_isolate/src/lib.rs
+++ b/crates/runtime_isolate/src/lib.rs
@@ -101,7 +101,8 @@ impl Heartbeat {
 pub struct Isolate {
     options: IsolateOptions,
     isolate: Option<v8::OwnedIsolate>,
-    handler: Option<v8::Global<v8::Function>>,
+    master_handler: Option<v8::Global<v8::Function>>,
+    handler: Option<v8::Global<v8::Value>>,
     compilation_error: Option<String>,
     stream_receiver: flume::Receiver<(u32, StreamResult)>,
     termination_result: Arc<RwLock<Option<RunResult>>>,
@@ -200,6 +201,7 @@ impl Isolate {
         let mut this = Self {
             options,
             isolate: Some(isolate),
+            master_handler: None,
             handler: None,
             compilation_error: None,
             stream_receiver,
@@ -367,6 +369,13 @@ impl Isolate {
                 }
 
                 if !self.options.snapshot {
+                    let namespace = module.get_module_namespace().to_object(try_catch).unwrap();
+                    let handler_key = v8_string(try_catch, "handler");
+                    let handler = namespace.get(try_catch, handler_key.into()).unwrap();
+                    let handler = v8::Global::new(try_catch, handler);
+
+                    self.handler = Some(handler);
+
                     let global = global.open(try_catch);
                     let global = global.global(try_catch);
                     let handler_key = v8_string(try_catch, "masterHandler");
@@ -374,7 +383,7 @@ impl Isolate {
                     let handler = v8::Local::<v8::Function>::try_from(handler).unwrap();
                     let handler = v8::Global::new(try_catch, handler);
 
-                    self.handler = Some(handler);
+                    self.master_handler = Some(handler);
                 }
             }
             None => {
@@ -400,8 +409,11 @@ impl Isolate {
                 );
                 let try_catch = &mut v8::TryCatch::new(scope);
 
+                let master_handler = self.master_handler.as_ref().unwrap();
+                let master_handler = master_handler.open(try_catch);
+
                 let handler = self.handler.as_ref().unwrap();
-                let handler = handler.open(try_catch);
+                let handler = v8::Local::new(try_catch, handler);
 
                 let global = global.open(try_catch);
                 let global = global.global(try_catch);
@@ -422,7 +434,11 @@ impl Isolate {
                     },
                 );
 
-                match handler.call(try_catch, global.into(), &[id.into(), request.into()]) {
+                match master_handler.call(
+                    try_catch,
+                    global.into(),
+                    &[id.into(), handler, request.into()],
+                ) {
                     Some(response) => {
                         let promise = v8::Local::<v8::Promise>::try_from(response)
                             .expect("Handler did not return a promise");
diff --git a/crates/runtime_isolate/src/options.rs b/crates/runtime_isolate/src/options.rs
index 54e33efda..c207094ec 100644
--- a/crates/runtime_isolate/src/options.rs
+++ b/crates/runtime_isolate/src/options.rs
@@ -127,8 +127,7 @@ impl IsolateOptions {
                     scope,
                     &format!(
                         r"{environment_variables}
-{code}
-globalThis.handler = handler;"
+{code}"
                     ),
                 ),
                 environment_variables.lines().count() + 1,
@@ -146,8 +145,7 @@ globalThis.handler = handler;"
                     &format!(
                         r"{JS_RUNTIME}
 {environment_variables}
-{code}
-globalThis.handler = handler;"
+{code}"
                     ),
                 ),
                 JS_RUNTIME.lines().count() + environment_variables.lines().count() + 2,
diff --git a/packages/js-runtime/src/index.ts b/packages/js-runtime/src/index.ts
index e4d87e0db..71007b5f7 100644
--- a/packages/js-runtime/src/index.ts
+++ b/packages/js-runtime/src/index.ts
@@ -113,9 +113,9 @@ declare global {
     TEXT_DECODER: TextDecoder;
   };
   var __storage__: Map<AsyncContext, unknown>;
-  var handler: (request: Request) => Promise<Response>;
   var masterHandler: (
     id: number,
+    handler: (request: Request) => Promise<Response>,
     request: {
       i: string;
       m: RequestInit['method'];
@@ -141,7 +141,7 @@ declare global {
   }
 }
 
-globalThis.masterHandler = async (id, request) => {
+globalThis.masterHandler = async (id, handler, request) => {
   if (typeof handler !== 'function') {
     throw new Error('Handler function is not defined or is not a function');
   }