github/dbx
4
0 Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-11-06 00:56:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

restore ssh_tunnel.rs file

Browse Source
This commit is contained in:
runstone committed 2026-06-12 08:53:54 +08:00
1 parent 91f1206783
commit ef1697d8a8
1 file changed
+75 -111
Download Patch File Download Diff File Expand all files Collapse all files
crates/dbx-core/src/db/ssh_tunnel.rs
+75 -111
View File
@@ -5,10 +5,7 @@ use std::sync::Arc;
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use base64::Engine; use base64::Engine;
use russh::client::{self, Config, Handle}; use russh::client::{self, Config, Handle};
use russh::keys::agent::{ use russh::keys::agent::{client::AgentClient, AgentIdentity};
client::{AgentClient, AgentStream},
AgentIdentity,
};
use russh::keys::{decode_secret_key, key::PrivateKeyWithHashAlg, PrivateKey}; use russh::keys::{decode_secret_key, key::PrivateKeyWithHashAlg, PrivateKey};
use russh::ChannelMsg; use russh::ChannelMsg;
use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::io::{AsyncReadExt, AsyncWriteExt};
@@ -46,12 +43,12 @@ impl client::Handler for SshClient {
} }
async fn connect_and_authenticate( async fn connect_and_authenticate(
ssh_host: String, ssh_host: &str,
ssh_port: u16, ssh_port: u16,
ssh_user: String, ssh_user: &str,
ssh_password: String, ssh_password: &str,
ssh_key_path: String, ssh_key_path: &str,
ssh_key_passphrase: String, ssh_key_passphrase: &str,
use_ssh_agent: bool, use_ssh_agent: bool,
connect_timeout_secs: u64, connect_timeout_secs: u64,
) -> Result<Handle<SshClient>, String> { ) -> Result<Handle<SshClient>, String> {
@@ -67,15 +64,15 @@ async fn connect_and_authenticate(
if !ssh_key_path.is_empty() { if !ssh_key_path.is_empty() {
// Validate SSH key file path // Validate SSH key file path
validate_file_path(&ssh_key_path, |_| false)?; validate_file_path(ssh_key_path, |_| false)?;
let passphrase = if ssh_key_passphrase.is_empty() { None } else { Some(ssh_key_passphrase.as_str()) }; let passphrase = if ssh_key_passphrase.is_empty() { None } else { Some(ssh_key_passphrase) };
let key_pair = let key_pair =
load_ssh_private_key(&ssh_key_path, passphrase).map_err(|e| format!("Failed to load SSH key: {e}"))?; load_ssh_private_key(ssh_key_path, passphrase).map_err(|e| format!("Failed to load SSH key: {e}"))?;
let auth_res = tokio::time::timeout( let auth_res = tokio::time::timeout(
connect_timeout, connect_timeout,
session.authenticate_publickey( session.authenticate_publickey(
&ssh_user, ssh_user,
PrivateKeyWithHashAlg::new( PrivateKeyWithHashAlg::new(
Arc::new(key_pair), Arc::new(key_pair),
session.best_supported_rsa_hash().await.ok().flatten().flatten(), session.best_supported_rsa_hash().await.ok().flatten().flatten(),
@@ -89,7 +86,7 @@ async fn connect_and_authenticate(
return Err("SSH public key authentication failed".to_string()); return Err("SSH public key authentication failed".to_string());
} }
} else if !ssh_password.is_empty() { } else if !ssh_password.is_empty() {
let auth_res = tokio::time::timeout(connect_timeout, session.authenticate_password(&ssh_user, &ssh_password)) let auth_res = tokio::time::timeout(connect_timeout, session.authenticate_password(ssh_user, ssh_password))
.await .await
.map_err(|_| format!("SSH password auth timed out ({connect_timeout_secs}s)"))? .map_err(|_| format!("SSH password auth timed out ({connect_timeout_secs}s)"))?
.map_err(|e| format!("SSH password auth failed: {e}"))?; .map_err(|e| format!("SSH password auth failed: {e}"))?;
@@ -97,7 +94,7 @@ async fn connect_and_authenticate(
return Err("SSH password authentication failed".to_string()); return Err("SSH password authentication failed".to_string());
} }
} else if use_ssh_agent { } else if use_ssh_agent {
match try_authenticate_with_agent(&mut session, &ssh_user, &connect_timeout).await { match try_authenticate_with_agent(&mut session, ssh_user, &connect_timeout).await {
Ok(()) => {} Ok(()) => {}
Err(agent_err) => return Err(agent_err), Err(agent_err) => return Err(agent_err),
} }
@@ -108,18 +105,6 @@ async fn connect_and_authenticate(
Ok(session) Ok(session)
} }
/// Connect to the platform-appropriate SSH agent.
#[cfg(unix)]
async fn connect_ssh_agent() -> Result<AgentClient<Box<dyn AgentStream + Send + Unpin>>, String> {
AgentClient::connect_env().await.map(|c| c.dynamic()).map_err(|e| e.to_string())
}
/// Connect to the platform-appropriate SSH agent (Pageant on Windows).
#[cfg(windows)]
async fn connect_ssh_agent() -> Result<AgentClient<Box<dyn AgentStream + Send + Unpin>>, String> {
AgentClient::connect_pageant().await.map(|c| c.dynamic()).map_err(|e| e.to_string())
}
/// Try to authenticate using ssh-agent identities. Returns `Ok(())` on success, /// Try to authenticate using ssh-agent identities. Returns `Ok(())` on success,
/// or an error describing why agent auth failed (unavailable, no identities, all rejected). /// or an error describing why agent auth failed (unavailable, no identities, all rejected).
async fn try_authenticate_with_agent( async fn try_authenticate_with_agent(
@@ -323,12 +308,7 @@ fn read_u32(bytes: &[u8], pos: &mut usize) -> Result<u32, String> {
/// Accept connections on the local listener and forward them through the SSH session. /// Accept connections on the local listener and forward them through the SSH session.
/// Returns when the SSH session dies (listener error or session.is_closed()). /// Returns when the SSH session dies (listener error or session.is_closed()).
async fn forward_loop( async fn forward_loop(session: &Handle<SshClient>, listener: &TcpListener, remote_host: &str, remote_port: u16) {
session: Arc<Handle<SshClient>>,
listener: Arc<TcpListener>,
remote_host: String,
remote_port: u16,
) {
let mut idle_check = tokio::time::interval(IDLE_SESSION_CHECK_INTERVAL); let mut idle_check = tokio::time::interval(IDLE_SESSION_CHECK_INTERVAL);
idle_check.set_missed_tick_behavior(MissedTickBehavior::Delay); idle_check.set_missed_tick_behavior(MissedTickBehavior::Delay);
@@ -370,7 +350,7 @@ async fn forward_loop(
let mut channel = match session let mut channel = match session
.channel_open_direct_tcpip( .channel_open_direct_tcpip(
remote_host.clone(), remote_host,
remote_port.into(), remote_port.into(),
peer_addr.ip().to_string(), peer_addr.ip().to_string(),
peer_addr.port().into(), peer_addr.port().into(),
@@ -427,8 +407,8 @@ async fn forward_loop(
/// Uses exponential backoff for reconnect attempts and gives up after /// Uses exponential backoff for reconnect attempts and gives up after
/// MAX_RECONNECT_ATTEMPTS to avoid log storms from permanent failures. /// MAX_RECONNECT_ATTEMPTS to avoid log storms from permanent failures.
#[allow(clippy::too_many_arguments)] #[allow(clippy::too_many_arguments)]
fn tunnel_reconnect_loop( async fn tunnel_reconnect_loop(
session: Handle<SshClient>, mut session: Handle<SshClient>,
connect_host: String, connect_host: String,
connect_port: u16, connect_port: u16,
ssh_user: String, ssh_user: String,
@@ -437,60 +417,66 @@ fn tunnel_reconnect_loop(
ssh_key_passphrase: String, ssh_key_passphrase: String,
use_ssh_agent: bool, use_ssh_agent: bool,
connect_timeout_secs: u64, connect_timeout_secs: u64,
listener: Arc<TcpListener>, listener: TcpListener,
remote_host: String, remote_host: String,
remote_port: u16, remote_port: u16,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = ()>>> { ) {
let mut session = Arc::new(session); loop {
let label = format!("{connect_host}:{connect_port} -> {remote_host}:{remote_port}"); log::info!("SSH tunnel active: {}:{} -> {}:{}", connect_host, connect_port, remote_host, remote_port);
Box::pin(async move {
forward_loop(&session, &listener, &remote_host, remote_port).await;
log::warn!("SSH tunnel connection lost ({}:{}), reconnecting...", connect_host, connect_port);
// Reconnect with exponential backoff
let mut delay = INITIAL_RECONNECT_DELAY;
let mut attempts: u32 = 0;
loop { loop {
log::info!("SSH tunnel active: {label}"); if attempts >= MAX_RECONNECT_ATTEMPTS {
log::error!(
"SSH tunnel ({connect_host}:{connect_port}): max reconnect attempts ({MAX_RECONNECT_ATTEMPTS}) exhausted, giving up"
);
return;
}
forward_loop(session.clone(), listener.clone(), remote_host.clone(), remote_port).await; tokio::time::sleep(delay).await;
log::warn!("SSH tunnel connection lost ({label}), reconnecting..."); match connect_and_authenticate(
&connect_host,
// Reconnect with exponential backoff connect_port,
let mut delay = INITIAL_RECONNECT_DELAY; &ssh_user,
let mut attempts: u32 = 0; &ssh_password,
&ssh_key_path,
loop { &ssh_key_passphrase,
if attempts >= MAX_RECONNECT_ATTEMPTS { use_ssh_agent,
log::error!( connect_timeout_secs,
"SSH tunnel ({label}): max reconnect attempts ({MAX_RECONNECT_ATTEMPTS}) exhausted, giving up" )
.await
{
Ok(new_session) => {
session = new_session;
log::info!(
"SSH tunnel reconnected to {}:{} (attempt {})",
connect_host,
connect_port,
attempts + 1
); );
return; break;
} }
Err(e) => {
tokio::time::sleep(delay).await; attempts += 1;
log::error!(
match connect_and_authenticate( "SSH reconnect failed ({}:{}, attempt {attempts}/{MAX_RECONNECT_ATTEMPTS}): {e}",
connect_host.clone(), connect_host,
connect_port, connect_port,
ssh_user.clone(), );
ssh_password.clone(), // Exponential backoff: double the delay, cap at MAX_RECONNECT_DELAY
ssh_key_path.clone(), delay = std::cmp::min(delay * 2, MAX_RECONNECT_DELAY);
ssh_key_passphrase.clone(),
use_ssh_agent,
connect_timeout_secs,
)
.await
{
Ok(new_session) => {
session = Arc::new(new_session);
log::info!("SSH tunnel reconnected to {label} (attempt {})", attempts + 1);
break;
}
Err(e) => {
attempts += 1;
log::error!("SSH reconnect failed ({label}, attempt {attempts}/{MAX_RECONNECT_ATTEMPTS}): {e}");
delay = std::cmp::min(delay * 2, MAX_RECONNECT_DELAY);
}
} }
} }
} }
}) }
} }
struct TunnelEntry { struct TunnelEntry {
@@ -546,7 +532,7 @@ impl TunnelManager {
return Ok(port); return Ok(port);
} }
} }
// Slow SSH connection 鈥?do this outside the lock. // Slow SSH connection — do this outside the lock.
let (handle, local_port) = spawn_tunnel( let (handle, local_port) = spawn_tunnel(
ssh_host, ssh_host,
ssh_port, ssh_port,
@@ -684,20 +670,18 @@ async fn spawn_tunnel(
// Initial connection: fail fast on bad credentials // Initial connection: fail fast on bad credentials
let session = connect_and_authenticate( let session = connect_and_authenticate(
connect_host.to_string(), connect_host,
connect_port, connect_port,
ssh_user.to_string(), ssh_user,
ssh_password.to_string(), ssh_password,
ssh_key_path.to_string(), ssh_key_path,
ssh_key_passphrase.to_string(), ssh_key_passphrase,
use_ssh_agent, use_ssh_agent,
connect_timeout_secs, connect_timeout_secs,
) )
.await?; .await?;
// Spawn the tunnel loop on a dedicated blocking thread to avoid let handle = tokio::spawn(tunnel_reconnect_loop(
// a known Rust async Send inference limitation (rust-lang/rust#102211).
let tunnel_args = (
session, session,
connect_host.to_string(), connect_host.to_string(),
connect_port, connect_port,
@@ -707,30 +691,10 @@ async fn spawn_tunnel(
ssh_key_passphrase.to_string(), ssh_key_passphrase.to_string(),
use_ssh_agent, use_ssh_agent,
connect_timeout_secs, connect_timeout_secs,
Arc::new(listener), listener,
remote_host.to_string(), remote_host.to_string(),
remote_port, remote_port,
); ));
let handle = tokio::task::spawn_blocking(move || {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("failed to create tunnel runtime");
rt.block_on(tunnel_reconnect_loop(
tunnel_args.0,
tunnel_args.1,
tunnel_args.2,
tunnel_args.3,
tunnel_args.4,
tunnel_args.5,
tunnel_args.6,
tunnel_args.7,
tunnel_args.8,
tunnel_args.9,
tunnel_args.10,
tunnel_args.11,
));
});
Ok((handle, local_port)) Ok((handle, local_port))
} }