From 5c5d9d0f043e73bed2081e9d52e03159acd57d97 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1=C5=A1=20Ml=C3=A1dek?= <t@mldk.cz>
Date: Sun, 23 Oct 2022 12:53:33 +0200
Subject: [PATCH] feat: add --allow-hosts CLI option

---
 src/main.rs | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/src/main.rs b/src/main.rs
index caa6ca0..4e4a380 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -116,6 +116,14 @@ fn main() -> Result<()> {
                 .long("key")
                 .env("UPEND_KEY")
                 .help("Authentication key users must supply."),
+        )
+        .arg(
+            Arg::with_name("ALLOW_HOST")
+                .takes_value(true)
+                .multiple(true)
+                .number_of_values(1)
+                .long("allow-host")
+                .help("Allowed host/domain name the API can serve."),
         );
 
     let matches = app.get_matches();
@@ -234,16 +242,29 @@ fn main() -> Result<()> {
 
     let mut cnt = 0;
     let ui_path = ui_path.ok();
+    let allowed_origins: Vec<_> = if let Some(matches) = matches.values_of("ALLOW_HOST") {
+        matches.map(String::from).collect()
+    } else {
+        vec![]
+    };
     let server = loop {
         let state = state.clone();
         let ui_path = ui_path.clone();
+        let allowed_origins = allowed_origins.clone();
 
         let server = HttpServer::new(move || {
+            let allowed_origins = allowed_origins.clone();
+
             let cors = Cors::default()
                 .allowed_origin("http://localhost")
                 .allowed_origin_fn(|origin, _req_head| {
                     origin.as_bytes().starts_with(b"http://localhost:")
                 })
+                .allowed_origin_fn(move |origin, _req_head| {
+                    allowed_origins
+                        .iter()
+                        .any(|allowed_origin| *allowed_origin == "*" || origin == allowed_origin)
+                })
                 .allow_any_method();
 
             let app = App::new()