fixed #1093 - allow multiple return domains for SSO, prefer host header over external_host

warpgate-common/src/config/mod.rs

@@ -5,7 +5,7 @@ use std::path::PathBuf;
 use std::time::Duration;
 
 use defaults::*;
-use poem::http::{self, uri};
+use poem::http::uri;
 use poem_openapi::{Object, Union};
 use serde::{Deserialize, Serialize};
 pub use target::*;
@@ -381,7 +381,7 @@ pub struct WarpgateConfig {
 }
 
 impl WarpgateConfig {
-    pub fn _external_host_from_config(&self) -> Option<(Scheme, String, Option<u16>)> {
+    pub fn external_host_from_config(&self) -> Option<(Scheme, String, Option<u16>)> {
         if let Some(external_host) = self.store.external_host.as_ref() {
             #[allow(clippy::unwrap_used)]
             let external_host = external_host.split(":").next().unwrap();
@@ -399,8 +399,8 @@ impl WarpgateConfig {
         }
     }
 
-    // Extract external host:port from request headers
-    pub fn _external_host_from_request(
+    /// Extract external host:port from request headers
+    pub fn external_host_from_request(
         &self,
         request: &poem::Request,
     ) -> Option<(Scheme, String, Option<u16>)> {
@@ -410,11 +410,10 @@ impl WarpgateConfig {
         // Try the Host header first
         scheme = request.uri().scheme().cloned().unwrap_or(scheme);
 
-        if let Some(host_header) = request.header(http::header::HOST).map(|x| x.to_string()) {
-            if let Ok(host_port) = Url::parse(&format!("https://{host_header}/")) {
-                host = host_port.host_str().map(Into::into).or(host);
-                port = host_port.port();
-            }
+        let original_url = request.original_uri();
+        if let Some(original_host) = original_url.host() {
+            host = Some(original_host.to_string());
+            port = original_url.port().map(|x| x.as_u16());
         }
 
         // But prefer X-Forwarded-* headers if enabled
@@ -443,14 +442,24 @@ impl WarpgateConfig {
     pub fn construct_external_url(
         &self,
         for_request: Option<&poem::Request>,
+        domain_whitelist: Option<&[String]>,
     ) -> Result<Url, WarpgateError> {
-        let Some((scheme, host, port)) = self
-            ._external_host_from_config()
-            .or(for_request.and_then(|r| self._external_host_from_request(r)))
+        let Some((scheme, host, port)) = for_request
+            .and_then(|r| self.external_host_from_request(r))
+            .or(self.external_host_from_config())
         else {
-            return Err(WarpgateError::ExternalHostNotSet);
+            return Err(WarpgateError::ExternalHostUnknown);
         };
 
+        if let Some(list) = domain_whitelist {
+            if !list.contains(&host) {
+                return Err(WarpgateError::ExternalHostNotWhitelisted(
+                    host.clone(),
+                    list.iter().map(|x| x.to_string()).collect(),
+                ));
+            }
+        }
+
         let mut url = format!("{scheme}://{host}");
         if let Some(port) = port {
             // can't `match` `Scheme`

warpgate-common/src/error.rs

@@ -21,8 +21,10 @@ pub enum WarpgateError {
     UrlParse(#[from] url::ParseError),
     #[error("deserialization failed: {0}")]
     DeserializeJson(#[from] serde_json::Error),
-    #[error("external_url config option is not set")]
-    ExternalHostNotSet,
+    #[error("no valid Host header found and `external_host` config option is not set")]
+    ExternalHostUnknown,
+    #[error("current hostname ({0}) is not on the whitelist ({1:?})")]
+    ExternalHostNotWhitelisted(String, Vec<String>),
     #[error("URL contains no host")]
     NoHostInUrl,
     #[error("Inconsistent state error")]

warpgate-protocol-http/src/api/info.rs

@@ -35,10 +35,6 @@ enum InstanceInfoResponse {
     Ok(Json<Info>),
 }
 
-fn strip_port(host: &str) -> Option<&str> {
-    host.split(':').next()
-}
-
 #[OpenApi]
 impl Api {
     #[oai(path = "/info", method = "get", operation_id = "get_info")]
@@ -50,18 +46,15 @@ impl Api {
     ) -> poem::Result<InstanceInfoResponse> {
         let config = services.config.lock().await;
         let external_host = config
-            .store
-            .external_host
-            .as_deref()
-            .and_then(strip_port)
-            .or_else(|| req.header(http::header::HOST).and_then(strip_port))
-            .or_else(|| req.original_uri().host());
+            .construct_external_url(Some(req), None)?
+            .host()
+            .map(|x| x.to_string());
 
         Ok(InstanceInfoResponse::Ok(Json(Info {
             version: env!("CARGO_PKG_VERSION").to_string(),
             username: session.get_username(),
             selected_target: session.get_target_name(),
-            external_host: external_host.map(str::to_string),
+            external_host,
             authorized_via_ticket: matches!(
                 session.get_auth(),
                 Some(SessionAuthorization::Ticket { .. })

warpgate-protocol-http/src/api/sso_provider_detail.rs

@@ -54,15 +54,18 @@ impl Api {
 
         let name = name.0;
 
-        let mut return_url = config.construct_external_url(Some(req))?;
-        return_url.set_path("@warpgate/api/sso/return");
-        debug!("Return URL: {}", &return_url);
-
         let Some(provider_config) = config.store.sso_providers.iter().find(|p| p.name == *name)
         else {
             return Ok(StartSsoResponse::NotFound);
         };
 
+        let mut return_url = config.construct_external_url(
+            Some(req),
+            provider_config.return_domain_whitelist.as_deref(),
+        )?;
+        return_url.set_path("@warpgate/api/sso/return");
+        debug!("Return URL: {}", &return_url);
+
         let client = SsoClient::new(provider_config.provider.clone());
 
         let sso_req = client

warpgate-protocol-http/src/api/sso_provider_list.rs

@@ -300,7 +300,7 @@ impl Api {
 
         let config = services.config.lock().await;
 
-        let return_url = config.construct_external_url(Some(req))?;
+        let return_url = config.construct_external_url(Some(req), None)?;
         debug!("Return URL: {}", &return_url);
 
         let Some(provider_config) = config

warpgate-protocol-http/src/lib.rs

@@ -124,9 +124,8 @@ impl ProtocolServer for HTTPProtocolServer {
                         let url = req.original_uri().clone();
                         let client_ip = get_client_ip(&req).await?;
 
-                        let response = ep.call(req).await.map_err(|e| {
-                            log_request_error(&method, &url, &client_ip, &e);
-                            e
+                        let response = ep.call(req).await.inspect_err(|e| {
+                            log_request_error(&method, &url, &client_ip, e);
                         })?;
 
                         log_request_result(&method, &url, &client_ip, &response.status());

warpgate-protocol-ssh/src/server/session.rs

@@ -1370,7 +1370,7 @@ impl ServerSession {
                         .config
                         .lock()
                         .await
-                        .construct_external_url(None)
+                        .construct_external_url(None, None)
                     {
                         Ok(url) => url,
                         Err(error) => {

warpgate-sso/src/config.rs

@@ -21,6 +21,7 @@ pub struct SsoProviderConfig {
     pub name: String,
     pub label: Option<String>,
     pub provider: SsoInternalProviderConfig,
+    pub return_domain_whitelist: Option<Vec<String>>,
 }
 
 impl SsoProviderConfig {