From 4f6f510bd47e25096c2fac07c66e1a6874623973 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20Garc=C3=ADa?=
 <dani-garcia@users.noreply.github.com>
Date: Thu, 12 Jul 2018 23:28:01 +0200
Subject: [PATCH] Improve domain detection, should fix attachment problems.
 Otherwise, set the `DOMAIN` env variable to the correct domain

---
 .env        |  6 ++++++
 src/auth.rs | 32 +++++++++++++++++++++++++-------
 src/main.rs |  6 +++++-
 3 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/.env b/.env
index 152f4ce..0c231d7 100644
--- a/.env
+++ b/.env
@@ -27,6 +27,12 @@
 ## The change only applies when the password is changed
 # PASSWORD_ITERATIONS=100000
 
+## Domain settings
+## The domain must match the address from where you access the server
+## Unless you are using U2F, or having problems with attachments not downloading, there is no need to change this
+## For U2F to work, the server must use HTTPS, you can use Let's Encrypt for free certs
+# DOMAIN=https://bw.domain.tld:8443
+
 ## Rocket specific settings, check Rocket documentation to learn more
 # ROCKET_ENV=staging
 # ROCKET_ADDRESS=0.0.0.0 # Enable this to test mobile app
diff --git a/src/auth.rs b/src/auth.rs
index 0e2a440..a2261ed 100644
--- a/src/auth.rs
+++ b/src/auth.rs
@@ -109,14 +109,32 @@ impl<'a, 'r> FromRequest<'a, 'r> for Headers {
     fn from_request(request: &'a Request<'r>) -> request::Outcome<Self, Self::Error> {
         let headers = request.headers();
 
+        println!("{:#?}", headers);
+
         // Get host
-        let host = match headers.get_one("Host") {
-            Some(host) => {
-                use std::env;
-                let protocol = if env::var("ROCKET_TLS").is_ok() {"https"} else {"http"};
-                format!("{}://{}", protocol, host)
-            },
-            _ => String::new()
+        let host = if CONFIG.domain_set {
+            CONFIG.domain.clone()
+        } else if let Some(referer) = headers.get_one("Referer") {
+            referer.to_string()
+        } else {   
+            // Try to guess from the headers
+            use std::env;
+
+            let protocol = if let Some(proto) = headers.get_one("X-Forwarded-Proto") {
+                proto
+            } else if env::var("ROCKET_TLS").is_ok() {
+                "https"
+            } else {
+                "http"
+            };
+
+            let host = if let Some(host) = headers.get_one("Host") {
+                host
+            } else {
+                ""
+            };
+
+            format!("{}://{}", protocol, host)
         };
 
         // Get access_token
diff --git a/src/main.rs b/src/main.rs
index 48f8a47..17b6379 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -165,6 +165,7 @@ pub struct Config {
     signups_allowed: bool,
     password_iterations: i32,
     domain: String,
+    domain_set: bool,
 }
 
 impl Config {
@@ -174,6 +175,8 @@ impl Config {
         let df = env::var("DATA_FOLDER").unwrap_or("data".into());
         let key = env::var("RSA_KEY_FILENAME").unwrap_or(format!("{}/{}", &df, "rsa_key"));
 
+        let domain = env::var("DOMAIN");
+
         Config {
             database_url: env::var("DATABASE_URL").unwrap_or(format!("{}/{}", &df, "db.sqlite3")),
             icon_cache_folder: env::var("ICON_CACHE_FOLDER").unwrap_or(format!("{}/{}", &df, "icon_cache")),
@@ -189,7 +192,8 @@ impl Config {
             local_icon_extractor: util::parse_option_string(env::var("LOCAL_ICON_EXTRACTOR").ok()).unwrap_or(false),
             signups_allowed: util::parse_option_string(env::var("SIGNUPS_ALLOWED").ok()).unwrap_or(true),
             password_iterations: util::parse_option_string(env::var("PASSWORD_ITERATIONS").ok()).unwrap_or(100_000),
-            domain: env::var("DOMAIN").unwrap_or("https://localhost".into()),
+            domain_set: domain.is_ok(),
+            domain: domain.unwrap_or("http://localhost".into()),
         }
     }
 }