TheLanc3/rustdesk
1
0
Fork 0
mirror of https://github.com/rustdesk/rustdesk.git synced 2026-09-08 05:20:59 +03:00
Code Issues Packages Projects Releases Wiki Activity

port forward: the tunnel's login is the raw pipe's, asked for by a window flag

Browse Source 
Master's fix for the shared login slots (#16069) keeps the target and
the challenge in the window's `LoginConfigHandler` and serializes the
mappings' logins with a turn lock, all inside `port_forward.rs`. This
branch had carried a broader shape of the same fix, a `with_port_forward`
on `Interface` and the target and `Hash` as parameters through the login
functions, which every caller had to follow. That is gone: `Interface`,
`Session`, `create_login_msg`, `send_login`, `handle_hash` and
`handle_login_from_ui` are as on master.

What the tunnel needs on top is one bit in the login, `multiplex`. It is
a window flag beside `port_forward` in the handler, set once in `io_loop`
before the window's mappings start, so an accept's claim and its login
read the same value; the setting takes effect for windows opened after
it changes. `connect_and_login_mux` is now master's `connect_and_login`
with the tunnel's three differences and the same `hash_arrived` and
`login_from_ui` calls. The raw pipe is master's, line for line.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EZ49AbZJYfm8NTp5yDPMab
This commit is contained in:
rustdesk
2026-09-04 19:55:55 +08:00
parent 72a7da3ff9
commit 0fa8e7fedd
4 changed files with 297 additions and 280 deletions
Download Patch File Download Diff File Expand all files Collapse all files

255
src/client.rs
View File

