Add background forwarding mode
ci / test (push) Has been cancelled
ci / remote-bench (push) Has been cancelled

This commit is contained in:
DuProcess
2026-06-13 15:55:14 -04:00
parent da16588bc5
commit 6d1e9d3065
2 changed files with 137 additions and 2 deletions
+84 -2
View File
@@ -27,7 +27,10 @@ use std::collections::HashMap;
use std::collections::HashSet;
use std::fs;
use std::io::{Read, Write};
use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, ToSocketAddrs};
use std::net::{
Ipv4Addr, Ipv6Addr, SocketAddr, TcpListener as StdTcpListener, TcpStream as StdTcpStream,
ToSocketAddrs,
};
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use std::sync::Arc;
@@ -70,6 +73,8 @@ struct Args {
dynamic_forward: Vec<String>,
#[arg(short = 'N', long)]
forward_only: bool,
#[arg(short = 'f', long)]
background: bool,
#[arg(long)]
predict: bool,
#[arg(long)]
@@ -157,7 +162,7 @@ async fn main() -> Result<()> {
if args.server.as_deref() == Some("trust") {
return run_trust_command(&config, &args);
}
let requested_server = args.server.unwrap_or_else(|| config.server.clone());
let requested_server = args.server.clone().unwrap_or_else(|| config.server.clone());
let hosts = load_hosts_config(None).unwrap_or_default();
let host = hosts
.hosts
@@ -216,6 +221,9 @@ async fn main() -> Result<()> {
if forwarding_requested && !auth_allows(&auth_preference, "native") {
return Err(anyhow!("port forwarding requires --auth native or auto"));
}
if args.background && std::env::var_os("DOSH_BACKGROUND_CHILD").is_none() {
return spawn_background_forwarder(&args, forwarding_requested);
}
let allow_cache = !args.no_cache && !forwarding_requested;
let started = Instant::now();
@@ -539,6 +547,79 @@ fn startup_command_from_string(command: &str) -> Vec<u8> {
command.into_bytes()
}
fn spawn_background_forwarder(args: &Args, forwarding_requested: bool) -> Result<()> {
if !forwarding_requested {
return Err(anyhow!(
"-f requires at least one -L, -R, or -D forwarding request"
));
}
if !args.forward_only {
return Err(anyhow!("-f currently requires -N for forward-only mode"));
}
let listener =
StdTcpListener::bind("127.0.0.1:0").context("bind background readiness socket")?;
listener
.set_nonblocking(true)
.context("configure background readiness socket")?;
let readiness_addr = listener.local_addr()?;
let token = URL_SAFE_NO_PAD.encode(crypto::random_16());
let exe = std::env::current_exe().context("locate current dosh-client executable")?;
let mut child = Command::new(exe)
.args(std::env::args_os().skip(1))
.env("DOSH_BACKGROUND_CHILD", "1")
.env("DOSH_BACKGROUND_READY_ADDR", readiness_addr.to_string())
.env("DOSH_BACKGROUND_READY_TOKEN", &token)
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.spawn()
.context("spawn background dosh forwarder")?;
if let Some(path) = std::env::var_os("DOSH_BACKGROUND_PID_FILE") {
fs::write(path, child.id().to_string()).context("write background pid file")?;
}
let deadline = Instant::now() + Duration::from_secs(30);
loop {
match listener.accept() {
Ok((mut stream, _)) => {
let mut raw = String::new();
stream.read_to_string(&mut raw).ok();
if raw.trim() == token {
eprintln!("dosh background forwarding ready");
return Ok(());
}
return Err(anyhow!("background forwarder sent invalid readiness token"));
}
Err(err) if err.kind() == std::io::ErrorKind::WouldBlock => {}
Err(err) => return Err(err).context("wait for background readiness"),
}
if let Some(status) = child.try_wait()? {
return Err(anyhow!(
"background dosh forwarder exited before readiness: {status}"
));
}
if Instant::now() >= deadline {
let _ = child.kill();
let _ = child.wait();
return Err(anyhow!("background dosh forwarder did not become ready"));
}
std::thread::sleep(Duration::from_millis(20));
}
}
fn signal_background_ready() {
let Some(addr) = std::env::var_os("DOSH_BACKGROUND_READY_ADDR") else {
return;
};
let Some(token) = std::env::var_os("DOSH_BACKGROUND_READY_TOKEN") else {
return;
};
if let Ok(mut stream) = StdTcpStream::connect(addr.to_string_lossy().as_ref()) {
let _ = stream.write_all(token.to_string_lossy().as_bytes());
let _ = stream.shutdown(std::net::Shutdown::Write);
}
}
fn parse_local_forwards(raw: &[String]) -> Result<Vec<LocalForward>> {
raw.iter().map(|value| parse_local_forward(value)).collect()
}
@@ -1513,6 +1594,7 @@ async fn run_terminal(
let stream_ids = Arc::new(AtomicU64::new(1));
start_local_forwards(&local_forwards, forward_tx.clone(), Arc::clone(&stream_ids)).await?;
start_dynamic_forwards(&dynamic_forwards, forward_tx.clone(), stream_ids).await?;
signal_background_ready();
let mut stream_writers: HashMap<u64, tokio::net::tcp::OwnedWriteHalf> = HashMap::new();
let mut pending_socks_replies: HashSet<u64> = HashSet::new();
if let Some(frame) = first_frame {
+53
View File
@@ -626,6 +626,59 @@ fn native_local_forward_echo_smoke() {
);
}
#[test]
fn native_local_forward_background_smoke() {
let dir = tempfile::tempdir().unwrap();
let port = free_udp_port();
let config = write_server_config(&dir, port);
write_native_client_auth(&dir, &config);
let echo_port = start_echo_server();
let local_port = free_tcp_port();
let pid_file = dir.path().join("background.pid");
let mut server = start_server(&dir, &config);
let client_bin = env!("CARGO_BIN_EXE_dosh-client");
let output = Command::new(client_bin)
.arg("--auth")
.arg("native")
.arg("--no-cache")
.arg("-f")
.arg("-N")
.arg("-L")
.arg(format!("{local_port}:127.0.0.1:{echo_port}"))
.arg("--session")
.arg("default")
.arg("--dosh-host")
.arg("127.0.0.1")
.arg("--dosh-port")
.arg(port.to_string())
.arg("local")
.env("HOME", dir.path())
.env("DOSH_BACKGROUND_PID_FILE", &pid_file)
.output()
.unwrap();
assert!(
output.status.success(),
"stderr={}",
String::from_utf8_lossy(&output.stderr)
);
let mut stream = TcpStream::connect(("127.0.0.1", local_port)).unwrap();
stream
.set_read_timeout(Some(Duration::from_secs(3)))
.unwrap();
stream.write_all(b"dosh-background-forward-ping").unwrap();
let mut buf = [0u8; 64];
let n = stream.read(&mut buf).unwrap();
if let Ok(pid) = fs::read_to_string(&pid_file) {
let _ = Command::new("kill").arg(pid.trim()).status();
}
let _ = server.kill();
let _ = server.wait();
assert_eq!(&buf[..n], b"dosh-background-forward-ping");
}
#[test]
fn native_remote_forward_echo_smoke() {
let dir = tempfile::tempdir().unwrap();