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::collections::HashSet;
|
||||||
use std::fs;
|
use std::fs;
|
||||||
use std::io::{Read, Write};
|
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::path::{Path, PathBuf};
|
||||||
use std::process::{Command, Stdio};
|
use std::process::{Command, Stdio};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
@@ -70,6 +73,8 @@ struct Args {
|
|||||||
dynamic_forward: Vec<String>,
|
dynamic_forward: Vec<String>,
|
||||||
#[arg(short = 'N', long)]
|
#[arg(short = 'N', long)]
|
||||||
forward_only: bool,
|
forward_only: bool,
|
||||||
|
#[arg(short = 'f', long)]
|
||||||
|
background: bool,
|
||||||
#[arg(long)]
|
#[arg(long)]
|
||||||
predict: bool,
|
predict: bool,
|
||||||
#[arg(long)]
|
#[arg(long)]
|
||||||
@@ -157,7 +162,7 @@ async fn main() -> Result<()> {
|
|||||||
if args.server.as_deref() == Some("trust") {
|
if args.server.as_deref() == Some("trust") {
|
||||||
return run_trust_command(&config, &args);
|
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 hosts = load_hosts_config(None).unwrap_or_default();
|
||||||
let host = hosts
|
let host = hosts
|
||||||
.hosts
|
.hosts
|
||||||
@@ -216,6 +221,9 @@ async fn main() -> Result<()> {
|
|||||||
if forwarding_requested && !auth_allows(&auth_preference, "native") {
|
if forwarding_requested && !auth_allows(&auth_preference, "native") {
|
||||||
return Err(anyhow!("port forwarding requires --auth native or auto"));
|
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 allow_cache = !args.no_cache && !forwarding_requested;
|
||||||
|
|
||||||
let started = Instant::now();
|
let started = Instant::now();
|
||||||
@@ -539,6 +547,79 @@ fn startup_command_from_string(command: &str) -> Vec<u8> {
|
|||||||
command.into_bytes()
|
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>> {
|
fn parse_local_forwards(raw: &[String]) -> Result<Vec<LocalForward>> {
|
||||||
raw.iter().map(|value| parse_local_forward(value)).collect()
|
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));
|
let stream_ids = Arc::new(AtomicU64::new(1));
|
||||||
start_local_forwards(&local_forwards, forward_tx.clone(), Arc::clone(&stream_ids)).await?;
|
start_local_forwards(&local_forwards, forward_tx.clone(), Arc::clone(&stream_ids)).await?;
|
||||||
start_dynamic_forwards(&dynamic_forwards, forward_tx.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 stream_writers: HashMap<u64, tokio::net::tcp::OwnedWriteHalf> = HashMap::new();
|
||||||
let mut pending_socks_replies: HashSet<u64> = HashSet::new();
|
let mut pending_socks_replies: HashSet<u64> = HashSet::new();
|
||||||
if let Some(frame) = first_frame {
|
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]
|
#[test]
|
||||||
fn native_remote_forward_echo_smoke() {
|
fn native_remote_forward_echo_smoke() {
|
||||||
let dir = tempfile::tempdir().unwrap();
|
let dir = tempfile::tempdir().unwrap();
|
||||||
|
|||||||
Reference in New Issue
Block a user