From b2cf587a6ba0698320ffcca21b5ec8f6255ec2d2 Mon Sep 17 00:00:00 2001
From: JacobLinCool <jacoblincool@gmail.com>
Date: Sun, 10 Sep 2023 23:58:12 +0800
Subject: [PATCH] feat: add cors headers

---
 src/config.rs      |  9 +++++++++
 src/server/core.rs |  7 +++++++
 src/server/cors.rs | 25 +++++++++++++++++++++++++
 src/server/mod.rs  |  1 +
 4 files changed, 42 insertions(+)
 create mode 100644 src/server/cors.rs

diff --git a/src/config.rs b/src/config.rs
index b6c9b1b..ff92e13 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -29,3 +29,12 @@ pub fn no_cache() -> bool {
         Err(_) => false,
     }
 }
+
+/// Disable CORS by setting the environment variable "NO_CORS" to "1" or "true".
+/// If the variable is not set, a default value of false is returned.
+pub fn no_cors() -> bool {
+    match env::var("NO_CORS") {
+        Ok(val) => val == "1" || val == "true",
+        Err(_) => false,
+    }
+}
diff --git a/src/server/core.rs b/src/server/core.rs
index fb91a70..87360a7 100644
--- a/src/server/core.rs
+++ b/src/server/core.rs
@@ -1,4 +1,5 @@
 use super::compress;
+use super::cors;
 use super::routes;
 use super::version;
 use crate::compile;
@@ -43,6 +44,12 @@ pub fn rocket() -> Rocket<Build> {
 
     let server = server.attach(version::fairing());
 
+    let server = if !no_cors() {
+        server.attach(cors::CORS)
+    } else {
+        server
+    };
+
     if cfg!(debug_assertions) {
         server
     } else {
diff --git a/src/server/cors.rs b/src/server/cors.rs
new file mode 100644
index 0000000..6a7577a
--- /dev/null
+++ b/src/server/cors.rs
@@ -0,0 +1,25 @@
+use rocket::fairing::{Fairing, Info, Kind};
+use rocket::http::Header;
+use rocket::{Request, Response};
+
+pub struct CORS;
+
+#[rocket::async_trait]
+impl Fairing for CORS {
+    fn info(&self) -> Info {
+        Info {
+            name: "Add CORS headers to responses",
+            kind: Kind::Response,
+        }
+    }
+
+    async fn on_response<'r>(&self, _req: &'r Request<'_>, res: &mut Response<'r>) {
+        res.set_header(Header::new("Access-Control-Allow-Origin", "*"));
+        res.set_header(Header::new(
+            "Access-Control-Allow-Methods",
+            "POST, GET, PATCH, OPTIONS",
+        ));
+        res.set_header(Header::new("Access-Control-Allow-Headers", "*"));
+        res.set_header(Header::new("Access-Control-Allow-Credentials", "true"));
+    }
+}
diff --git a/src/server/mod.rs b/src/server/mod.rs
index 0448ef4..719715e 100644
--- a/src/server/mod.rs
+++ b/src/server/mod.rs
@@ -1,5 +1,6 @@
 pub mod compress;
 pub mod core;
+pub mod cors;
 pub mod jwt;
 pub mod routes;
 pub mod system;