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.
distant/src/cli/commands/client.rs

852 lines
31 KiB
Rust

use crate::{
cli::{
client::{MsgReceiver, MsgSender},
Cache, Client,
},
config::{
ClientActionConfig, ClientConfig, ClientConnectConfig, ClientLaunchConfig,
ClientReplConfig, NetworkConfig,
},
paths::user::CACHE_FILE_PATH_STR,
CliError, CliResult,
};
use anyhow::Context;
use clap::{Subcommand, ValueHint};
use dialoguer::{console::Term, theme::ColorfulTheme, Select};
use distant_core::{
data::{ChangeKindSet, Environment},
net::{IntoSplit, Request, Response, TypedAsyncRead, TypedAsyncWrite},
ConnectionId, Destination, DistantManagerClient, DistantMsg, DistantRequestData,
DistantResponseData, Host, Map, RemoteCommand, Searcher, Watcher,
};
use log::*;
use serde_json::{json, Value};
use std::{
io,
path::{Path, PathBuf},
time::Duration,
};
mod buf;
mod format;
mod link;
mod lsp;
mod shell;
mod stdin;
pub use format::Format;
use format::Formatter;
use link::RemoteProcessLink;
use lsp::Lsp;
use shell::Shell;
#[derive(Debug, Subcommand)]
pub enum ClientSubcommand {
/// Performs some action on a remote machine
Action {
/// Location to store cached data
#[clap(
long,
value_hint = ValueHint::FilePath,
value_parser,
default_value = CACHE_FILE_PATH_STR.as_str()
)]
cache: PathBuf,
#[clap(flatten)]
config: ClientActionConfig,
/// Specify a connection being managed
#[clap(long)]
connection: Option<ConnectionId>,
#[clap(flatten)]
network: NetworkConfig,
#[clap(subcommand)]
request: DistantRequestData,
},
/// Requests that active manager connects to the server at the specified destination
Connect {
/// Location to store cached data
#[clap(
long,
value_hint = ValueHint::FilePath,
value_parser,
default_value = CACHE_FILE_PATH_STR.as_str()
)]
cache: PathBuf,
#[clap(flatten)]
config: ClientConnectConfig,
#[clap(flatten)]
network: NetworkConfig,
#[clap(short, long, default_value_t, value_enum)]
format: Format,
destination: Box<Destination>,
},
/// Launches the server-portion of the binary on a remote machine
Launch {
/// Location to store cached data
#[clap(
long,
value_hint = ValueHint::FilePath,
value_parser,
default_value = CACHE_FILE_PATH_STR.as_str()
)]
cache: PathBuf,
#[clap(flatten)]
config: ClientLaunchConfig,
#[clap(flatten)]
network: NetworkConfig,
#[clap(short, long, default_value_t, value_enum)]
format: Format,
destination: Box<Destination>,
},
/// Specialized treatment of running a remote LSP process
Lsp {
/// Location to store cached data
#[clap(
long,
value_hint = ValueHint::FilePath,
value_parser,
default_value = CACHE_FILE_PATH_STR.as_str()
)]
cache: PathBuf,
/// Specify a connection being managed
#[clap(long)]
connection: Option<ConnectionId>,
#[clap(flatten)]
network: NetworkConfig,
/// If provided, will run in persist mode, meaning that the process will not be killed if the
/// client disconnects from the server
#[clap(long)]
persist: bool,
/// If provided, will run LSP in a pty
#[clap(long)]
pty: bool,
cmd: String,
},
/// Runs actions in a read-eval-print loop
Repl {
/// Location to store cached data
#[clap(
long,
value_hint = ValueHint::FilePath,
value_parser,
default_value = CACHE_FILE_PATH_STR.as_str()
)]
cache: PathBuf,
#[clap(flatten)]
config: ClientReplConfig,
/// Specify a connection being managed
#[clap(long)]
connection: Option<ConnectionId>,
#[clap(flatten)]
network: NetworkConfig,
/// Format used for input into and output from the repl
#[clap(short, long, default_value_t, value_enum)]
format: Format,
},
/// Select the active connection
Select {
/// Location to store cached data
#[clap(
long,
value_hint = ValueHint::FilePath,
value_parser,
default_value = CACHE_FILE_PATH_STR.as_str()
)]
cache: PathBuf,
/// Connection to use, otherwise will prompt to select
connection: Option<ConnectionId>,
#[clap(short, long, default_value_t, value_enum)]
format: Format,
#[clap(flatten)]
network: NetworkConfig,
},
/// Specialized treatment of running a remote shell process
Shell {
/// Location to store cached data
#[clap(
long,
value_hint = ValueHint::FilePath,
value_parser,
default_value = CACHE_FILE_PATH_STR.as_str()
)]
cache: PathBuf,
/// Specify a connection being managed
#[clap(long)]
connection: Option<ConnectionId>,
#[clap(flatten)]
network: NetworkConfig,
/// Environment variables to provide to the shell
#[clap(long, default_value_t)]
environment: Environment,
/// If provided, will run in persist mode, meaning that the process will not be killed if the
/// client disconnects from the server
#[clap(long)]
persist: bool,
/// Optional command to run instead of $SHELL
cmd: Option<String>,
},
}
impl ClientSubcommand {
pub fn run(self, config: ClientConfig) -> CliResult {
let rt = tokio::runtime::Runtime::new().context("Failed to start up runtime")?;
rt.block_on(Self::async_run(self, config))
}
fn cache_path(&self) -> &Path {
match self {
Self::Action { cache, .. } => cache.as_path(),
Self::Connect { cache, .. } => cache.as_path(),
Self::Launch { cache, .. } => cache.as_path(),
Self::Lsp { cache, .. } => cache.as_path(),
Self::Repl { cache, .. } => cache.as_path(),
Self::Select { cache, .. } => cache.as_path(),
Self::Shell { cache, .. } => cache.as_path(),
}
}
async fn async_run(self, config: ClientConfig) -> CliResult {
// If we get an error, just default anyway
let mut cache = Cache::read_from_disk_or_default(self.cache_path().to_path_buf())
.await
.unwrap_or_else(|_| Cache::new(self.cache_path().to_path_buf()));
match self {
Self::Action {
config: action_config,
connection,
network,
request,
..
} => {
let network = network.merge(config.network);
debug!("Connecting to manager");
let mut client = Client::new(network)
.connect()
.await
.context("Failed to connect to manager")?;
let connection_id =
use_or_lookup_connection_id(&mut cache, connection, &mut client).await?;
debug!("Opening channel to connection {}", connection_id);
let mut channel = client.open_channel(connection_id).await.with_context(|| {
format!("Failed to open channel to connection {connection_id}")
})?;
let timeout = action_config.timeout.or(config.action.timeout);
debug!(
"Timeout configured to be {}",
match timeout {
Some(secs) => format!("{}s", secs),
None => "none".to_string(),
}
);
let mut formatter = Formatter::shell();
debug!("Sending request {:?}", request);
match request {
DistantRequestData::ProcSpawn {
cmd,
environment,
current_dir,
persist,
pty,
} => {
debug!("Special request spawning {:?}", cmd);
let mut proc = RemoteCommand::new()
.environment(environment)
.current_dir(current_dir)
.persist(persist)
.pty(pty)
.spawn(channel, cmd.as_str())
.await
.with_context(|| format!("Failed to spawn {cmd}"))?;
// Now, map the remote process' stdin/stdout/stderr to our own process
let link = RemoteProcessLink::from_remote_pipes(
proc.stdin.take(),
proc.stdout.take().unwrap(),
proc.stderr.take().unwrap(),
);
let status = proc.wait().await.context("Failed to wait for process")?;
// Shut down our link
link.shutdown().await;
if !status.success {
if let Some(code) = status.code {
return Err(CliError::Exit(code as u8));
} else {
return Err(CliError::FAILURE);
}
}
}
DistantRequestData::Search { query } => {
debug!("Special request creating searcher for {:?}", query);
let mut searcher = Searcher::search(channel, query)
.await
.context("Failed to start search")?;
// Continue to receive and process matches
while let Some(m) = searcher.next().await {
// TODO: Provide a cleaner way to print just a match
let res = Response::new(
"".to_string(),
DistantMsg::Single(DistantResponseData::SearchResults {
id: 0,
matches: vec![m],
}),
);
formatter.print(res).context("Failed to print match")?;
}
}
DistantRequestData::Watch {
path,
recursive,
only,
except,
} => {
debug!("Special request creating watcher for {:?}", path);
let mut watcher = Watcher::watch(
channel,
path.as_path(),
recursive,
only.into_iter().collect::<ChangeKindSet>(),
except.into_iter().collect::<ChangeKindSet>(),
)
.await
.with_context(|| format!("Failed to watch {path:?}"))?;
// Continue to receive and process changes
while let Some(change) = watcher.next().await {
// TODO: Provide a cleaner way to print just a change
let res = Response::new(
"".to_string(),
DistantMsg::Single(DistantResponseData::Changed(change)),
);
formatter.print(res).context("Failed to print change")?;
}
}
request => {
let response = channel
.send_timeout(
DistantMsg::Single(request),
timeout
.or(config.action.timeout)
.map(Duration::from_secs_f32),
)
.await
.context("Failed to send request")?;
debug!("Got response {:?}", response);
// NOTE: We expect a single response, and if that is an error then
// we want to pass that error up the stack
let id = response.id;
let origin_id = response.origin_id;
match response.payload {
DistantMsg::Single(DistantResponseData::Error(x)) => {
return Err(CliError::Error(anyhow::anyhow!(x)));
}
payload => formatter
.print(Response {
id,
origin_id,
payload,
})
.context("Failed to print response")?,
}
}
}
}
Self::Connect {
config: connect_config,
network,
format,
destination,
..
} => {
let network = network.merge(config.network);
debug!("Connecting to manager");
let mut client = {
let client = match format {
Format::Shell => Client::new(network),
Format::Json => Client::new(network).using_msg_stdin_stdout(),
};
client
.connect()
.await
.context("Failed to connect to manager")?
};
// Merge our connect configs, overwriting anything in the config file with our cli
// arguments
let mut options = Map::from(config.connect);
options.extend(Map::from(connect_config).into_map());
// Trigger our manager to connect to the launched server
debug!("Connecting to server at {} with {}", destination, options);
let id = client
.connect(*destination, options)
.await
.context("Failed to connect to server")?;
// Mark the server's id as the new default
debug!("Updating selected connection id in cache to {}", id);
*cache.data.selected = id;
cache.write_to_disk().await?;
match format {
Format::Shell => println!("{}", id),
Format::Json => println!(
"{}",
serde_json::to_string(&json!({
"type": "connected",
"id": id,
}))
.unwrap()
),
}
}
Self::Launch {
config: launch_config,
network,
format,
mut destination,
..
} => {
let network = network.merge(config.network);
debug!("Connecting to manager");
let mut client = {
let client = match format {
Format::Shell => Client::new(network),
Format::Json => Client::new(network).using_msg_stdin_stdout(),
};
client
.connect()
.await
.context("Failed to connect to manager")?
};
// Merge our launch configs, overwriting anything in the config file
// with our cli arguments
let mut options = Map::from(config.launch);
options.extend(Map::from(launch_config).into_map());
// Grab the host we are connecting to for later use
let host = destination.host.to_string();
// If we have no scheme on launch, we need to fill it in with something
//
// TODO: Can we have the server support this instead of the client? Right now, the
// server is failing because it cannot parse //localhost/ as it fails with
// an invalid IPv4 or registered name character error on host
if destination.scheme.is_none() {
destination.scheme = Some("ssh".to_string());
}
// Start the server using our manager
debug!("Launching server at {} with {}", destination, options);
let mut new_destination = client
.launch(*destination, options)
.await
.context("Failed to launch server")?;
// Update the new destination with our previously-used host if the
// new host is not globally-accessible
if !new_destination.host.is_global() {
trace!(
"Updating host to {:?} from non-global {:?}",
host,
new_destination.host.to_string()
);
new_destination.host = host
.parse::<Host>()
.map_err(|x| anyhow::anyhow!(x))
.context("Failed to replace host")?;
} else {
trace!("Host {:?} is global", new_destination.host.to_string());
}
// Trigger our manager to connect to the launched server
debug!("Connecting to server at {}", new_destination);
let id = client
.connect(new_destination, Map::new())
.await
.context("Failed to connect to server")?;
// Mark the server's id as the new default
debug!("Updating selected connection id in cache to {}", id);
*cache.data.selected = id;
cache.write_to_disk().await?;
match format {
Format::Shell => println!("{}", id),
Format::Json => println!(
"{}",
serde_json::to_string(&json!({
"type": "launched",
"id": id,
}))
.unwrap()
),
}
}
Self::Lsp {
connection,
network,
persist,
pty,
cmd,
..
} => {
let network = network.merge(config.network);
debug!("Connecting to manager");
let mut client = Client::new(network)
.connect()
.await
.context("Failed to connect to manager")?;
let connection_id =
use_or_lookup_connection_id(&mut cache, connection, &mut client).await?;
debug!("Opening channel to connection {}", connection_id);
let channel = client.open_channel(connection_id).await.with_context(|| {
format!("Failed to open channel to connection {connection_id}")
})?;
debug!(
"Spawning LSP server (persist = {}, pty = {}): {}",
persist, pty, cmd
);
Lsp::new(channel).spawn(cmd, persist, pty).await?;
}
Self::Repl {
config: repl_config,
connection,
network,
format,
..
} => {
let network = network.merge(config.network);
debug!("Connecting to manager");
let mut client = Client::new(network)
.using_msg_stdin_stdout()
.connect()
.await
.context("Failed to connect to manager")?;
let connection_id =
use_or_lookup_connection_id(&mut cache, connection, &mut client).await?;
let timeout = repl_config.timeout.or(config.repl.timeout);
debug!("Opening raw channel to connection {}", connection_id);
let channel = client
.open_raw_channel(connection_id)
.await
.with_context(|| {
format!("Failed to open raw channel to connection {connection_id}")
})?;
debug!(
"Timeout configured to be {}",
match timeout {
Some(secs) => format!("{}s", secs),
None => "none".to_string(),
}
);
// TODO: Support shell format?
if !format.is_json() {
return Err(CliError::Error(anyhow::anyhow!(
"Only JSON format is supported"
)));
}
debug!("Starting repl using format {:?}", format);
let (mut writer, mut reader) = channel.transport.into_split();
let response_task = tokio::task::spawn(async move {
let tx = MsgSender::from_stdout();
while let Some(response) = reader.read().await? {
debug!("Received response {:?}", response);
tx.send_blocking(&response)?;
}
io::Result::Ok(())
});
let request_task = tokio::spawn(async move {
let mut rx = MsgReceiver::from_stdin()
.into_rx::<Request<DistantMsg<DistantRequestData>>>();
loop {
match rx.recv().await {
Some(Ok(request)) => {
debug!("Sending request {:?}", request);
writer.write(request).await?;
}
Some(Err(x)) => error!("{}", x),
None => {
debug!("Shutting down repl");
break;
}
}
}
io::Result::Ok(())
});
let (r1, r2) = tokio::join!(request_task, response_task);
match r1 {
Err(x) => error!("{}", x),
Ok(Err(x)) => error!("{}", x),
_ => (),
}
match r2 {
Err(x) => error!("{}", x),
Ok(Err(x)) => error!("{}", x),
_ => (),
}
debug!("Shutting down repl");
}
Self::Select {
connection,
format,
network,
..
} => match connection {
Some(id) => {
*cache.data.selected = id;
cache.write_to_disk().await?;
}
None => {
let network = network.merge(config.network);
debug!("Connecting to manager");
let mut client = Client::new(network)
.connect()
.await
.context("Failed to connect to manager")?;
let list = client
.list()
.await
.context("Failed to get a list of managed connections")?;
if list.is_empty() {
return Err(CliError::Error(anyhow::anyhow!(
"No connection available in manager"
)));
}
// Figure out the current selection
let current = list
.iter()
.enumerate()
.find_map(|(i, (id, _))| {
if *cache.data.selected == *id {
Some(i)
} else {
None
}
})
.unwrap_or_default();
trace!("Building selection prompt of {} choices", list.len());
let items: Vec<String> = list
.iter()
.map(|(_, destination)| {
format!(
"{}{}{}",
destination
.scheme
.as_ref()
.map(|scheme| format!(r"{scheme}://"))
.unwrap_or_default(),
destination.host,
destination
.port
.map(|port| format!(":{port}"))
.unwrap_or_default()
)
})
.collect();
// Prompt for a selection, with None meaning no change
let selected = match format {
Format::Shell => {
trace!("Rendering prompt");
Select::with_theme(&ColorfulTheme::default())
.items(&items)
.default(current)
.interact_on_opt(&Term::stderr())
.context("Failed to render prompt")?
}
Format::Json => {
// Print out choices
MsgSender::from_stdout()
.send_blocking(&json!({
"type": "select",
"choices": items,
"current": current,
}))
.context("Failed to send JSON choices")?;
// Wait for a response
let msg = MsgReceiver::from_stdin()
.recv_blocking::<Value>()
.context("Failed to receive JSON selection")?;
// Verify the response type is "selected"
match msg.get("type") {
Some(value) if value == "selected" => msg
.get("choice")
.and_then(|value| value.as_u64())
.map(|choice| choice as usize),
Some(value) => {
return Err(CliError::Error(anyhow::anyhow!(
"Unexpected 'type' field value: {value}"
)))
}
None => {
return Err(CliError::Error(anyhow::anyhow!(
"Missing 'type' field"
)))
}
}
}
};
match selected {
Some(index) => {
trace!("Selected choice {}", index);
if let Some((id, _)) = list.iter().nth(index) {
debug!("Updating selected connection id in cache to {}", id);
*cache.data.selected = *id;
cache.write_to_disk().await?;
}
}
None => {
debug!("No change in selection of default connection id");
}
}
}
},
Self::Shell {
connection,
network,
environment,
persist,
cmd,
..
} => {
let network = network.merge(config.network);
debug!("Connecting to manager");
let mut client = Client::new(network)
.connect()
.await
.context("Failed to connect to manager")?;
let connection_id =
use_or_lookup_connection_id(&mut cache, connection, &mut client).await?;
debug!("Opening channel to connection {}", connection_id);
let channel = client.open_channel(connection_id).await.with_context(|| {
format!("Failed to open channel to connection {connection_id}")
})?;
debug!(
"Spawning shell (environment = {:?}, persist = {}): {}",
environment,
persist,
cmd.as_deref().unwrap_or(r"$SHELL")
);
Shell::new(channel).spawn(cmd, environment, persist).await?;
}
}
Ok(())
}
}
async fn use_or_lookup_connection_id(
cache: &mut Cache,
connection: Option<ConnectionId>,
client: &mut DistantManagerClient,
) -> anyhow::Result<ConnectionId> {
match connection {
Some(id) => {
trace!("Using specified connection id: {}", id);
Ok(id)
}
None => {
trace!("Looking up connection id");
let list = client
.list()
.await
.context("Failed to retrieve list of available connections")?;
if list.contains_key(&cache.data.selected) {
trace!("Using cached connection id: {}", cache.data.selected);
Ok(*cache.data.selected)
} else if list.is_empty() {
trace!("Cached connection id is invalid as there are no connections");
anyhow::bail!("There are no connections being managed! You need to start one!");
} else if list.len() > 1 {
trace!("Cached connection id is invalid and there are multiple connections");
anyhow::bail!(
"There are multiple connections being managed! You need to pick one!"
);
} else {
trace!("Cached connection id is invalid");
*cache.data.selected = *list.keys().next().unwrap();
trace!(
"Detected singular connection id, so updating cache: {}",
cache.data.selected
);
cache.write_to_disk().await?;
Ok(*cache.data.selected)
}
}
}
}