Add background forwarding mode
This commit is contained in:
+84
-2
@@ -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 {
|
||||
|
||||
@@ -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();
|
||||
|
||||
Reference in New Issue
Block a user