From 7639a2b03d212097feba2cd012c2eaa872ad9210 Mon Sep 17 00:00:00 2001
From: BlockListed <44610569+BlockListed@users.noreply.github.com>
Date: Sat, 9 Sep 2023 14:33:07 +0200
Subject: [PATCH] make identity support multi-domains

---
 src/api/identity.rs | 22 +++++++++++++---------
 1 file changed, 13 insertions(+), 9 deletions(-)

diff --git a/src/api/identity.rs b/src/api/identity.rs
index 9f3cd1bf..0155bc28 100644
--- a/src/api/identity.rs
+++ b/src/api/identity.rs
@@ -17,7 +17,7 @@ use crate::{
         push::register_push_device,
         ApiResult, EmptyResult, JsonResult, JsonUpcase,
     },
-    auth::{generate_organization_api_key_login_claims, ClientHeaders, ClientIp},
+    auth::{generate_organization_api_key_login_claims, ClientHeaders, ClientIp, HostInfo},
     db::{models::*, DbConn},
     error::MapResult,
     mail, util, CONFIG,
@@ -28,7 +28,7 @@ pub fn routes() -> Vec<Route> {
 }
 
 #[post("/connect/token", data = "<data>")]
-async fn login(data: Form<ConnectData>, client_header: ClientHeaders, mut conn: DbConn) -> JsonResult {
+async fn login(data: Form<ConnectData>, client_header: ClientHeaders, host_info: HostInfo, mut conn: DbConn) -> JsonResult {
     let data: ConnectData = data.into_inner();
 
     let mut user_uuid: Option<String> = None;
@@ -48,7 +48,7 @@ async fn login(data: Form<ConnectData>, client_header: ClientHeaders, mut conn:
             _check_is_some(&data.device_name, "device_name cannot be blank")?;
             _check_is_some(&data.device_type, "device_type cannot be blank")?;
 
-            _password_login(data, &mut user_uuid, &mut conn, &client_header.ip).await
+            _password_login(data, &mut user_uuid, &mut conn, &client_header.ip, &host_info.base_url, &host_info.origin).await
         }
         "client_credentials" => {
             _check_is_some(&data.client_id, "client_id cannot be blank")?;
@@ -140,6 +140,8 @@ async fn _password_login(
     user_uuid: &mut Option<String>,
     conn: &mut DbConn,
     ip: &ClientIp,
+    base_url: &str,
+    origin: &str,
 ) -> JsonResult {
     // Validate scope
     let scope = data.scope.as_ref().unwrap();
@@ -250,7 +252,7 @@ async fn _password_login(
 
     let (mut device, new_device) = get_device(&data, conn, &user).await;
 
-    let twofactor_token = twofactor_auth(&user, &data, &mut device, ip, conn).await?;
+    let twofactor_token = twofactor_auth(&user, &data, &mut device, ip, base_url, origin, conn).await?;
 
     if CONFIG.mail_enabled() && new_device {
         if let Err(e) = mail::send_new_device_logged_in(&user.email, &ip.ip.to_string(), &now, &device.name).await {
@@ -480,6 +482,8 @@ async fn twofactor_auth(
     data: &ConnectData,
     device: &mut Device,
     ip: &ClientIp,
+    base_url: &str,
+    origin: &str,
     conn: &mut DbConn,
 ) -> ApiResult<Option<String>> {
     let twofactors = TwoFactor::find_by_user(&user.uuid, conn).await;
@@ -497,7 +501,7 @@ async fn twofactor_auth(
 
     let twofactor_code = match data.two_factor_token {
         Some(ref code) => code,
-        None => err_json!(_json_err_twofactor(&twofactor_ids, &user.uuid, conn).await?, "2FA token not provided"),
+        None => err_json!(_json_err_twofactor(&twofactor_ids, &user.uuid, base_url, origin, conn).await?, "2FA token not provided"),
     };
 
     let selected_twofactor = twofactors.into_iter().find(|tf| tf.atype == selected_id && tf.enabled);
@@ -511,7 +515,7 @@ async fn twofactor_auth(
         Some(TwoFactorType::Authenticator) => {
             authenticator::validate_totp_code_str(&user.uuid, twofactor_code, &selected_data?, ip, conn).await?
         }
-        Some(TwoFactorType::Webauthn) => webauthn::validate_webauthn_login(&user.uuid, twofactor_code, conn).await?,
+        Some(TwoFactorType::Webauthn) => webauthn::validate_webauthn_login(&user.uuid, twofactor_code, base_url, origin, conn).await?,
         Some(TwoFactorType::YubiKey) => yubikey::validate_yubikey_login(twofactor_code, &selected_data?).await?,
         Some(TwoFactorType::Duo) => {
             duo::validate_duo_login(data.username.as_ref().unwrap().trim(), twofactor_code, conn).await?
@@ -527,7 +531,7 @@ async fn twofactor_auth(
                 }
                 _ => {
                     err_json!(
-                        _json_err_twofactor(&twofactor_ids, &user.uuid, conn).await?,
+                        _json_err_twofactor(&twofactor_ids, &user.uuid, base_url, origin, conn).await?,
                         "2FA Remember token not provided"
                     )
                 }
@@ -555,7 +559,7 @@ fn _selected_data(tf: Option<TwoFactor>) -> ApiResult<String> {
     tf.map(|t| t.data).map_res("Two factor doesn't exist")
 }
 
-async fn _json_err_twofactor(providers: &[i32], user_uuid: &str, conn: &mut DbConn) -> ApiResult<Value> {
+async fn _json_err_twofactor(providers: &[i32], user_uuid: &str, base_url: &str, origin: &str, conn: &mut DbConn) -> ApiResult<Value> {
     let mut result = json!({
         "error" : "invalid_grant",
         "error_description" : "Two factor required.",
@@ -570,7 +574,7 @@ async fn _json_err_twofactor(providers: &[i32], user_uuid: &str, conn: &mut DbCo
             Some(TwoFactorType::Authenticator) => { /* Nothing to do for TOTP */ }
 
             Some(TwoFactorType::Webauthn) if CONFIG.domain_set() => {
-                let request = webauthn::generate_webauthn_login(user_uuid, conn).await?;
+                let request = webauthn::generate_webauthn_login(user_uuid, base_url, origin, conn).await?;
                 result["TwoFactorProviders2"][provider.to_string()] = request.0;
             }