From dbcf752e67e60d54d94c3571569bcf8baed2016b Mon Sep 17 00:00:00 2001
From: Narayan Bhat <48803246+Narayanbhat166@users.noreply.github.com>
Date: Mon, 28 Aug 2023 21:02:37 +0530
Subject: [PATCH] feat: add `--cors` flag to enable cors for workers (#202)

---
 crates/server/src/handlers/worker.rs | 20 ++++++++++++++++++--
 crates/server/src/lib.rs             |  6 +++++-
 src/main.rs                          |  5 +++++
 3 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/crates/server/src/handlers/worker.rs b/crates/server/src/handlers/worker.rs
index 75369e15..12453017 100644
--- a/crates/server/src/handlers/worker.rs
+++ b/crates/server/src/handlers/worker.rs
@@ -5,13 +5,15 @@ use super::{assets::handle_assets, not_found::handle_not_found};
 use crate::DataConnectors;
 use actix_web::{
     http::StatusCode,
-    web::{Bytes, Data},
+    web::{self, Bytes, Data},
     HttpRequest, HttpResponse,
 };
 use std::{fs::File, io::Write, sync::RwLock};
 use wws_router::Routes;
 use wws_worker::io::WasmOutput;
 
+const CORS_HEADER: &str = "Access-Control-Allow-Origin";
+
 /// Process an HTTP request by passing it to the right Runner. The Runner
 /// will prepare the WASI environment and call the Wasm module with the data.
 ///
@@ -29,7 +31,11 @@ use wws_worker::io::WasmOutput;
 ///
 /// For these reasons, we are selecting the right handler at this point and not
 /// allowing Actix to select it for us.
-pub async fn handle_worker(req: HttpRequest, body: Bytes) -> HttpResponse {
+pub async fn handle_worker(
+    req: HttpRequest,
+    body: Bytes,
+    cors_origins: web::Data<Option<Vec<String>>>,
+) -> HttpResponse {
     let routes = req.app_data::<Data<Routes>>().unwrap();
     let stderr_file = req.app_data::<Data<Option<File>>>().unwrap();
     let data_connectors = req
@@ -97,6 +103,16 @@ pub async fn handle_worker(req: HttpRequest, body: Bytes) -> HttpResponse {
         // Default content type
         builder.insert_header(("Content-Type", "text/html"));
 
+        // Check if cors config has any origins to register
+        if let Some(origins) = cors_origins.as_ref() {
+            // Check if worker has overridden the header, if not
+            if !handler_result.headers.contains_key(CORS_HEADER) {
+                // insert those origins in 'Access-Control-Allow-Origin' header
+                let header_value = origins.join(",");
+                builder.insert_header((CORS_HEADER, header_value));
+            }
+        }
+
         for (key, val) in handler_result.headers.iter() {
             // Note that QuickJS is replacing the "-" character
             // with "_" on property keys. Here, we rollback it
diff --git a/crates/server/src/lib.rs b/crates/server/src/lib.rs
index ecce2418..94b027dc 100644
--- a/crates/server/src/lib.rs
+++ b/crates/server/src/lib.rs
@@ -39,6 +39,7 @@ pub async fn serve(
     port: u16,
     panel: bool,
     stderr: Option<&Path>,
+    cors_origins: Option<Vec<String>>,
 ) -> Result<Server> {
     // Initializes the data connectors. For now, just KV
     let data = Data::new(RwLock::new(DataConnectors::default()));
@@ -59,6 +60,8 @@ pub async fn serve(
         stderr_file = Data::new(None);
     }
 
+    let cors_data = Data::new(cors_origins);
+
     let server = HttpServer::new(move || {
         let mut app = App::new()
             // enable logger
@@ -68,7 +71,8 @@ pub async fn serve(
             .app_data(Data::clone(&routes))
             .app_data(Data::clone(&data))
             .app_data(Data::clone(&root_path))
-            .app_data(Data::clone(&stderr_file));
+            .app_data(Data::clone(&stderr_file))
+            .app_data(Data::clone(&cors_data));
 
         // Configure panel
         if panel {
diff --git a/src/main.rs b/src/main.rs
index 999ac0a1..80c0c9cf 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -68,6 +68,10 @@ pub struct Args {
     /// Manage language runtimes in your project
     #[command(subcommand)]
     commands: Option<Main>,
+
+    /// CORS headers to add to all workers if not already set by the worker
+    #[arg(long)]
+    cors: Option<Vec<String>>,
 }
 
 #[actix_web::main]
@@ -198,6 +202,7 @@ async fn main() -> std::io::Result<()> {
             args.port,
             args.enable_panel,
             None,
+            args.cors,
         )
         .await
         .map_err(|err| Error::new(ErrorKind::AddrInUse, err))?;