From 6154e03c05656aa2a6210f17454403bb46cb2917 Mon Sep 17 00:00:00 2001
From: BlackDex <black.dex@gmail.com>
Date: Thu, 15 Dec 2022 15:57:30 +0100
Subject: [PATCH] Fix recover-2fa not working.

When audit logging was introduced there entered a small bug preventing
the recover-2fa from working.

This PR fixes that by add a new headers check to extract the device-type
when possible and use that for the logging.

Fixes #2985
---
 src/api/core/two_factor/mod.rs | 11 ++++++++---
 src/api/identity.rs            | 18 +++++++++++-------
 src/auth.rs                    | 22 ++++++++++++++++++++++
 3 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/src/api/core/two_factor/mod.rs b/src/api/core/two_factor/mod.rs
index ce3cfb72..a2bbc806 100644
--- a/src/api/core/two_factor/mod.rs
+++ b/src/api/core/two_factor/mod.rs
@@ -6,7 +6,7 @@ use serde_json::Value;
 
 use crate::{
     api::{core::log_user_event, JsonResult, JsonUpcase, NumberOrString, PasswordData},
-    auth::{ClientIp, Headers},
+    auth::{ClientHeaders, ClientIp, Headers},
     crypto,
     db::{models::*, DbConn, DbPool},
     mail, CONFIG,
@@ -73,7 +73,12 @@ struct RecoverTwoFactor {
 }
 
 #[post("/two-factor/recover", data = "<data>")]
-async fn recover(data: JsonUpcase<RecoverTwoFactor>, headers: Headers, mut conn: DbConn, ip: ClientIp) -> JsonResult {
+async fn recover(
+    data: JsonUpcase<RecoverTwoFactor>,
+    client_headers: ClientHeaders,
+    mut conn: DbConn,
+    ip: ClientIp,
+) -> JsonResult {
     let data: RecoverTwoFactor = data.into_inner().data;
 
     use crate::db::models::User;
@@ -97,7 +102,7 @@ async fn recover(data: JsonUpcase<RecoverTwoFactor>, headers: Headers, mut conn:
     // Remove all twofactors from the user
     TwoFactor::delete_all_by_user(&user.uuid, &mut conn).await?;
 
-    log_user_event(EventType::UserRecovered2fa as i32, &user.uuid, headers.device.atype, &ip.ip, &mut conn).await;
+    log_user_event(EventType::UserRecovered2fa as i32, &user.uuid, client_headers.device_type, &ip.ip, &mut conn).await;
 
     // Remove the recovery code, not needed without twofactors
     user.totp_recover = None;
diff --git a/src/api/identity.rs b/src/api/identity.rs
index 6499ee38..05700d62 100644
--- a/src/api/identity.rs
+++ b/src/api/identity.rs
@@ -14,7 +14,7 @@ use crate::{
         core::two_factor::{duo, email, email::EmailTokenData, yubikey},
         ApiResult, EmptyResult, JsonResult, JsonUpcase,
     },
-    auth::ClientIp,
+    auth::{ClientHeaders, ClientIp},
     db::{models::*, DbConn},
     error::MapResult,
     mail, util, CONFIG,
@@ -25,11 +25,10 @@ pub fn routes() -> Vec<Route> {
 }
 
 #[post("/connect/token", data = "<data>")]
-async fn login(data: Form<ConnectData>, mut conn: DbConn, ip: ClientIp) -> JsonResult {
+async fn login(data: Form<ConnectData>, client_header: ClientHeaders, mut conn: DbConn, ip: ClientIp) -> JsonResult {
     let data: ConnectData = data.into_inner();
 
     let mut user_uuid: Option<String> = None;
-    let device_type = data.device_type.clone();
 
     let login_result = match data.grant_type.as_ref() {
         "refresh_token" => {
@@ -59,15 +58,20 @@ async fn login(data: Form<ConnectData>, mut conn: DbConn, ip: ClientIp) -> JsonR
     };
 
     if let Some(user_uuid) = user_uuid {
-        // When unknown or unable to parse, return 14, which is 'Unknown Browser'
-        let device_type = util::try_parse_string(device_type).unwrap_or(14);
         match &login_result {
             Ok(_) => {
-                log_user_event(EventType::UserLoggedIn as i32, &user_uuid, device_type, &ip.ip, &mut conn).await;
+                log_user_event(
+                    EventType::UserLoggedIn as i32,
+                    &user_uuid,
+                    client_header.device_type,
+                    &ip.ip,
+                    &mut conn,
+                )
+                .await;
             }
             Err(e) => {
                 if let Some(ev) = e.get_event() {
-                    log_user_event(ev.event as i32, &user_uuid, device_type, &ip.ip, &mut conn).await
+                    log_user_event(ev.event as i32, &user_uuid, client_header.device_type, &ip.ip, &mut conn).await
                 }
             }
         }
diff --git a/src/auth.rs b/src/auth.rs
index f3e3af8b..69c5203d 100644
--- a/src/auth.rs
+++ b/src/auth.rs
@@ -315,6 +315,28 @@ impl<'r> FromRequest<'r> for Host {
     }
 }
 
+pub struct ClientHeaders {
+    pub host: String,
+    pub device_type: i32,
+}
+
+#[rocket::async_trait]
+impl<'r> FromRequest<'r> for ClientHeaders {
+    type Error = &'static str;
+
+    async fn from_request(request: &'r Request<'_>) -> Outcome<Self, Self::Error> {
+        let host = try_outcome!(Host::from_request(request).await).host;
+        // When unknown or unable to parse, return 14, which is 'Unknown Browser'
+        let device_type: i32 =
+            request.headers().get_one("device-type").map(|d| d.parse().unwrap_or(14)).unwrap_or_else(|| 14);
+
+        Outcome::Success(ClientHeaders {
+            host,
+            device_type,
+        })
+    }
+}
+
 pub struct Headers {
     pub host: String,
     pub device: Device,