diff --git a/src/net/mod.rs b/src/net/mod.rs index 9cf4f61..92634aa 100644 --- a/src/net/mod.rs +++ b/src/net/mod.rs @@ -47,7 +47,6 @@ impl Transport { pub async fn send(&mut self, data: T) -> Result<(), TransportError> { // Serialize, encrypt, and then (TODO) sign // NOTE: Cannot used packed implementation for now due to issues with deserialization - // let data = serde_cbor::ser::to_vec_packed(&data)?; let data = serde_cbor::to_vec(&data)?; let data = aead::seal(&self.key, &data)?; @@ -57,24 +56,18 @@ impl Transport { .map_err(TransportError::CodecError) } - /// Receives some data from out on the wire, waiting until it's available - pub async fn receive(&mut self) -> Result { - loop { - if let Some(data) = self.try_receive().await? { - break Ok(data); - } - } - } - - /// Attempts to receive some data from out on the wire, returning that data if available - /// or none if unavailable - pub async fn try_receive(&mut self) -> Result, TransportError> { + /// Receives some data from out on the wire, waiting until it's available, + /// returning none if the transport is now closed + pub async fn receive(&mut self) -> Result, TransportError> { + // If data is received, we process like usual if let Some(data) = self.inner.next().await { - // Validate (TODO), decrypt, and then deserialize + // Validate (TODO) signature, decrypt, and then deserialize let data = data?; let data = aead::open(&self.key, &data)?; let data = serde_cbor::from_slice(&data)?; Ok(Some(data)) + + // Otherwise, if no data is received, this means that our socket has closed } else { Ok(None) } diff --git a/src/opt.rs b/src/opt.rs index 7177f9a..7f324dd 100644 --- a/src/opt.rs +++ b/src/opt.rs @@ -187,7 +187,7 @@ pub struct LaunchSubcommand { #[structopt(short, long, default_value = "ssh")] pub ssh_program: String, - /// Control the IP address that the mosh-server binds to. + /// Control the IP address that the server binds to. /// /// The default is `ssh', in which case the server will reply from the IP address that the SSH /// connection came from (as found in the SSH_CONNECTION environment variable). This is @@ -201,14 +201,9 @@ pub struct LaunchSubcommand { #[structopt(long, value_name = "ssh|any|IP", default_value = "ssh")] pub bind_server: BindAddress, - /// If specified, will write server logs to a file instead of discarding them - #[structopt(long)] - pub server_log_file: Option, - - /// If specified, will set the server's log level (0 is warning and above, 1 is info, 2 is - /// debug, and 3 or higher is trace) - #[structopt(long, default_value = "0")] - pub server_log_level: u8, + /// Additional arguments to provide to the server + #[structopt(long, allow_hyphen_values(true))] + pub extra_server_args: Option, /// If specified, will bind server to the ipv6 interface if host is "any" instead of ipv4 #[structopt(short = "6", long)] @@ -333,7 +328,7 @@ pub struct ListenSubcommand { short, long, value_name = "PORT[:PORT2]", - default_value = "60000:61000" + default_value = "8080:8099" )] pub port: PortRange, } diff --git a/src/subcommand/launch.rs b/src/subcommand/launch.rs index bcf10b9..ccffa53 100644 --- a/src/subcommand/launch.rs +++ b/src/subcommand/launch.rs @@ -4,7 +4,6 @@ use crate::{ }; use derive_more::{Display, Error, From}; use hex::FromHexError; -use log::*; use orion::{aead::SecretKey, errors::UnknownCryptoError}; use std::string::FromUtf8Error; use tokio::{io, process::Command}; @@ -27,18 +26,10 @@ pub fn run(cmd: LaunchSubcommand, opt: CommonOpt) -> Result<(), Error> { async fn run_async(cmd: LaunchSubcommand, _opt: CommonOpt) -> Result<(), Error> { let remote_command = format!( - "{} listen --daemon --host {} {} {} {}", + "{} listen --daemon --host {} {}", cmd.remote_program, cmd.bind_server, - if cmd.use_ipv6 { "-6" } else { "" }, - cmd.server_log_file - .as_ref() - .map(|path| format!("--log-file {:?}", path)) - .unwrap_or_default(), - match cmd.server_log_level { - 0 => String::new(), - n => format!("-{}", "v".repeat(n as usize)), - }, + cmd.extra_server_args.unwrap_or_default(), ); let ssh_command = format!( "{} -o StrictHostKeyChecking=no ssh://{}@{}:{} {} {}", @@ -98,7 +89,7 @@ async fn run_async(cmd: LaunchSubcommand, _opt: CommonOpt) -> Result<(), Error> session.save().await?; if cmd.print_startup_data { - info!("DISTANT DATA {} {}", port, session.to_hex_key()); + println!("DISTANT DATA {} {}", port, session.to_hex_key()); } Ok(()) diff --git a/src/subcommand/listen.rs b/src/subcommand/listen.rs index d259569..2a20f9e 100644 --- a/src/subcommand/listen.rs +++ b/src/subcommand/listen.rs @@ -68,16 +68,14 @@ async fn run_async(cmd: ListenSubcommand, _opt: CommonOpt, is_forked: bool) -> R } } - // Begin our listen loop - loop { - // Wait for a client connection - let (client, _) = listener.accept().await?; - + // Wait for a client connection, then spawn a new task to handle + // receiving data from the client + while let Ok((client, _)) = listener.accept().await { // Grab the client's remote address for later logging purposes let addr_string = match client.peer_addr() { Ok(addr) => { let addr_string = addr.to_string(); - info!("New client from {}", addr_string); + info!(" Established connection", addr_string); addr_string } Err(x) => { @@ -93,9 +91,9 @@ async fn run_async(cmd: ListenSubcommand, _opt: CommonOpt, is_forked: bool) -> R tokio::spawn(async move { loop { match transport.receive::().await { - Ok(request) => { + Ok(Some(request)) => { trace!( - "[{}] Received request of type {}", + " Received request of type {}", addr_string.as_str(), request.as_ref() ); @@ -105,12 +103,16 @@ async fn run_async(cmd: ListenSubcommand, _opt: CommonOpt, is_forked: bool) -> R }; if let Err(x) = transport.send(response).await { - error!("{}", x); + error!(" {}", addr_string.as_str(), x); break; } } + Ok(None) => { + info!(" Closed connection", addr_string.as_str()); + break; + } Err(x) => { - error!("{}", x); + error!(" {}", addr_string.as_str(), x); break; } } @@ -118,7 +120,6 @@ async fn run_async(cmd: ListenSubcommand, _opt: CommonOpt, is_forked: bool) -> R }); } - #[allow(unreachable_code)] Ok(()) }