chore: improve cors configuration

gabotechs · gabotechs · commit 81d3df54b912 · 2023-06-20T18:22:49.000+02:00
diff --git a/server/src/lib.rs b/server/src/lib.rs
@@ -9,7 +9,6 @@ mod _test_tools;
 
 mod body;
 mod gateway_callbacks;
-mod route_cors;
 mod route_gateway;
 mod secret_getter;
 mod server;
diff --git a/server/src/route_cors.rs b/server/src/route_cors.rs
diff --git a/server/src/server.rs b/server/src/server.rs
@@ -6,8 +6,12 @@ use std::u16;
 use anyhow::Result;
 use async_trait::async_trait;
 use hyper::client::HttpConnector;
+use hyper::header::{
+    ACCESS_CONTROL_ALLOW_HEADERS, ACCESS_CONTROL_ALLOW_METHODS, ACCESS_CONTROL_ALLOW_ORIGIN,
+};
+use hyper::http::HeaderValue;
 use hyper::service::service_fn;
-use hyper::{Body, Method, Request, Response};
+use hyper::{Body, Method, Request, Response, StatusCode};
 use hyper_tls::HttpsConnector;
 use tokio::net::TcpListener;
 use tracing::{error, info};
@@ -16,6 +20,13 @@ use crate::gateway_callbacks::{CallbackResult, OnRequest, OnSuccess};
 use crate::secret_getter::SecretGetter;
 use crate::{BytesTransferredInfo, OnBytesTransferred};
 
+fn ok() -> Response<Body> {
+    Response::builder()
+        .status(StatusCode::OK)
+        .body(Body::empty())
+        .unwrap()
+}
+
 pub struct SignwayServer<T: SecretGetter + 'static> {
     pub port: u16,
     pub secret_getter: T,
@@ -24,6 +35,9 @@ pub struct SignwayServer<T: SecretGetter + 'static> {
     pub on_bytes_transferred: Arc<dyn OnBytesTransferred>,
     pub(crate) monitor_bytes: bool,
     pub(crate) client: hyper::Client<HttpsConnector<HttpConnector>, Body>,
+    pub(crate) access_control_allow_origin: HeaderValue,
+    pub(crate) access_control_allow_methods: HeaderValue,
+    pub(crate) access_control_allow_headers: HeaderValue,
 }
 
 pub(crate) struct NoneCallback;
@@ -60,6 +74,9 @@ impl<T: SecretGetter> SignwayServer<T> {
             on_success: Box::new(NoneCallback {}),
             on_bytes_transferred: Arc::new(NoneCallback {}),
             monitor_bytes: false,
+            access_control_allow_origin: HeaderValue::from_static("*"),
+            access_control_allow_headers: HeaderValue::from_static("*"),
+            access_control_allow_methods: HeaderValue::from_static("*"),
             client,
         }
     }
@@ -74,6 +91,9 @@ impl<T: SecretGetter> SignwayServer<T> {
             on_success: Box::new(NoneCallback {}),
             on_bytes_transferred: Arc::new(NoneCallback {}),
             monitor_bytes: false,
+            access_control_allow_origin: HeaderValue::from_static("*"),
+            access_control_allow_headers: HeaderValue::from_static("*"),
+            access_control_allow_methods: HeaderValue::from_static("*"),
             client,
         }
     }
@@ -94,6 +114,38 @@ impl<T: SecretGetter> SignwayServer<T> {
         self
     }
 
