From 21860484d1ee62bdbed1e3142c165478526837e3 Mon Sep 17 00:00:00 2001
From: Chris Collier <czcollier@google.com>
Date: Tue, 12 Nov 2024 22:05:03 +0000
Subject: [PATCH] First pass at implementing cors

Change-Id: I4da216d52191baabab9f63cc3ee6ae749add906f
---
 mesop/features/page.py              | 5 ++++-
 mesop/runtime/runtime.py            | 2 ++
 mesop/security/cors.py              | 5 +++++
 mesop/server/static_file_serving.py | 6 ++++++
 4 files changed, 17 insertions(+), 1 deletion(-)
 create mode 100644 mesop/security/cors.py

diff --git a/mesop/features/page.py b/mesop/features/page.py
index d6f209c5..1f87c6f1 100644
--- a/mesop/features/page.py
+++ b/mesop/features/page.py
@@ -4,6 +4,7 @@
 
 from mesop.runtime import OnLoadHandler, PageConfig, runtime
 from mesop.security.security_policy import SecurityPolicy
+from mesop.security.cors import CORS
 
 
 def page(
@@ -12,6 +13,7 @@ def page(
   title: str | None = None,
   stylesheets: list[str] | None = None,
   security_policy: SecurityPolicy | None = None,
+  cors: CORS | None = None,
   on_load: OnLoadHandler | None = None,
 ) -> Callable[[Callable[[], None]], Callable[[], None]]:
   """Defines a page in a Mesop application.
@@ -42,7 +44,8 @@ def wrapper() -> None:
         page_fn=wrapper,
         title=title or f"Mesop: {path}",
         stylesheets=deepcopy(stylesheets or default_stylesheets),
-        security_policy=deepcopy(security_policy)
+        security_policy=deepcopy(security_policy),
+        cors = deepcopy(cors)
         if security_policy
         else SecurityPolicy(),
         on_load=on_load,
diff --git a/mesop/runtime/runtime.py b/mesop/runtime/runtime.py
index f17571ef..60a78e09 100644
--- a/mesop/runtime/runtime.py
+++ b/mesop/runtime/runtime.py
@@ -10,6 +10,7 @@
 from mesop.exceptions import MesopDeveloperException, MesopUserException
 from mesop.key import Key
 from mesop.security.security_policy import SecurityPolicy
+from mesop.security.cors import CORS
 from mesop.utils.backoff import exponential_backoff
 from mesop.warn import warn
 
@@ -33,6 +34,7 @@ class PageConfig:
   title: str
   stylesheets: list[str]
   security_policy: SecurityPolicy
+  cors: CORS
   on_load: OnLoadHandler | None
 
 
diff --git a/mesop/security/cors.py b/mesop/security/cors.py
new file mode 100644
index 00000000..077ce7bb
--- /dev/null
+++ b/mesop/security/cors.py
@@ -0,0 +1,5 @@
+from dataclasses import dataclass, field
+
+@dataclass(kw_only=True)
+class CORS:
+  allowed_origins: list[str] = field(default_factory=list)
\ No newline at end of file
diff --git a/mesop/server/static_file_serving.py b/mesop/server/static_file_serving.py
index f0899cf4..6a1f734a 100644
--- a/mesop/server/static_file_serving.py
+++ b/mesop/server/static_file_serving.py
@@ -228,6 +228,12 @@ def add_security_headers(response: Response):
         "report-uri": "/__csp__",
       }
     )
+
+    if page_config and page_config.cors:
+      cors = page_config.cors
+      if cors.allowed_origins:
+        response.headers["Access-Control-Allow-Origin"] = ", ".join(cors.allowed_origins)
+
     if page_config and page_config.stylesheets:
       csp["style-src"] += " " + " ".join(
         [sanitize_url_for_csp(url) for url in page_config.stylesheets]