mirror of
https://github.com/NVIDIA/OpenShell.git
synced 2026-10-24 06:34:33 +08:00
krishicks
19ae2b7ed7
Forward tests retried on a new port only when the CLI's port check found the port in use. Another test can also take the port after that check passes and before the fake ssh binds it. The fake then panicked with exit status 101, which the retry did not recognize, so the test failed. The fake ssh now exits 255 when its port is taken, as ssh -o ExitOnForwardFailure=yes does, and the tests retry on that exit status too. Signed-off-by: Kris Hicks <khicks@nvidia.com>
74 lines
2.3 KiB
Rust
74 lines
2.3 KiB
Rust
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
use std::net::TcpListener;
|
|
use std::process::Command;
|
|
use std::thread;
|
|
use std::time::Duration;
|
|
|
|
fn main() {
|
|
if std::env::args_os().next().as_deref() != Some(std::ffi::OsStr::new("ssh")) {
|
|
reexec_as_ssh();
|
|
}
|
|
|
|
match std::env::var("OPENSHELL_FAKE_FORWARD_MODE").as_deref() {
|
|
Ok("listen") => run_listener(),
|
|
Ok("sleep") => loop {
|
|
thread::sleep(Duration::from_secs(60));
|
|
},
|
|
_ => std::process::exit(2),
|
|
}
|
|
}
|
|
|
|
fn reexec_as_ssh() -> ! {
|
|
use std::os::unix::process::CommandExt;
|
|
|
|
let executable = std::env::current_exe().expect("fake forward must resolve its executable");
|
|
let error = Command::new(executable)
|
|
.arg0("ssh")
|
|
.args(std::env::args_os().skip(1))
|
|
.exec();
|
|
panic!("fake forward failed to re-execute as ssh: {error}");
|
|
}
|
|
|
|
fn run_listener() {
|
|
let port = forward_port().expect("fake forward must receive an SSH -L argument");
|
|
// Like `ssh -o ExitOnForwardFailure=yes`, exit 255 when another process
|
|
// holds the forward port.
|
|
let listener = match TcpListener::bind(("127.0.0.1", port)) {
|
|
Ok(listener) => listener,
|
|
Err(error) if error.kind() == std::io::ErrorKind::AddrInUse => {
|
|
eprintln!("bind [127.0.0.1]:{port}: Address already in use");
|
|
std::process::exit(255);
|
|
}
|
|
Err(error) => panic!("fake forward must bind: {error}"),
|
|
};
|
|
for stream in listener.incoming() {
|
|
let _ = stream;
|
|
}
|
|
}
|
|
|
|
fn forward_port() -> Option<u16> {
|
|
let args = std::env::args().skip(1).collect::<Vec<_>>();
|
|
let mut index = 0;
|
|
while index < args.len() {
|
|
let arg = &args[index];
|
|
if arg == "-L" {
|
|
return args.get(index + 1).and_then(|value| local_port(value));
|
|
}
|
|
if let Some(value) = arg.strip_prefix("-L").filter(|value| !value.is_empty()) {
|
|
return local_port(value);
|
|
}
|
|
index += 1;
|
|
}
|
|
None
|
|
}
|
|
|
|
fn local_port(forward: &str) -> Option<u16> {
|
|
let (first, rest) = forward.split_once(':')?;
|
|
if first.bytes().all(|byte| byte.is_ascii_digit()) {
|
|
return first.parse().ok();
|
|
}
|
|
rest.split_once(':')?.0.parse().ok()
|
|
}
|