Multi-domain fix #502

Author: joepio

server/src/config.rs

@@ -209,22 +209,30 @@ impl Config {
     /// Returns the server URL for a given request.
     /// If multi-tenancy is enabled and the host matches a subdomain of the base domain, it returns the host URL.
     pub fn get_server_url_for_request(&self, req: &actix_web::HttpRequest) -> String {
-        if let Some(base) = &self.base_domain {
-            if let Some(host) = req.head().headers.get("Host") {
-                if let Ok(host_str) = host.to_str() {
-                    // Remove port if present
-                    let domain = host_str.split(':').next().unwrap_or(host_str);
-                    if domain.ends_with(base) {
-                        let schema =
-                            if let Some(proto) = req.head().headers.get("X-Forwarded-Proto") {
-                                proto.to_str().unwrap_or("http")
-                            } else if self.opts.https {
-                                "https"
-                            } else {
-                                "http"
-                            };
-                        return format!("{}://{}", schema, host_str);
-                    }
+        let host_header = req
+            .head()
+            .headers
+            .get("X-Forwarded-Host")
+            .or_else(|| req.head().headers.get("Host"));
+
+        if let Some(host) = host_header {
+            if let Ok(host_str) = host.to_str() {
+                let domain = host_str.split(':').next().unwrap_or(host_str);
+                let allowed = if let Some(base) = &self.base_domain {
+                    domain.ends_with(base)
+                } else {
+                    true
+                };
+
+                if allowed {
+                    let schema = if let Some(proto) = req.head().headers.get("X-Forwarded-Proto") {
+                        proto.to_str().unwrap_or("http")
+                    } else if self.opts.https {
+                        "https"
+                    } else {
+                        "http"
+                    };
+                    return format!("{}://{}", schema, host_str);
                 }
             }
         }

server/src/handlers/single_page_app.rs

@@ -9,14 +9,15 @@ use actix_web::HttpResponse;
 pub async fn single_page(
     appstate: actix_web::web::Data<AppState>,
     path: actix_web::web::Path<String>,
+    req: actix_web::HttpRequest,
 ) -> AtomicServerResult<HttpResponse> {
+    let server_url = appstate.config.get_server_url_for_request(&req);
+    let store = appstate.store.clone_with_url(server_url.clone());
     let template = include_str!("../../assets_tmp/index.html");
     let csp_nonce = generate_nonce().map_err(|_e| "Failed to generate nonce")?;
-    let subject = format!("{}/{}", appstate.store.get_server_url()?, path);
+    let subject = format!("{}/{}", server_url, path);
     let meta_tags: MetaTags = if let Ok(resource_response) =
-        appstate
-            .store
-            .get_resource_extended(&subject, true, &ForAgent::Public)
+        store.get_resource_extended(&subject, true, &ForAgent::Public)
     {
         resource_response.into()
     } else {

server/src/handlers/web_sockets.rs

@@ -41,13 +41,16 @@ pub async fn web_socket_handler(
     )?;
     tracing::debug!("Starting websocket for {}", for_agent);
 
+    let server_url = appstate.config.get_server_url_for_request(&req);
+    let store = appstate.store.clone_with_url(server_url);
+
     let result = ws::start(
         WebSocketConnection::new(
             appstate.commit_monitor.clone(),
             appstate.y_sync_broadcaster.clone(),
             for_agent,
             // We need to make sure this is easily clone-able
-            appstate.store.clone(),
+            store,
         ),
         &req,
         stream,