use bytes::Bytes;
use russh::Channel;
use russh::client::Msg;
use std::io;
use std::net::SocketAddr;
use tokio::sync::mpsc::{Receiver, Sender, channel};
use tokio::task::JoinHandle;
use super::ToSocketAddrsWithHostname;
use super::address_family::AddressFamily;
use super::connection::Client;
use crate::security::{SudoPassword, contains_sudo_failure, contains_sudo_prompt};
const OUTPUT_EVENTS_CHANNEL_SIZE: usize = 100;
const MAX_SUDO_PROMPT_BUFFER_SIZE: usize = 64 * 1024;
const MAX_SUDO_PASSWORD_SENDS: u32 = 10;
#[derive(Debug, Clone, PartialEq, Eq)]
struct DirectTcpipRequestTarget {
host: String,
port: u32,
}
#[derive(Debug, Clone)]
pub enum CommandOutput {
StdOut(Bytes),
StdErr(Bytes),
ExitCode(u32),
}
pub(crate) struct CommandOutputBuffer {
pub(crate) sender: Sender<CommandOutput>,
pub(crate) receiver_task: JoinHandle<(Vec<u8>, Vec<u8>)>,
}
impl CommandOutputBuffer {
pub(crate) fn new() -> Self {
let (sender, mut receiver): (Sender<CommandOutput>, Receiver<CommandOutput>) =
channel(OUTPUT_EVENTS_CHANNEL_SIZE);
let receiver_task = tokio::task::spawn(async move {
let mut stdout = Vec::with_capacity(1024); let mut stderr = Vec::with_capacity(256);
while let Some(output) = receiver.recv().await {
match output {
CommandOutput::StdOut(buffer) => {
let required = stdout.len() + buffer.len();
if stdout.capacity() < required {
let new_capacity =
required.max(stdout.capacity() + stdout.capacity() / 2);
stdout.reserve(new_capacity - stdout.capacity());
}
stdout.extend_from_slice(&buffer);
}
CommandOutput::StdErr(buffer) => {
let required = stderr.len() + buffer.len();
if stderr.capacity() < required {
let new_capacity =
required.max(stderr.capacity() + stderr.capacity() / 2);
stderr.reserve(new_capacity - stderr.capacity());
}
stderr.extend_from_slice(&buffer);
}
CommandOutput::ExitCode(_) => {
}
}
}
(stdout, stderr)
});
Self {
sender,
receiver_task,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct CommandExecutedResult {
pub stdout: String,
pub stderr: String,
pub exit_status: u32,
}
impl Client {
fn direct_tcpip_targets<T: ToSocketAddrsWithHostname>(
target: &T,
address_family: AddressFamily,
) -> Result<Vec<SocketAddr>, super::Error> {
let resolved = target
.to_socket_addrs()
.map_err(super::Error::AddressInvalid)?;
let targets = address_family.filter(resolved);
if address_family.is_forced() && targets.is_empty() {
return Err(super::Error::NoAddressForFamily {
host: target.hostname(),
family: address_family,
});
}
Ok(targets)
}
fn direct_tcpip_request_targets<T: ToSocketAddrsWithHostname>(
target: &T,
address_family: AddressFamily,
) -> Result<Vec<DirectTcpipRequestTarget>, super::Error> {
if !address_family.is_forced() {
let (host, port) = target.host_port().map_err(super::Error::AddressInvalid)?;
return Ok(vec![DirectTcpipRequestTarget {
host,
port: port.into(),
}]);
}
Self::direct_tcpip_targets(target, address_family).map(|targets| {
targets
.into_iter()
.map(|target| DirectTcpipRequestTarget {
host: target.ip().to_string(),
port: target.port().into(),
})
.collect()
})
}
pub async fn get_channel(&self) -> Result<Channel<Msg>, super::Error> {
self.connection_handle
.channel_open_session()
.await
.map_err(super::Error::SshError)
}
pub async fn open_direct_tcpip_channel<
T: ToSocketAddrsWithHostname,
S: Into<Option<SocketAddr>>,
>(
&self,
target: T,
src: S,
) -> Result<Channel<Msg>, super::Error> {
self.open_direct_tcpip_channel_with_family(target, src, AddressFamily::Any)
.await
}
pub async fn open_direct_tcpip_channel_with_family<
T: ToSocketAddrsWithHostname,
S: Into<Option<SocketAddr>>,
>(
&self,
target: T,
src: S,
address_family: AddressFamily,
) -> Result<Channel<Msg>, super::Error> {
let targets = Self::direct_tcpip_request_targets(&target, address_family)?;
let src = src
.into()
.map(|src| (src.ip().to_string(), src.port().into()))
.unwrap_or_else(|| ("127.0.0.1".to_string(), 22));
let mut connect_err = super::Error::AddressInvalid(io::Error::new(
io::ErrorKind::InvalidInput,
"could not resolve to any addresses",
));
for DirectTcpipRequestTarget { host, port } in targets {
match self
.connection_handle
.channel_open_direct_tcpip(host, port, src.0.clone(), src.1)
.await
{
Ok(channel) => return Ok(channel),
Err(err) => connect_err = super::Error::SshError(err),
}
}
Err(connect_err)
}
pub async fn execute_streaming(
&self,
command: &str,
sender: Sender<CommandOutput>,
) -> Result<u32, super::Error> {
let sanitized_command = crate::utils::sanitize_command(command)
.map_err(|e| super::Error::CommandValidationFailed(e.to_string()))?;
let mut channel = self.connection_handle.channel_open_session().await?;
channel.exec(true, sanitized_command.as_str()).await?;
let mut result: Option<u32> = None;
while let Some(msg) = channel.wait().await {
match msg {
russh::ChannelMsg::Data { ref data } => {
match sender.try_send(CommandOutput::StdOut(data.clone())) {
Ok(_) => {}
Err(tokio::sync::mpsc::error::TrySendError::Full(output)) => {
tracing::trace!("Channel full, applying backpressure for stdout");
if sender.send(output).await.is_err() {
tracing::debug!("Receiver dropped, stopping stdout processing");
break;
}
}
Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
tracing::debug!("Channel closed, stopping stdout processing");
break;
}
}
}
russh::ChannelMsg::ExtendedData { ref data, ext: 1 } => {
match sender.try_send(CommandOutput::StdErr(data.clone())) {
Ok(_) => {}
Err(tokio::sync::mpsc::error::TrySendError::Full(output)) => {
tracing::trace!("Channel full, applying backpressure for stderr");
if sender.send(output).await.is_err() {
tracing::debug!("Receiver dropped, stopping stderr processing");
break;
}
}
Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
tracing::debug!("Channel closed, stopping stderr processing");
break;
}
}
}
russh::ChannelMsg::ExitStatus { exit_status } => result = Some(exit_status),
_ => {}
}
}
drop(sender);
if let Some(result) = result {
Ok(result)
} else {
Err(super::Error::CommandDidntExit)
}
}
pub async fn execute_with_sudo(
&self,
command: &str,
sender: Sender<CommandOutput>,
sudo_password: &SudoPassword,
) -> Result<u32, super::Error> {
let sanitized_command = crate::utils::sanitize_command(command)
.map_err(|e| super::Error::CommandValidationFailed(e.to_string()))?;
let mut channel = self.connection_handle.channel_open_session().await?;
channel
.request_pty(
true, "xterm", 80, 24, 0, 0, &[], )
.await?;
channel.exec(true, sanitized_command.as_str()).await?;
let mut result: Option<u32> = None;
let mut password_send_count: u32 = 0;
let mut accumulated_output = String::new();
while let Some(msg) = channel.wait().await {
match msg {
russh::ChannelMsg::Data { ref data } => {
let text = String::from_utf8_lossy(data);
accumulated_output.push_str(&text);
if accumulated_output.len() > MAX_SUDO_PROMPT_BUFFER_SIZE {
let truncate_at = accumulated_output.len() - MAX_SUDO_PROMPT_BUFFER_SIZE;
accumulated_output = accumulated_output[truncate_at..].to_string();
tracing::debug!(
"Sudo prompt buffer exceeded limit, truncated to {} bytes",
MAX_SUDO_PROMPT_BUFFER_SIZE
);
}
match sender.try_send(CommandOutput::StdOut(data.clone())) {
Ok(_) => {}
Err(tokio::sync::mpsc::error::TrySendError::Full(output)) => {
tracing::trace!("Channel full, applying backpressure for stdout");
if sender.send(output).await.is_err() {
tracing::debug!("Receiver dropped, stopping stdout processing");
break;
}
}
Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
tracing::debug!("Channel closed, stopping stdout processing");
break;
}
}
if password_send_count < MAX_SUDO_PASSWORD_SENDS
&& contains_sudo_prompt(&accumulated_output)
{
password_send_count += 1;
tracing::debug!(
"Sudo prompt detected, sending password (attempt {}/{})",
password_send_count,
MAX_SUDO_PASSWORD_SENDS
);
let password_data = sudo_password.with_newline();
if let Err(e) = channel.data(&password_data[..]).await {
tracing::error!("Failed to send sudo password: {}", e);
return Err(super::Error::SshError(e));
}
accumulated_output.clear();
}
if password_send_count > 0 && contains_sudo_failure(&accumulated_output) {
tracing::debug!(
"Sudo authentication failed after {} attempt(s), closing channel",
password_send_count
);
let error_msg = format!(
"\n[bssh] Sudo authentication failed after {} attempt(s). \
Please verify your sudo password is correct.\n",
password_send_count
);
let _ = sender
.send(CommandOutput::StdErr(Bytes::from(error_msg.into_bytes())))
.await;
let _ = sender.send(CommandOutput::ExitCode(1)).await;
let _ = channel.eof().await;
let _ = channel.close().await;
drop(sender);
return Ok(1);
}
}
russh::ChannelMsg::ExtendedData { ref data, ext: 1 } => {
let text = String::from_utf8_lossy(data);
accumulated_output.push_str(&text);
if accumulated_output.len() > MAX_SUDO_PROMPT_BUFFER_SIZE {
let truncate_at = accumulated_output.len() - MAX_SUDO_PROMPT_BUFFER_SIZE;
accumulated_output = accumulated_output[truncate_at..].to_string();
tracing::debug!(
"Sudo prompt buffer exceeded limit (stderr), truncated to {} bytes",
MAX_SUDO_PROMPT_BUFFER_SIZE
);
}
match sender.try_send(CommandOutput::StdErr(data.clone())) {
Ok(_) => {}
Err(tokio::sync::mpsc::error::TrySendError::Full(output)) => {
tracing::trace!("Channel full, applying backpressure for stderr");
if sender.send(output).await.is_err() {
tracing::debug!("Receiver dropped, stopping stderr processing");
break;
}
}
Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
tracing::debug!("Channel closed, stopping stderr processing");
break;
}
}
if password_send_count < MAX_SUDO_PASSWORD_SENDS
&& contains_sudo_prompt(&accumulated_output)
{
password_send_count += 1;
tracing::debug!(
"Sudo prompt detected on stderr, sending password (attempt {}/{})",
password_send_count,
MAX_SUDO_PASSWORD_SENDS
);
let password_data = sudo_password.with_newline();
if let Err(e) = channel.data(&password_data[..]).await {
tracing::error!("Failed to send sudo password: {}", e);
return Err(super::Error::SshError(e));
}
accumulated_output.clear();
}
if password_send_count > 0 && contains_sudo_failure(&accumulated_output) {
tracing::debug!(
"Sudo authentication failed on stderr after {} attempt(s), closing channel",
password_send_count
);
let error_msg = format!(
"\n[bssh] Sudo authentication failed after {} attempt(s). \
Please verify your sudo password is correct.\n",
password_send_count
);
let _ = sender
.send(CommandOutput::StdErr(Bytes::from(error_msg.into_bytes())))
.await;
let _ = sender.send(CommandOutput::ExitCode(1)).await;
let _ = channel.eof().await;
let _ = channel.close().await;
drop(sender);
return Ok(1);
}
}
russh::ChannelMsg::ExitStatus { exit_status } => result = Some(exit_status),
_ => {}
}
}
drop(sender);
if let Some(result) = result {
Ok(result)
} else {
Err(super::Error::CommandDidntExit)
}
}
pub async fn execute(&self, command: &str) -> Result<CommandExecutedResult, super::Error> {
let output_buffer = CommandOutputBuffer::new();
let sender = output_buffer.sender.clone();
let exit_status = self.execute_streaming(command, sender).await?;
drop(output_buffer.sender);
let (stdout_bytes, stderr_bytes) = output_buffer.receiver_task.await.map_err(|e| {
super::Error::JoinError(e)
})?;
Ok(CommandExecutedResult {
stdout: String::from_utf8_lossy(&stdout_bytes).to_string(),
stderr: String::from_utf8_lossy(&stderr_bytes).to_string(),
exit_status,
})
}
pub async fn request_interactive_shell(
&self,
_term_type: &str,
_width: u32,
_height: u32,
) -> Result<Channel<Msg>, super::Error> {
let channel = self.connection_handle.channel_open_session().await?;
Ok(channel)
}
pub async fn resize_pty(
&self,
channel: &mut Channel<Msg>,
width: u32,
height: u32,
) -> Result<(), super::Error> {
channel
.window_change(width, height, 0, 0)
.await
.map_err(super::Error::SshError)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn v4(s: &str) -> SocketAddr {
s.parse().expect("valid IPv4 socket address")
}
fn v6(s: &str) -> SocketAddr {
s.parse().expect("valid IPv6 socket address")
}
fn multi_hop_candidates() -> Vec<SocketAddr> {
vec![
v6("[2001:db8::10]:22"),
v4("192.0.2.10:22"),
v6("[2001:db8::11]:22"),
v4("192.0.2.11:22"),
]
}
#[test]
fn direct_tcpip_targets_preserve_any_family_candidates_for_tunneled_hops() {
let candidates = multi_hop_candidates();
let targets = Client::direct_tcpip_targets(&candidates.as_slice(), AddressFamily::Any)
.expect("unconstrained direct-tcpip target selection must succeed");
assert_eq!(targets, candidates);
}
#[test]
fn direct_tcpip_targets_filter_ipv4_candidates_for_tunneled_hops() {
let candidates = multi_hop_candidates();
let targets = Client::direct_tcpip_targets(&candidates.as_slice(), AddressFamily::V4)
.expect("IPv4 direct-tcpip target selection must succeed");
assert_eq!(targets, vec![v4("192.0.2.10:22"), v4("192.0.2.11:22")]);
}
#[test]
fn direct_tcpip_targets_filter_ipv6_candidates_for_tunneled_hops() {
let candidates = multi_hop_candidates();
let targets = Client::direct_tcpip_targets(&candidates.as_slice(), AddressFamily::V6)
.expect("IPv6 direct-tcpip target selection must succeed");
assert_eq!(
targets,
vec![v6("[2001:db8::10]:22"), v6("[2001:db8::11]:22")]
);
}
#[test]
fn direct_tcpip_request_targets_send_unforced_hostname_without_resolution() {
let targets =
Client::direct_tcpip_request_targets(&"server-only.internal:5432", AddressFamily::Any)
.expect("unforced direct-tcpip targets must not require local DNS");
assert_eq!(
targets,
vec![DirectTcpipRequestTarget {
host: "server-only.internal".to_string(),
port: 5432,
}]
);
}
#[test]
fn direct_tcpip_request_targets_send_unforced_tuple_hostname_for_jump_hops() {
let targets = Client::direct_tcpip_request_targets(
&("jump-private.internal", 2222),
AddressFamily::Any,
)
.expect("unforced jump-hop targets must not require local DNS");
assert_eq!(
targets,
vec![DirectTcpipRequestTarget {
host: "jump-private.internal".to_string(),
port: 2222,
}]
);
}
#[test]
fn direct_tcpip_request_targets_send_forced_ipv4_address() {
let candidates = multi_hop_candidates();
let targets =
Client::direct_tcpip_request_targets(&candidates.as_slice(), AddressFamily::V4)
.expect("forced IPv4 direct-tcpip targets must resolve to numeric addresses");
assert_eq!(
targets,
vec![
DirectTcpipRequestTarget {
host: "192.0.2.10".to_string(),
port: 22,
},
DirectTcpipRequestTarget {
host: "192.0.2.11".to_string(),
port: 22,
},
]
);
}
#[test]
fn direct_tcpip_request_targets_send_forced_ipv6_address() {
let candidates = multi_hop_candidates();
let targets =
Client::direct_tcpip_request_targets(&candidates.as_slice(), AddressFamily::V6)
.expect("forced IPv6 direct-tcpip targets must resolve to numeric addresses");
assert_eq!(
targets,
vec![
DirectTcpipRequestTarget {
host: "2001:db8::10".to_string(),
port: 22,
},
DirectTcpipRequestTarget {
host: "2001:db8::11".to_string(),
port: 22,
},
]
);
}
}