From ccd23a2fdcd58693dbd3c4c0ee8b4159832d141d Mon Sep 17 00:00:00 2001 From: Chip Senkbeil Date: Mon, 26 Jul 2021 03:42:59 -0500 Subject: [PATCH] Support port range binding --- src/opt.rs | 110 ++++++++++++++++++++++++++++++++++----- src/subcommand/listen.rs | 43 ++++++++------- 2 files changed, 120 insertions(+), 33 deletions(-) diff --git a/src/opt.rs b/src/opt.rs index 11f5643..f7a4f46 100644 --- a/src/opt.rs +++ b/src/opt.rs @@ -1,8 +1,9 @@ use crate::subcommand; -use derive_more::Display; +use derive_more::{Display, Error, From}; use lazy_static::lazy_static; use std::{ - net::{AddrParseError, IpAddr, Ipv4Addr}, + env, + net::{AddrParseError, IpAddr, Ipv4Addr, SocketAddr}, path::PathBuf, str::FromStr, }; @@ -77,6 +78,32 @@ pub enum BindAddress { Ip(IpAddr), } +#[derive(Clone, Debug, Display, From, Error, PartialEq, Eq)] +pub enum ConvertToIpAddrError { + ClientIpParseError(AddrParseError), + MissingClientIp, + VarError(env::VarError), +} + +impl BindAddress { + /// Converts address into valid IP + pub fn to_ip_addr(&self) -> Result { + match self { + Self::Ssh => { + let ssh_connection = env::var("SSH_CONNECTION")?; + let ip_str = ssh_connection + .split(' ') + .next() + .ok_or(ConvertToIpAddrError::MissingClientIp)?; + let ip = ip_str.parse::()?; + Ok(ip) + } + Self::Any => Ok(IpAddr::V4(Ipv4Addr::UNSPECIFIED)), + Self::Ip(addr) => Ok(*addr), + } + } +} + impl FromStr for BindAddress { type Err = AddrParseError; @@ -136,6 +163,62 @@ pub struct LaunchSubcommand { pub host: String, } +/// Represents some range of ports +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct PortRange { + pub start: u16, + pub end: Option, +} + +impl PortRange { + /// Builds a collection of `SocketAddr` instances from the port range and given ip address + pub fn make_socket_addrs(&self, addr: impl Into) -> Vec { + let mut socket_addrs = Vec::new(); + let addr = addr.into(); + + for port in self.start..=self.end.unwrap_or(self.start) { + socket_addrs.push(SocketAddr::from((addr, port))); + } + + socket_addrs + } +} + +#[derive(Copy, Clone, Debug, Display, Error, PartialEq, Eq)] +pub enum PortRangeParseError { + InvalidPort, + MissingPort, +} + +impl FromStr for PortRange { + type Err = PortRangeParseError; + + /// Parses PORT into single range or PORT1:PORTN into full range + fn from_str(s: &str) -> Result { + let mut tokens = s.trim().split(':'); + let start = tokens + .next() + .ok_or(PortRangeParseError::MissingPort)? + .parse::() + .map_err(|_| PortRangeParseError::InvalidPort)?; + let end = if let Some(token) = tokens.next() { + Some( + token + .parse::() + .map_err(|_| PortRangeParseError::InvalidPort)?, + ) + } else { + None + }; + + if tokens.next().is_some() { + return Err(PortRangeParseError::InvalidPort); + } + + Ok(Self { start, end }) + } +} + /// Represents subcommand to operate in listen mode for incoming requests #[derive(Debug, StructOpt)] pub struct ListenSubcommand { @@ -151,7 +234,9 @@ pub struct ListenSubcommand { #[structopt(long)] pub bind_ssh_connection: bool, - /// Control the IP address that the distant binds to. There are three options here: + /// Control the IP address that the distant binds to + /// + /// There are three options here: /// /// 1. `ssh`: the server will reply from the IP address that the SSH /// connection came from (as found in the SSH_CONNECTION environment variable). This is @@ -163,15 +248,14 @@ pub struct ListenSubcommand { /// /// 3. `IP`: the server will attempt to bind to the specified IP address. #[structopt(short, long, default_value = "localhost")] - pub host: String, - - /// Represents the port to bind to when listening - #[structopt(short, long, default_value = "60000")] - pub port: u16, + pub host: BindAddress, - /// Represents total range of ports to try if a port is already taken - /// when binding, applying range incrementally against the specified - /// port (e.g. 60000-61000 inclusively if range is 1000) - #[structopt(long, default_value = "1000")] - pub port_range: u16, + // Set the port(s) that the server will attempt to bind to + // + // This can be in the form of PORT1 or PORT1:PORTN to provide a range of ports. + // With -p 0, the server will let the operating system pick an available TCP port. + // + // Please note that this option does not affect the server-side port used by SSH + #[structopt(short, long, default_value = "60000:61000")] + pub port: PortRange, } diff --git a/src/subcommand/listen.rs b/src/subcommand/listen.rs index 0db0032..f1253f6 100644 --- a/src/subcommand/listen.rs +++ b/src/subcommand/listen.rs @@ -1,4 +1,4 @@ -use crate::opt::ListenSubcommand; +use crate::opt::{ConvertToIpAddrError, ListenSubcommand}; use derive_more::{Display, Error, From}; use fork::{daemon, Fork}; use orion::aead::SecretKey; @@ -9,49 +9,52 @@ pub type Result = std::result::Result<(), Error>; #[derive(Debug, Display, Error, From)] pub enum Error { + ConvertToIpAddrError(ConvertToIpAddrError), ForkError, IoError(io::Error), Utf8Error(FromUtf8Error), } pub fn run(cmd: ListenSubcommand) -> Result { - // TODO: Determine actual port bound to pre-fork if possible... - // - // 1. See if we can bind to a tcp port and then fork - // 2. If not, we can still output to stdout in the child process (see publish_data); so, - // would just bind early in the child process - let port = cmd.port; - let key = SecretKey::default(); - if cmd.daemon { // NOTE: We keep the stdin, stdout, stderr open so we can print out the pid with the parent match daemon(false, true) { Ok(Fork::Child) => { - publish_data(port, &key); - - // For the child, we want to fully disconnect it from pipes, which we do now + let rt = tokio::runtime::Runtime::new()?; + rt.block_on(async { run_async(cmd, true).await })?; + } + Ok(Fork::Parent(pid)) => { + eprintln!("[distant detached, pid = {}]", pid); if let Err(_) = fork::close_fd() { return Err(Error::ForkError); } - - let rt = tokio::runtime::Runtime::new()?; - rt.block_on(async { run_async(cmd).await })?; } - Ok(Fork::Parent(pid)) => eprintln!("[distant detached, pid = {}]", pid), Err(_) => return Err(Error::ForkError), } } else { - publish_data(port, &key); - let rt = tokio::runtime::Runtime::new()?; - rt.block_on(async { run_async(cmd).await })?; + rt.block_on(async { run_async(cmd, false).await })?; } // MAC -> Decrypt Ok(()) } -async fn run_async(_cmd: ListenSubcommand) -> Result { +async fn run_async(cmd: ListenSubcommand, is_forked: bool) -> Result { + let addr = cmd.host.to_ip_addr()?; + let socket_addrs = cmd.port.make_socket_addrs(addr); + let listener = tokio::net::TcpListener::bind(socket_addrs.as_slice()).await?; + let port = listener.local_addr()?.port(); + let key = SecretKey::default(); + publish_data(port, &key); + + // For the child, we want to fully disconnect it from pipes, which we do now + if is_forked { + if let Err(_) = fork::close_fd() { + return Err(Error::ForkError); + } + } + // TODO: Implement server logic Ok(()) }