@@ -1752,6 +1752,14 @@ pub struct LoginConfigHandler {
password: Vec<u8>, // remember password for reconnect
pub remember: bool,
config: PeerConfig,
pub port_forward: (String, i32),
/// Set once per window, before its mappings start: the claim and the
/// login of every accept must agree on it.
pub(crate) port_forward_mux: bool,
/// Held by a port-forward mapping from filling `port_forward` and `hash`
/// until its login is built from them; a window's mappings log in
/// concurrently.
pub(crate) port_forward_login_turn: Arc<hbb_common::tokio::sync::Mutex<()>>,
pub version: i64,
features: Option<Features>,
pub session_id: u64, // used for local <-> server communication
@@ -2665,7 +2673,6 @@ impl LoginConfigHandler {
os_username: String,
os_password: String,
password: Vec<u8>,
port_forward: PortForward,
) -> Message {
let my_id = Config::get_id();
let (my_id, pure_id) = if let Some((id, _, _)) = self.other_server.as_ref() {
@@ -2762,7 +2769,12 @@ impl LoginConfigHandler {
..Default::default()
}),
ConnType::VIEW_CAMERA => lr.set_view_camera(Default::default()),
ConnType::PORT_FORWARD | ConnType::RDP => lr.set_port_forward(port_forward),
ConnType::PORT_FORWARD | ConnType::RDP => lr.set_port_forward(PortForward {
host: self.port_forward.0.clone(),
port: self.port_forward.1,
multiplex: self.port_forward_mux,
..Default::default()
}),
ConnType::TERMINAL => {
let mut terminal = Terminal::new();
terminal.service_id = self.get_option(self.get_key_terminal_service_id());
@@ -3496,7 +3508,6 @@ pub async fn handle_hash(
lc: Arc<RwLock<LoginConfigHandler>>,
password_preset: &str,
hash: Hash,
port_forward: PortForward,
interface: &impl Interface,
peer: &mut Stream,
) -> bool {
@@ -3616,7 +3627,7 @@ pub async fn handle_hash(
hasher.finalize()[..].into()
};
send_login(lc.clone(), String::new(), String::new(), password, port_forward, peer).await;
send_login(lc.clone(), String::new(), String::new(), password, peer).await;
lc.write().unwrap().hash = hash;
true
}
@@ -3654,20 +3665,18 @@ fn try_get_password_from_personal_ab(lc: Arc<RwLock<LoginConfigHandler>>, passwo
/// * `os_username` - OS username.
/// * `os_password` - OS password.
/// * `password` - Password.
/// * `port_forward` - Target of a port-forward login; ignored by other types.
/// * `peer` - [`Stream`] for communicating with peer.
async fn send_login(
lc: Arc<RwLock<LoginConfigHandler>>,
os_username: String,
os_password: String,
password: Vec<u8>,
port_forward: PortForward,
peer: &mut Stream,
) {
let msg_out = lc
.read()
.unwrap()
.create_login_msg(os_username, os_password, password, port_forward);
.create_login_msg(os_username, os_password, password);
allow_err!(peer.send(&msg_out).await);
}
@@ -3680,8 +3689,6 @@ async fn send_login(
/// * `os_password` - OS password.
/// * `password` - Password.
/// * `remember` - Whether to remember password.
/// * `port_forward` - Target of a port-forward login; ignored by other types.
/// * `hash` - The challenge this connection was given, if it has one yet.
/// * `peer` - [`Stream`] for communicating with peer.
pub async fn handle_login_from_ui(
lc: Arc<RwLock<LoginConfigHandler>>,
@@ -3689,20 +3696,8 @@ pub async fn handle_login_from_ui(
os_password: String,
password: String,
remember: bool,
port_forward: PortForward,
hash: Option<Hash>,
peer: &mut Stream,
) {
// The window's password prompt is broadcast to every port-forward
// mapping, and can reach one whose own `Hash` has not arrived. It has
// nothing to answer with: a digest over an empty challenge is refused
// and counted as a failed attempt. The mapping that prompted stores the
// salted password in the shared handler, and `handle_hash` logs this one
// in with it when its `Hash` comes.
let Some(hash) = hash else {
log::info!("login from UI before this connection's hash, waiting for it");
return;
};
let mut hash_password = if password.is_empty() {
let mut password2 = lc.read().unwrap().password.clone();
if password2.is_empty() {
@@ -3716,7 +3711,7 @@ pub async fn handle_login_from_ui(
lc.write().unwrap().password_source = Default::default();
let mut hasher = Sha256::new();
hasher.update(password);
hasher.update(&hash.salt);
hasher.update(&lc.read().unwrap().hash.salt);
let res = hasher.finalize();
lc.write().unwrap().remember = remember;
res[..].into()
@@ -3724,10 +3719,10 @@ pub async fn handle_login_from_ui(
lc.write().unwrap().password = hash_password.clone();
let mut hasher2 = Sha256::new();
hasher2.update(&hash_password[..]);
hasher2.update(&hash.challenge);
hasher2.update(&lc.read().unwrap().hash.challenge);
hash_password = hasher2.finalize()[..].to_vec();
send_login(lc.clone(), os_username, os_password, hash_password, port_forward, peer).await;
send_login(lc.clone(), os_username, os_password, hash_password, peer).await;
}
async fn send_switch_login_request(
@@ -3741,7 +3736,7 @@ async fn send_switch_login_request(
lr: hbb_common::protobuf::MessageField::some(
lc.read()
.unwrap()
.create_login_msg("".to_owned(), "".to_owned(), vec![], Default::default())
.create_login_msg("".to_owned(), "".to_owned(), vec![])
.login_request()
.to_owned(),
),
@@ -3774,11 +3769,6 @@ pub trait Interface: Send + Clone + 'static + Sized {
async fn handle_test_delay(&self, t: TestDelay, peer: &mut Stream);
fn get_lch(&self) -> Arc<RwLock<LoginConfigHandler>>;
/// A clone whose port-forward login asks for this target. The target
/// travels with the clone, never through the shared
/// `LoginConfigHandler`, so mappings logging in at the same time cannot
/// overwrite each other's target between the accept and the peer's `Hash`.
fn with_port_forward(&self, port_forward: PortForward) -> Self;
fn get_id(&self) -> String {
self.get_lch().read().unwrap().id.clone()
@@ -4081,207 +4071,22 @@ mod retry_tests {
}
#[cfg(test)]
mod login_scope_tests {
mod port_forward_mux_tests {
use super::*;
use hbb_common::{
tcp::FramedStream,
tokio::{
self,
time::{timeout, Duration},
},
Stream,
};
fn target(host: &str, port: i32) -> PortForward {
PortForward {
host: host.to_owned(),
port,
..Default::default()
}
}
fn hash(challenge: &str) -> Hash {
Hash {
salt: "salt".to_owned(),
challenge: challenge.to_owned(),
..Default::default()
}
}
/// What the peer expects for password `pw` under `hash(challenge)`.
fn digest(challenge: &str) -> Vec<u8> {
let mut h = Sha256::new();
h.update("pw");
h.update("salt");
let salted = h.finalize();
let mut h2 = Sha256::new();
h2.update(&salted[..]);
h2.update(challenge);
h2.finalize()[..].to_vec()
}
/// (our end, the peer's end) of one connection.
async fn loopback() -> (Stream, Stream) {
let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = l.local_addr().unwrap();
let client = tokio::net::TcpStream::connect(addr).await.unwrap();
let (server, _) = l.accept().await.unwrap();
(
Stream::Tcp(FramedStream::from(client, addr)),
Stream::Tcp(FramedStream::from(server, addr)),
)
}
async fn next_login(peer: &mut Stream) -> LoginRequest {
let bytes = peer.next().await.unwrap().unwrap();
Message::parse_from_bytes(&bytes)
.unwrap()
.login_request()
.clone()
}
/// An `Interface` that records the dialogs it was asked to show.
#[derive(Clone, Default)]
struct NoUi(Arc<Mutex<Vec<String>>>);
#[async_trait]
impl Interface for NoUi {
fn send(&self, _data: Data) {}
fn msgbox(&self, msgtype: &str, _title: &str, _text: &str, _link: &str) {
self.0.lock().unwrap().push(msgtype.to_owned());
}
fn handle_login_error(&self, _err: &str) -> bool {
false
}
fn handle_peer_info(&self, _pi: PeerInfo) {}
fn set_multiple_windows_session(&self, _sessions: Vec<WindowsSession>) {}
async fn handle_hash(&self, _pass: &str, _hash: Hash, _peer: &mut Stream) -> bool {
false
}
async fn handle_login_from_ui(
&self,
_os_username: String,
_os_password: String,
_password: String,
_remember: bool,
_peer: &mut Stream,
) {
}
async fn handle_test_delay(&self, _t: TestDelay, _peer: &mut Stream) {}
fn get_lch(&self) -> Arc<RwLock<LoginConfigHandler>> {
Default::default()
}
fn with_port_forward(&self, _port_forward: PortForward) -> Self {
self.clone()
}
}
/// The target rides in the call, never in the shared handler: two accepts
/// building their logins off one handler cannot see each other's target.
#[test]
fn a_port_forward_login_carries_the_callers_target() {
fn a_multiplexed_window_asks_for_the_tunnel() {
let mut lc = LoginConfigHandler::default();
lc.conn_type = ConnType::PORT_FORWARD;
let muxed = |host: &str, port: i32| PortForward {
multiplex: true,
..target(host, port)
let asks = |lc: &LoginConfigHandler| {
lc.create_login_msg(String::new(), String::new(), vec![])
.login_request()
.port_forward()
.multiplex
};
let a = lc.create_login_msg(String::new(), String::new(), vec![], muxed("a", 1));
let b = lc.create_login_msg(String::new(), String::new(), vec![], muxed("b", 2));
let pf = |m: &Message| m.login_request().port_forward().clone();
assert_eq!((pf(&a).host.as_str(), pf(&a).port), ("a", 1));
assert_eq!((pf(&b).host.as_str(), pf(&b).port), ("b", 2));
assert!(pf(&a).multiplex);
}
/// Each connection answers its own challenge: the `Hash` is a parameter
/// of the login, not a field two accepts could overwrite in the handler.
#[test]
fn a_ui_login_answers_the_challenge_it_was_given() {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
rt.block_on(async {
let (mut ours, mut peer) = loopback().await;
let lc = Arc::new(RwLock::new(LoginConfigHandler::default()));
for challenge in ["a", "b"] {
handle_login_from_ui(
lc.clone(),
String::new(),
String::new(),
"pw".to_owned(),
false,
Default::default(),
Some(hash(challenge)),
&mut ours,
)
.await;
assert_eq!(next_login(&mut peer).await.password, digest(challenge));
}
});
}
/// The window's password prompt is broadcast to every mapping, and it can
/// reach one whose own `Hash` has not arrived. That mapping sends nothing:
/// a digest over an empty challenge would only be refused and counted as
/// a failed attempt. It logs in when its `Hash` comes, with the salted
/// password the mapping that prompted stored in the shared handler, and
/// without prompting again.
#[test]
fn a_mapping_still_waiting_for_its_hash_logs_in_when_it_comes() {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
rt.block_on(async {
let (mut a, mut a_peer) = loopback().await;
let (mut b, mut b_peer) = loopback().await;
let lc = Arc::new(RwLock::new(LoginConfigHandler::default()));
lc.write().unwrap().conn_type = ConnType::PORT_FORWARD;
// A has its hash and answers the prompt.
handle_login_from_ui(
lc.clone(),
String::new(),
String::new(),
"pw".to_owned(),
false,
target("a", 1),
Some(hash("a")),
&mut a,
)
.await;
assert_eq!(next_login(&mut a_peer).await.password, digest("a"));
// The same broadcast reaches B, whose hash is still on its way.
handle_login_from_ui(
lc.clone(),
String::new(),
String::new(),
"pw".to_owned(),
false,
target("b", 2),
None,
&mut b,
)
.await;
assert!(
timeout(Duration::from_millis(200), b_peer.next())
.await
.is_err(),
"B logged in before it had a challenge"
);
// B's hash arrives.
let ui = NoUi::default();
assert!(handle_hash(lc.clone(), "", hash("b"), target("b", 2), &ui, &mut b).await);
let login = next_login(&mut b_peer).await;
assert_eq!(login.password, digest("b"));
let pf = login.port_forward();
assert_eq!((pf.host.as_str(), pf.port), ("b", 2));
assert!(ui.0.lock().unwrap().is_empty(), "B prompted again");
});
assert!(!asks(&lc));
lc.port_forward_mux = true;
assert!(asks(&lc));
}
}

298
src/port_forward.rs
View File

@@ -96,11 +96,11 @@ pub async fn listen(
tokio::select! {
Ok((forward, addr)) = listener.accept() => {
log::info!("new connection from {:?}", addr);
// A multiplexed mapping takes the connection on its tunnel, or
// probes for one on its first accept. Everything else, the
// setting off or a peer without the feature, is the raw pipe
// below, as it always was.
let claim = if mux_enabled() { tunnel.claim() } else { Claim::Legacy };
// A multiplexed window takes the connection on the mapping's
// tunnel, or probes for one on its first accept. Everything
// else, the setting off or a peer without the feature, is the
// raw pipe below, as it always was.
let claim = if lc.read().unwrap().port_forward_mux { tunnel.claim() } else { Claim::Legacy };
match claim {
Claim::Muxed(handle) => {
if let Err(e) = handle.open(&remote_host, remote_port, forward, Vec::new()) {
@@ -116,14 +116,11 @@ pub async fn listen(
}
Claim::Legacy => {}
}
// The target rides with this accept's interface clone, so two
// mappings logging in at once cannot overwrite each other's.
let interface = interface.with_port_forward(login_target(&remote_host, remote_port, false));
let id = id.clone();
let password = password.clone();
let mut forward = Framed::new(forward, BytesCodec::new());
let mut close_port_forward = false;
match connect_and_login(&id, &password, &mut ui_receiver, interface.clone(), &mut forward, key, token, is_rdp, &mut close_port_forward).await {
match connect_and_login(&id, &password, &mut ui_receiver, interface.clone(), &mut forward, key, token, is_rdp, &mut close_port_forward, &remote_host, remote_port).await {
Ok(Some(stream)) => {
let interface = interface.clone();
tokio::spawn(async move {
@@ -169,6 +166,8 @@ async fn connect_and_login(
token: &str,
is_rdp: bool,
close_port_forward: &mut bool,
remote_host: &str,
remote_port: i32,
) -> ResultType<Option<Stream>> {
let conn_type = if is_rdp {
ConnType::RDP
@@ -186,6 +185,8 @@ async fn connect_and_login(
}
let mut buffer = Vec::new();
let mut received = false;
let mut challenge = None;
let mut pending_login = None;
let _keep_it = hc_connection(feedback, rendezvous_server, token).await;
@@ -203,7 +204,8 @@ async fn connect_and_login(
let msg_in = Message::parse_from_bytes(&bytes)?;
match msg_in.union {
Some(message::Union::Hash(hash)) => {
if !interface.handle_hash(password, hash, &mut stream).await {
challenge = Some(hash.clone());
if !hash_arrived(&interface, password, hash, pending_login.take(), remote_host, remote_port, &mut stream).await {
return Ok(None);
}
}
@@ -234,9 +236,10 @@ async fn connect_and_login(
},
d = ui_receiver.recv() => {
match d {
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::Login(login)) => match &challenge {
Some(hash) => login_from_ui(&interface, hash, login, remote_host, remote_port, &mut stream).await,
None => pending_login = Some(login),
},
Some(Data::Message(msg)) => {
allow_err!(stream.send(&msg).await);
}
@@ -259,6 +262,76 @@ async fn connect_and_login(
Ok(Some(stream))
}
/// A mapping's login is built from the window's shared handler:
/// `create_login_msg` reads `port_forward` and `handle_login_from_ui` reads
/// `hash`. Mappings log in concurrently, so each fills them and sends under
/// the window's turn lock, or one login carried another mapping's target or
/// answered another's challenge.
async fn login_with_hash(
interface: &impl Interface,
password: &str,
hash: Hash,
remote_host: &str,
remote_port: i32,
stream: &mut Stream,
) -> bool {
let lc = interface.get_lch();
let turn = lc.read().unwrap().port_forward_login_turn.clone();
let _turn = turn.lock().await;
lc.write().unwrap().port_forward = (remote_host.to_owned(), remote_port);
interface.handle_hash(password, hash, stream).await
}
type UiLogin = (String, String, String, bool);
/// This connection's `Hash`. The window's password prompt is broadcast to
/// every mapping and can reach this one first, so a password typed while
/// the `Hash` was on its way is kept and answers it now, rather than being
/// dropped in the hope that the mapping which prompted has already stored
/// it in the shared handler.
async fn hash_arrived(
interface: &impl Interface,
password: &str,
hash: Hash,
pending_login: Option<UiLogin>,
remote_host: &str,
remote_port: i32,
stream: &mut Stream,
) -> bool {
match pending_login {
Some(login) => {
login_from_ui(interface, &hash, login, remote_host, remote_port, stream).await;
true
}
None => login_with_hash(interface, password, hash, remote_host, remote_port, stream).await,
}
}
/// The window's password prompt is broadcast to every mapping; this one
/// answers it with its own challenge.
async fn login_from_ui(
interface: &impl Interface,
hash: &Hash,
login: UiLogin,
remote_host: &str,
remote_port: i32,
stream: &mut Stream,
) {
let lc = interface.get_lch();
let turn = lc.read().unwrap().port_forward_login_turn.clone();
let _turn = turn.lock().await;
{
let mut lc = lc.write().unwrap();
lc.port_forward = (remote_host.to_owned(), remote_port);
lc.set_hash(hash.clone());
}
let (os_username, os_password, password, remember) = login;
interface
.handle_login_from_ui(os_username, os_password, password, remember, stream)
.await;
}
/// The first accept of a multiplexed mapping. It logs in asking for the
/// tunnel, and the peer's answer fixes this listener's mode until it closes:
/// a peer with the feature gets a tunnel every later accept joins, one
@@ -280,10 +353,9 @@ async fn establish_tunnel(
remote_host: &str,
remote_port: i32,
) -> bool {
let interface = interface.with_port_forward(login_target(remote_host, remote_port, true));
let mut forward = Framed::new(forward, BytesCodec::new());
let mut close_port_forward = false;
match connect_and_login_mux(id, password, ui_receiver, interface.clone(), &mut forward, key, token, is_rdp, &mut close_port_forward).await {
match connect_and_login_mux(id, password, ui_receiver, interface.clone(), &mut forward, key, token, is_rdp, &mut close_port_forward, remote_host, remote_port).await {
Ok(Some(outcome)) if outcome.mux => {
let handle = tunnel.set_muxed(outcome.stream, interface.clone());
if !outcome.local_eof {
@@ -314,11 +386,12 @@ async fn establish_tunnel(
false
}
/// `connect_and_login` for a mapping that wants the tunnel: the login asks
/// for it, the pre-read stops at one window rather than growing without
/// bound, and a local EOF no longer ends the login, since the tunnel may
/// still be wanted. It reports what the peer answered rather than a raw
/// stream, because the caller's next step depends on it.
/// `connect_and_login` for a mapping that wants the tunnel: the pre-read
/// stops at one window rather than growing without bound, and a local EOF
/// no longer ends the login, since the tunnel may still be wanted. It
/// reports what the peer answered rather than a raw stream, because the
/// caller's next step depends on it. The login itself is the raw pipe's:
/// the window's `port_forward_mux` is what makes it ask for the tunnel.
async fn connect_and_login_mux(
id: &str,
password: &str,
@@ -329,6 +402,8 @@ async fn connect_and_login_mux(
token: &str,
is_rdp: bool,
close_port_forward: &mut bool,
remote_host: &str,
remote_port: i32,
) -> ResultType<Option<LoginOutcome>> {
let conn_type = if is_rdp {
ConnType::RDP
@@ -458,18 +533,6 @@ struct LoginOutcome {
local_eof: bool,
}
/// The target this accept's login asks for. It travels with the interface
/// clone rather than through the shared `LoginConfigHandler`, so mappings
/// logging in at the same time cannot overwrite each other's target.
fn login_target(host: &str, port: i32, multiplex: bool) -> PortForward {
PortForward {
host: host.to_owned(),
port,
multiplex,
..Default::default()
}
}
fn peer_supports_mux(pi: &PeerInfo) -> bool {
pi.features.as_ref().map(|f| f.port_forward_mux).unwrap_or(false)
}
@@ -482,7 +545,7 @@ fn take_socket(forward: Framed<TcpStream, BytesCodec>, mut prebuf: Vec<u8>) -> (
}
/// The controlling side's `enable-port-forward-mux`: on unless set to `N`.
fn mux_enabled() -> bool {
pub fn mux_enabled() -> bool {
use hbb_common::config::{keys, option2bool, LocalConfig};
option2bool(
keys::OPTION_ENABLE_PORT_FORWARD_MUX,
@@ -515,6 +578,175 @@ async fn run_forward(forward: Framed<TcpStream, BytesCodec>, stream: Stream) ->
Ok(())
}
#[cfg(test)]
mod login_tests {
use super::*;
use async_trait::async_trait;
use hbb_common::{
tcp::FramedStream,
tokio::time::{sleep, Duration},
};
use sha2::{Digest, Sha256};
/// A window's interface over its shared handler. `handle_hash` can pause
/// before building the login, where the real one looks passwords up.
#[derive(Clone)]
struct Ui {
lc: Arc<RwLock<LoginConfigHandler>>,
pause: Duration,
}
#[async_trait]
impl Interface for Ui {
fn send(&self, _data: Data) {}
fn msgbox(&self, _msgtype: &str, _title: &str, _text: &str, _link: &str) {}
fn handle_login_error(&self, _err: &str) -> bool {
false
}
fn handle_peer_info(&self, _pi: PeerInfo) {}
fn set_multiple_windows_session(&self, _sessions: Vec<WindowsSession>) {}
async fn handle_hash(&self, pass: &str, hash: Hash, peer: &mut Stream) -> bool {
sleep(self.pause).await;
crate::client::handle_hash(self.lc.clone(), pass, hash, self, peer).await
}
async fn handle_login_from_ui(
&self,
os_username: String,
os_password: String,
password: String,
remember: bool,
peer: &mut Stream,
) {
crate::client::handle_login_from_ui(
self.lc.clone(),
os_username,
os_password,
password,
remember,
peer,
)
.await
}
async fn handle_test_delay(&self, _t: TestDelay, _peer: &mut Stream) {}
fn get_lch(&self) -> Arc<RwLock<LoginConfigHandler>> {
self.lc.clone()
}
}
fn window() -> Ui {
let mut lc = LoginConfigHandler::default();
lc.conn_type = ConnType::PORT_FORWARD;
Ui {
lc: Arc::new(RwLock::new(lc)),
pause: Duration::ZERO,
}
}
/// (our end, the peer's end) of one connection.
async fn loopback() -> (Stream, Stream) {
let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = l.local_addr().unwrap();
let client = tokio::net::TcpStream::connect(addr).await.unwrap();
let (server, _) = l.accept().await.unwrap();
(
Stream::Tcp(FramedStream::from(client, addr)),
Stream::Tcp(FramedStream::from(server, addr)),
)
}
async fn login_at(peer: &mut Stream) -> LoginRequest {
let bytes = peer.next().await.unwrap().unwrap();
Message::parse_from_bytes(&bytes)
.unwrap()
.login_request()
.clone()
}
fn target(lr: &LoginRequest) -> (String, i32) {
(lr.port_forward().host.clone(), lr.port_forward().port)
}
fn hash(challenge: &str) -> Hash {
Hash {
salt: "salt".to_owned(),
challenge: challenge.to_owned(),
..Default::default()
}
}
/// What the peer expects for password `pw` under `hash(challenge)`.
fn digest(challenge: &str) -> Vec<u8> {
let mut h = Sha256::new();
h.update("pw");
h.update("salt");
let salted = h.finalize();
let mut h2 = Sha256::new();
h2.update(&salted[..]);
h2.update(challenge);
h2.finalize()[..].to_vec()
}
#[test]
fn mappings_logging_in_at_once_each_carry_their_own_target() {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
rt.block_on(async {
let mut ui = window();
ui.pause = Duration::from_millis(50);
let (mut a, mut a_peer) = loopback().await;
let (mut b, mut b_peer) = loopback().await;
tokio::join!(
login_with_hash(&ui, "pw", hash("a"), "a", 1, &mut a),
login_with_hash(&ui, "pw", hash("b"), "b", 2, &mut b),
);
assert_eq!(target(&login_at(&mut a_peer).await), ("a".to_owned(), 1));
assert_eq!(target(&login_at(&mut b_peer).await), ("b".to_owned(), 2));
});
}
#[test]
fn a_mapping_answers_the_prompt_with_its_own_challenge() {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
rt.block_on(async {
let ui = window();
let (mut a, mut a_peer) = loopback().await;
let (mut b, mut b_peer) = loopback().await;
// A's hash arrived last, so it is the one the handler holds.
assert!(login_with_hash(&ui, "pw", hash("a"), "a", 1, &mut a).await);
login_at(&mut a_peer).await;
let typed = (String::new(), String::new(), "pw".to_owned(), false);
login_from_ui(&ui, &hash("b"), typed, "b", 2, &mut b).await;
let lr = login_at(&mut b_peer).await;
assert_eq!(lr.password, digest("b"));
assert_eq!(target(&lr), ("b".to_owned(), 2));
});
}
#[test]
fn a_password_typed_before_this_connections_hash_answers_it_when_it_comes() {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
rt.block_on(async {
let ui = window();
let (mut b, mut b_peer) = loopback().await;
// The prompt's password reached B before its hash, and no other
// mapping has stored it in the handler yet.
let typed = (String::new(), String::new(), "pw".to_owned(), false);
assert!(hash_arrived(&ui, "", hash("b"), Some(typed), "b", 2, &mut b).await);
let lr = login_at(&mut b_peer).await;
assert_eq!(lr.password, digest("b"));
assert_eq!(target(&lr), ("b".to_owned(), 2));
});
}
}
#[cfg(test)]
mod tests {
use super::*;

3
src/port_forward_mux.rs
View File

@@ -1355,9 +1355,6 @@ mod tests {
fn handle_login_error(&self, _err: &str) -> bool {
false
}
fn with_port_forward(&self, _port_forward: PortForward) -> Self {
self.clone()
}
fn handle_peer_info(&self, _pi: PeerInfo) {}
fn set_multiple_windows_session(&self, _sessions: Vec<WindowsSession>) {}
async fn handle_hash(&self, _pass: &str, _hash: Hash, _peer: &mut Stream) -> bool {

21
src/ui_session_interface.rs
View File

@@ -57,13 +57,6 @@ const CHANGE_RESOLUTION_VALID_TIMEOUT_SECS: u64 = 15;
pub struct Session<T: InvokeUiSession> {
pub password: String,
pub args: Vec<String>,
/// Per clone, set by `Interface::with_port_forward`: the target a
/// port-forward login asks for.
pub port_forward: PortForward,
/// The `Hash` this connection was challenged with. A session's clones
/// share it, as they share the connection; a port-forward accept's clone
/// gets its own, since every accept is a connection of its own.
pub login_hash: Arc<RwLock<Option<Hash>>>,
pub lc: Arc<RwLock<LoginConfigHandler>>,
pub sender: Arc<RwLock<Option<mpsc::UnboundedSender<Data>>>>,
pub thread: Arc<Mutex<Option<std::thread::JoinHandle<()>>>>,
@@ -1875,16 +1868,8 @@ impl<T: InvokeUiSession> Interface for Session<T> {
}
}
fn with_port_forward(&self, port_forward: PortForward) -> Self {
let mut scoped = self.clone();
scoped.port_forward = port_forward;
scoped.login_hash = Default::default();
scoped
}
async fn handle_hash(&self, pass: &str, hash: Hash, peer: &mut Stream) -> bool {
*self.login_hash.write().unwrap() = Some(hash.clone());
handle_hash(self.lc.clone(), pass, hash, self.port_forward.clone(), self, peer).await
handle_hash(self.lc.clone(), pass, hash, self, peer).await
}
async fn handle_login_from_ui(
@@ -1895,15 +1880,12 @@ impl<T: InvokeUiSession> Interface for Session<T> {
remember: bool,
peer: &mut Stream,
) {
let hash = self.login_hash.read().unwrap().clone();
handle_login_from_ui(
self.lc.clone(),
os_username,
os_password,
password,
remember,
self.port_forward.clone(),
hash,
peer,
)
.await;
@@ -1962,6 +1944,7 @@ pub async fn io_loop<T: InvokeUiSession>(handler: Session<T>, round: u32) {
let key = crate::get_key(false).await;
#[cfg(not(any(target_os = "android", target_os = "ios")))]
if handler.is_port_forward() {
handler.lc.write().unwrap().port_forward_mux = crate::port_forward::mux_enabled();
if handler.is_rdp() {
let port = handler
.get_option("rdp_port".to_owned())