From 2e16a2be567bf3b548496aaf00e7987f899a4832 Mon Sep 17 00:00:00 2001
From: 21pages <pages21@163.com>
Date: Mon, 22 Jan 2024 19:57:23 +0800
Subject: [PATCH] fix port forward 2fa (#6956)

Signed-off-by: 21pages <pages21@163.com>
---
 src/port_forward.rs      | 8 ++++++--
 src/server/connection.rs | 9 +++------
 2 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/src/port_forward.rs b/src/port_forward.rs
index f71ae2579..f9d38f4c4 100644
--- a/src/port_forward.rs
+++ b/src/port_forward.rs
@@ -141,8 +141,9 @@ async fn connect_and_login(
                         }
                         Some(message::Union::LoginResponse(lr)) => match lr.union {
                             Some(login_response::Union::Error(err)) => {
-                                interface.handle_login_error(&err);
-                                return Ok(None);
+                                if !interface.handle_login_error(&err) {
+                                    return Ok(None);
+                                }
                             }
                             Some(login_response::Union::PeerInfo(pi)) => {
                                 interface.handle_peer_info(pi);
@@ -168,6 +169,9 @@ async fn connect_and_login(
                     Some(Data::Login((os_username, os_password, password, remember))) => {
                         interface.handle_login_from_ui(os_username, os_password, password, remember, &mut stream).await;
                     }
+                    Some(Data::Message(msg)) => {
+                        allow_err!(stream.send(&msg).await);
+                    }
                     _ => {}
                 }
             },
diff --git a/src/server/connection.rs b/src/server/connection.rs
index 8e9e20014..c4156c545 100644
--- a/src/server/connection.rs
+++ b/src/server/connection.rs
@@ -578,6 +578,9 @@ impl Connection {
                                     if !conn.on_message(msg_in).await {
                                         break;
                                     }
+                                    if conn.port_forward_socket.is_some() && conn.authorized {
+                                        break;
+                                    }
                                 }
                             }
                         }
@@ -1627,9 +1630,6 @@ impl Connection {
                     self.linux_headless_handle.wait_desktop_cm_ready().await;
                     self.send_logon_response().await;
                     self.try_start_cm(lr.my_id.clone(), lr.my_name.clone(), self.authorized);
-                    if self.port_forward_socket.is_some() {
-                        return false;
-                    }
                 } else {
                     self.send_login_error(err_msg).await;
                 }
@@ -1667,9 +1667,6 @@ impl Connection {
                         self.linux_headless_handle.wait_desktop_cm_ready().await;
                         self.send_logon_response().await;
                         self.try_start_cm(lr.my_id, lr.my_name, self.authorized);
-                        if self.port_forward_socket.is_some() {
-                            return false;
-                        }
                     } else {
                         self.send_login_error(err_msg).await;
                     }