From 3d0cef135a204c67063b5d8aec07afac53bef9f5 Mon Sep 17 00:00:00 2001
From: Oliver Stenbom <oliver@stenbom.eu>
Date: Mon, 23 Oct 2023 16:32:39 +0200
Subject: [PATCH] Make iframes work cross domain

---
 linkup-cli/src/local_server.rs   |  5 +---
 linkup-cli/src/services/caddy.rs |  2 +-
 linkup/src/headers.rs            | 14 +++++++++-
 linkup/src/lib.rs                | 45 +++++++++++++++++++++++++++++++-
 4 files changed, 59 insertions(+), 7 deletions(-)

diff --git a/linkup-cli/src/local_server.rs b/linkup-cli/src/local_server.rs
index d162632..2a156dc 100644
--- a/linkup-cli/src/local_server.rs
+++ b/linkup-cli/src/local_server.rs
@@ -190,14 +190,11 @@ async fn linkup_request_handler(
     };
 
     let mut extra_headers = get_additional_headers(&url, &headers, &session_name, &target_service);
-    extra_headers.insert(
-        LinkupHeaderName::Host,
-        Url::parse(&target_service.url).unwrap(),
-    );
 
     // Proxy the request using the destination_url and the merged headers
     let client = reqwest::Client::new();
     headers.extend(&extra_headers);
+    headers.remove(LinkupHeaderName::Host);
 
     let response_result = client
         .request(req.method().clone(), &target_service.url)
diff --git a/linkup-cli/src/services/caddy.rs b/linkup-cli/src/services/caddy.rs
index d5d8fd1..024e3a6 100644
--- a/linkup-cli/src/services/caddy.rs
+++ b/linkup-cli/src/services/caddy.rs
@@ -23,7 +23,7 @@ pub fn start(local_config: &YamlLocalConfig) -> Result<()> {
     let domains: Vec<String> = local_config
         .top_level_domains()
         .iter()
-        .map(|domain| format!("*.{}", domain))
+        .map(|domain| format!("{}, *.{}", domain, domain))
         .collect();
 
     write_caddyfile(&domains)?;
diff --git a/linkup/src/headers.rs b/linkup/src/headers.rs
index 1473749..8cb9892 100644
--- a/linkup/src/headers.rs
+++ b/linkup/src/headers.rs
@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, fmt};
 
 use unicase::UniCase;
 
@@ -79,6 +79,10 @@ impl HeaderMap {
         self.0.extend(iter)
     }
 
+    pub fn remove(&mut self, key: impl Into<UniCase<String>>) -> Option<String> {
+        self.0.remove(&key.into())
+    }
+
     #[cfg(feature = "actix")]
     pub fn from_actix_request(req: &actix_web::HttpRequest) -> Self {
         req.headers().into()
@@ -90,6 +94,14 @@ impl HeaderMap {
     }
 }
 
+impl fmt::Debug for HeaderMap {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_map()
+            .entries(self.0.iter().map(|(k, v)| (k.as_ref(), v.as_str())))
+            .finish()
+    }
+}
+
 #[cfg(feature = "reqwest")]
 impl From<HeaderMap> for reqwest::header::HeaderMap {
     fn from(value: HeaderMap) -> Self {
diff --git a/linkup/src/lib.rs b/linkup/src/lib.rs
index e18c8ab..56bfd51 100644
--- a/linkup/src/lib.rs
+++ b/linkup/src/lib.rs
@@ -107,7 +107,7 @@ pub fn additional_response_headers() -> HeaderMap {
     headers
 }
 
-#[cfg_attr(test, derive(Debug, PartialEq))]
+#[derive(Debug, PartialEq)]
 pub struct TargetService {
     pub name: String,
     pub url: String,
@@ -271,6 +271,8 @@ fn extrace_tracestate(tracestate: &str, linkup_key: String) -> String {
 
 #[cfg(test)]
 mod tests {
+    use std::fmt::format;
+
     use super::*;
 
     const CONF_STR: &str = r#"
@@ -287,6 +289,10 @@ mod tests {
                     }
                 ]
             },
+            {
+                "name": "other-frontend",
+                "location": "http://localhost:5000"
+            },
             {
                 "name": "backend",
                 "rewrites": [
@@ -316,6 +322,10 @@ mod tests {
             {
                 "domain": "api.example.com",
                 "default_service": "backend"
+            },
+            {
+                "domain": "other-example.com",
+                "default_service": "other-frontend"
             }
         ]
     }
@@ -592,4 +602,37 @@ mod tests {
         assert_eq!(target.name, "backend");
         assert_eq!(target.url, "http://localhost:8001/user");
     }
+
+    #[tokio::test]
+    async fn test_iframable() {
+        let string_store = MemoryStringStore::new();
+        let sessions = SessionAllocator::new(&string_store);
+
+        let input_config_value: serde_json::Value = serde_json::from_str(CONF_STR).unwrap();
+        let input_config: Session = input_config_value.try_into().unwrap();
+
+        let name = sessions
+            .store_session(input_config, NameKind::Animal, "".to_string())
+            .await
+            .unwrap();
+
+        let mut headers = HeaderMap::new();
+        headers.insert(HeaderName::Referer, format!("{}.example.com", name));
+
+        let (name, config) = sessions
+            .get_request_session("other-example.com", &headers)
+            .await
+            .unwrap();
+
+        let target = get_target_service(
+            "http://other-example.com",
+            &headers,
+            &config,
+            &name,
+        )
+        .unwrap();
+
+        assert_eq!(target.name, "other-frontend");
+        assert_eq!(target.url, "http://localhost:5000/");
+    }
 }