mirror of https://github.com/chipsenkbeil/distant
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
219 lines
7.3 KiB
Rust
219 lines
7.3 KiB
Rust
use crate::{
|
|
cli::{
|
|
opt::{CommonOpt, Format, LaunchSubcommand, SessionOutput},
|
|
CliSession, ExitCode, ExitCodeError,
|
|
},
|
|
core::{
|
|
client::{self, Session, SessionInfo, SessionInfoFile},
|
|
server::RelayServer,
|
|
},
|
|
};
|
|
use derive_more::{Display, Error, From};
|
|
use fork::{daemon, Fork};
|
|
use log::*;
|
|
use std::{path::Path, string::FromUtf8Error};
|
|
use tokio::{io, process::Command, runtime::Runtime, time::Duration};
|
|
|
|
#[derive(Debug, Display, Error, From)]
|
|
pub enum Error {
|
|
#[display(fmt = "Missing data for session")]
|
|
MissingSessionData,
|
|
|
|
ForkError(#[error(not(source))] i32),
|
|
IoError(io::Error),
|
|
Utf8Error(FromUtf8Error),
|
|
}
|
|
|
|
impl ExitCodeError for Error {
|
|
fn to_exit_code(&self) -> ExitCode {
|
|
match self {
|
|
Self::MissingSessionData => ExitCode::NoInput,
|
|
Self::ForkError(_) => ExitCode::OsErr,
|
|
Self::IoError(x) => x.to_exit_code(),
|
|
Self::Utf8Error(_) => ExitCode::DataErr,
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn run(cmd: LaunchSubcommand, opt: CommonOpt) -> Result<(), Error> {
|
|
let rt = Runtime::new()?;
|
|
let session_output = cmd.session;
|
|
let format = cmd.format;
|
|
let is_daemon = cmd.daemon;
|
|
|
|
let session_file = cmd.session_data.session_file.clone();
|
|
let session_socket = cmd.session_data.session_socket.clone();
|
|
let fail_if_socket_exists = cmd.fail_if_socket_exists;
|
|
let timeout = opt.to_timeout_duration();
|
|
let shutdown_after = cmd.to_shutdown_after_duration();
|
|
|
|
let session = rt.block_on(async { spawn_remote_server(cmd, opt).await })?;
|
|
|
|
// Handle sharing resulting session in different ways
|
|
match session_output {
|
|
SessionOutput::File => {
|
|
debug!("Outputting session to {:?}", session_file);
|
|
rt.block_on(async { SessionInfoFile::new(session_file, session).save().await })?
|
|
}
|
|
SessionOutput::Keep => {
|
|
debug!("Entering interactive loop over stdin");
|
|
rt.block_on(async { keep_loop(session, format, timeout).await })?
|
|
}
|
|
SessionOutput::Pipe => {
|
|
debug!("Piping session to stdout");
|
|
println!("{}", session.to_unprotected_string())
|
|
}
|
|
SessionOutput::Socket if is_daemon => {
|
|
debug!(
|
|
"Forking and entering interactive loop over unix socket {:?}",
|
|
session_socket
|
|
);
|
|
|
|
// Force runtime shutdown by dropping it BEFORE forking as otherwise
|
|
// this produces a garbage process that won't die
|
|
drop(rt);
|
|
|
|
match daemon(false, false) {
|
|
Ok(Fork::Child) => {
|
|
// NOTE: We need to create a runtime within the forked process as
|
|
// tokio's runtime doesn't support being transferred from
|
|
// parent to child in a fork
|
|
let rt = Runtime::new()?;
|
|
rt.block_on(async {
|
|
socket_loop(
|
|
session_socket,
|
|
session,
|
|
timeout,
|
|
fail_if_socket_exists,
|
|
shutdown_after,
|
|
)
|
|
.await
|
|
})?
|
|
}
|
|
Ok(_) => {}
|
|
Err(x) => return Err(Error::ForkError(x)),
|
|
}
|
|
}
|
|
#[cfg(unix)]
|
|
SessionOutput::Socket => {
|
|
debug!(
|
|
"Entering interactive loop over unix socket {:?}",
|
|
session_socket
|
|
);
|
|
rt.block_on(async {
|
|
socket_loop(
|
|
session_socket,
|
|
session,
|
|
timeout,
|
|
fail_if_socket_exists,
|
|
shutdown_after,
|
|
)
|
|
.await
|
|
})?
|
|
}
|
|
#[cfg(not(unix))]
|
|
SessionOutput::Socket => {
|
|
debug!(concat!(
|
|
"Trying to enter interactive loop over unix socket, ",
|
|
"but not on unix platform!"
|
|
));
|
|
unreachable!()
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
async fn keep_loop(info: SessionInfo, format: Format, duration: Duration) -> io::Result<()> {
|
|
match Session::tcp_connect_timeout(info, duration).await {
|
|
Ok(session) => {
|
|
let cli_session = CliSession::new(client::new_tenant(), session, format);
|
|
cli_session.wait().await
|
|
}
|
|
Err(x) => Err(x),
|
|
}
|
|
}
|
|
|
|
async fn socket_loop(
|
|
socket_path: impl AsRef<Path>,
|
|
info: SessionInfo,
|
|
duration: Duration,
|
|
fail_if_socket_exists: bool,
|
|
shutdown_after: Option<Duration>,
|
|
) -> io::Result<()> {
|
|
// We need to form a connection with the actual server to forward requests
|
|
// and responses between connections
|
|
debug!("Connecting to {} {}", info.host, info.port);
|
|
let session = Session::tcp_connect_timeout(info, duration).await?;
|
|
|
|
// Remove the socket file if it already exists
|
|
if !fail_if_socket_exists && socket_path.as_ref().exists() {
|
|
debug!("Removing old unix socket instance");
|
|
tokio::fs::remove_file(socket_path.as_ref()).await?;
|
|
}
|
|
|
|
// Continue to receive connections over the unix socket, store them in our
|
|
// connection mapping
|
|
debug!("Binding to unix socket: {:?}", socket_path.as_ref());
|
|
let listener = tokio::net::UnixListener::bind(socket_path)?;
|
|
|
|
let server = RelayServer::initialize(session, listener, shutdown_after).await?;
|
|
server
|
|
.wait()
|
|
.await
|
|
.map_err(|x| io::Error::new(io::ErrorKind::Other, x))
|
|
}
|
|
|
|
/// Spawns a remote server that listens for requests
|
|
///
|
|
/// Returns the session associated with the server
|
|
async fn spawn_remote_server(cmd: LaunchSubcommand, _opt: CommonOpt) -> Result<SessionInfo, Error> {
|
|
let distant_command = format!(
|
|
"{} listen --daemon --host {} {}",
|
|
cmd.distant,
|
|
cmd.bind_server,
|
|
cmd.extra_server_args.unwrap_or_default(),
|
|
);
|
|
let ssh_command = format!(
|
|
"{} -o StrictHostKeyChecking=no ssh://{}@{}:{} {} '{}'",
|
|
cmd.ssh,
|
|
cmd.username,
|
|
cmd.host.as_str(),
|
|
cmd.port,
|
|
cmd.identity_file
|
|
.map(|f| format!("-i {}", f.as_path().display()))
|
|
.unwrap_or_default(),
|
|
if cmd.no_shell {
|
|
distant_command.trim().to_string()
|
|
} else {
|
|
// TODO: Do we need to try to escape single quotes here because of extra_server_args?
|
|
// TODO: Replace this with the ssh2 library shell exec once we integrate that
|
|
format!("echo {} | $SHELL -l", distant_command.trim())
|
|
},
|
|
);
|
|
let out = Command::new("sh")
|
|
.arg("-c")
|
|
.arg(ssh_command)
|
|
.output()
|
|
.await?;
|
|
|
|
// If our attempt to run the program via ssh failed, report it
|
|
if !out.status.success() {
|
|
return Err(Error::from(io::Error::new(
|
|
io::ErrorKind::Other,
|
|
String::from_utf8(out.stderr)?.trim().to_string(),
|
|
)));
|
|
}
|
|
|
|
// Parse our output for the specific session line
|
|
// NOTE: The host provided on this line isn't valid, so we fill it in with our actual host
|
|
let out = String::from_utf8(out.stdout)?.trim().to_string();
|
|
let mut info = out
|
|
.lines()
|
|
.find_map(|line| line.parse::<SessionInfo>().ok())
|
|
.ok_or(Error::MissingSessionData)?;
|
|
info.host = cmd.host;
|
|
|
|
Ok(info)
|
|
}
|