+    pub fn access_control_allow_origin(mut self, value: &str) -> Result<Self> {
+        self.access_control_allow_origin = value.parse()?;
+        Ok(self)
+    }
+
+    pub fn access_control_allow_methods(mut self, value: &str) -> Result<Self> {
+        self.access_control_allow_methods = value.parse()?;
+        Ok(self)
+    }
+
+    pub fn access_control_allow_headers(mut self, value: &str) -> Result<Self> {
+        self.access_control_allow_headers = value.parse()?;
+        Ok(self)
+    }
+
+    fn with_cors_headers<B>(&self, mut res: Response<B>) -> Response<B> {
+        let h = res.headers_mut();
+        h.insert(
+            ACCESS_CONTROL_ALLOW_ORIGIN,
+            self.access_control_allow_origin.clone(),
+        );
+        h.insert(
+            ACCESS_CONTROL_ALLOW_METHODS,
+            self.access_control_allow_methods.clone(),
+        );
+        h.insert(
+            ACCESS_CONTROL_ALLOW_HEADERS,
+            self.access_control_allow_headers.clone(),
+        );
+        res
+    }
+
     pub async fn start(self) -> Result<()> {
         let in_addr: SocketAddr = ([0, 0, 0, 0], self.port).into();
 
@@ -109,10 +161,15 @@ impl<T: SecretGetter> SignwayServer<T> {
             let service = service_fn(move |req| {
                 let arc_self = arc_self.clone();
                 async move {
-                    if req.method() == Method::OPTIONS {
-                        Self::route_cors(&arc_self, req).await
+                    let res = if req.method() == Method::OPTIONS {
+                        Ok(ok())
+                    } else {
+                        arc_self.route_gateway(req).await
+                    };
+                    if let Ok(res) = res {
+                        Ok(arc_self.with_cors_headers(res))
                     } else {
-                        Self::route_gateway(&arc_self, req).await
+                        res
                     }
                 }
             });
@@ -192,6 +249,13 @@ mod tests {
 
         let status = response.status();
         assert_eq!(status, StatusCode::OK);
+        assert_eq!(
+            response
+                .headers()
+                .get("access-control-allow-origin")
+                .unwrap_or(&HeaderValue::from_str("NONE").unwrap()),
+            "*"
+        );
     }
 
     #[tokio::test]
diff --git a/src/main.rs b/src/main.rs
@@ -12,7 +12,7 @@ use signway_server::{
     SecretGetterResult, SignwayServer,
 };
 
-#[derive(Parser, Debug)]
+#[derive(Parser, Debug, Clone)]
 #[command(author, version, about, long_about = None)]
 struct Args {
     #[arg(help = "access id that is expected to sign urls for this server")]
@@ -32,6 +32,24 @@ struct Args {
         help = "disables the bytes transferred monitoring. This feature is experimental, because hyper does not put things easy for tracking IO results in the responses, and the current implementation might have some performance implications. https://github.com/hyperium/hyper/issues/2181"
     )]
     no_bytes_monitor: bool,
+
+    #[arg(
+        long,
+        help = "sets the Access-Control-Allow-Origin that will be answered in each request"
+    )]
+    access_control_allow_origin: Option<String>,
+
+    #[arg(
+        long,
+        help = "sets the Access-Control-Allow-Methods that will be answered in each request"
+    )]
+    access_control_allow_methods: Option<String>,
+
+    #[arg(
+        long,
+        help = "sets the Access-Control-Allow-Headers that will be answered in each request"
+    )]
+    access_control_allow_headers: Option<String>,
 }
 
 struct Config {
@@ -98,14 +116,24 @@ impl OnBytesTransferred for BytesTransferredLogger {
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
-    let args: Args = Args::parse();
-    let disable_monitoring = args.no_bytes_monitor;
-    let config: Config = args.try_into()?;
     tracing_subscriber::fmt().json().init();
+
+    let args: Args = Args::parse();
+    let config: Config = args.clone().try_into()?;
     let mut server = SignwayServer::from_env(config);
-    if !disable_monitoring {
+
+    if !args.no_bytes_monitor {
         server = server.on_bytes_transferred(BytesTransferredLogger {});
     }
+    if let Some(value) = args.access_control_allow_headers {
+        server = server.access_control_allow_headers(&value)?;
+    }
+    if let Some(value) = args.access_control_allow_methods {
+        server = server.access_control_allow_methods(&value)?;
+    }
+    if let Some(value) = args.access_control_allow_origin {
+        server = server.access_control_allow_origin(&value)?;
+    }
 
     tokio::select! {
         result = server.start() => {
