From 273e55fa0b0058d85a04924aca89368afc40f0d2 Mon Sep 17 00:00:00 2001 From: Chip Senkbeil Date: Mon, 26 Jul 2021 02:45:06 -0500 Subject: [PATCH] Refactor to create tokio runtimes within commands instead of at main --- src/lib.rs | 9 +++++ src/main.rs | 11 ++---- src/opt.rs | 72 +++++++++++++++++++++++++++++++++++---- src/subcommand/execute.rs | 9 +++-- src/subcommand/launch.rs | 15 ++++++-- src/subcommand/listen.rs | 52 ++++++++++++++++++++++++---- 6 files changed, 141 insertions(+), 27 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 5755ea6..9d01bf8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -14,6 +14,15 @@ lazy_static::lazy_static! { static ref SESSION_PATH: PathBuf = PROJECT_DIRS.cache_dir().join("session"); } +/// Main entrypoint into the program +pub fn run() { + let opt = Opt::load(); + init_logging(&opt.common); + if let Err(x) = opt.subcommand.run() { + eprintln!("{}", x); + } +} + pub fn init_logging(opt: &opt::CommonOpt) { stderrlog::new() .module("distant") diff --git a/src/main.rs b/src/main.rs index d7a450f..0cdcf85 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,10 +1,3 @@ -use distant::Opt; - -#[tokio::main] -async fn main() { - let opt = Opt::load(); - distant::init_logging(&opt.common); - if let Err(x) = opt.subcommand.run().await { - eprintln!("{}", x); - } +fn main() { + distant::run(); } diff --git a/src/opt.rs b/src/opt.rs index 95496bc..11f5643 100644 --- a/src/opt.rs +++ b/src/opt.rs @@ -1,6 +1,11 @@ use crate::subcommand; +use derive_more::Display; use lazy_static::lazy_static; -use std::path::PathBuf; +use std::{ + net::{AddrParseError, IpAddr, Ipv4Addr}, + path::PathBuf, + str::FromStr, +}; use structopt::StructOpt; lazy_static! { @@ -47,11 +52,11 @@ pub enum Subcommand { impl Subcommand { /// Runs the subcommand, returning the result - pub async fn run(self) -> Result<(), Box> { + pub fn run(self) -> Result<(), Box> { match self { - Self::Execute(cmd) => subcommand::execute::run(cmd).await?, - Self::Launch(cmd) => subcommand::launch::run(cmd).await?, - Self::Listen(cmd) => subcommand::listen::run(cmd).await?, + Self::Execute(cmd) => subcommand::execute::run(cmd)?, + Self::Launch(cmd) => subcommand::launch::run(cmd)?, + Self::Listen(cmd) => subcommand::listen::run(cmd)?, } Ok(()) @@ -62,6 +67,29 @@ impl Subcommand { #[derive(Debug, StructOpt)] pub struct ExecuteSubcommand {} +/// Represents options for binding a server to an IP address +#[derive(Copy, Clone, Debug, Display, PartialEq, Eq)] +pub enum BindAddress { + #[display(fmt = "ssh")] + Ssh, + #[display(fmt = "any")] + Any, + Ip(IpAddr), +} + +impl FromStr for BindAddress { + type Err = AddrParseError; + + fn from_str(s: &str) -> Result { + match s.trim() { + "ssh" => Ok(Self::Ssh), + "any" => Ok(Self::Any), + "localhost" => Ok(Self::Ip(IpAddr::V4(Ipv4Addr::LOCALHOST))), + x => x.parse(), + } + } +} + /// Represents subcommand to launch a remote server #[derive(Debug, StructOpt)] pub struct LaunchSubcommand { @@ -73,6 +101,24 @@ pub struct LaunchSubcommand { #[structopt(short, long, default_value = "distant")] pub remote_program: String, + /// Path to ssh program to execute + #[structopt(short, long, default_value = "ssh")] + pub ssh_program: String, + + /// Control the IP address that the mosh-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 + /// useful for multihomed servers. + /// + /// With --bind-server=any, the server will reply on the default interface and will not bind to + /// a particular IP address. This can be useful if the connection is made through sslh or + /// another tool that makes the SSH connection appear to come from localhost. + /// + /// With --bind-server=IP, the server will attempt to bind to the specified IP address. + #[structopt(long, default_value = "ssh")] + pub bind_server: BindAddress, + /// Username to use when sshing into remote machine #[structopt(short, long, default_value = &USERNAME)] pub username: String, @@ -101,7 +147,21 @@ pub struct ListenSubcommand { #[structopt(long)] pub no_print_startup_info: bool, - /// Represents the host to bind to when listening + /// If specified, will attempt to bind to SSH_CONNECTION instead of host + #[structopt(long)] + pub bind_ssh_connection: bool, + + /// 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 + /// useful for multihomed servers. + /// + /// 2. `any`: the server will reply on the default interface and will not bind to + /// a particular IP address. This can be useful if the connection is made through sslh or + /// another tool that makes the SSH connection appear to come from localhost. + /// + /// 3. `IP`: the server will attempt to bind to the specified IP address. #[structopt(short, long, default_value = "localhost")] pub host: String, diff --git a/src/subcommand/execute.rs b/src/subcommand/execute.rs index 7cd5a87..62aec10 100644 --- a/src/subcommand/execute.rs +++ b/src/subcommand/execute.rs @@ -25,7 +25,13 @@ pub enum Error { NoSessionFile, } -pub async fn run(_cmd: ExecuteSubcommand) -> Result { +pub fn run(cmd: ExecuteSubcommand) -> Result { + let rt = tokio::runtime::Runtime::new()?; + + rt.block_on(async { run_async(cmd).await }) +} + +async fn run_async(_cmd: ExecuteSubcommand) -> Result { // Load our session file's port and key let (port, key) = { let text = tokio::fs::read_to_string(SESSION_PATH.as_path()) @@ -51,6 +57,5 @@ pub async fn run(_cmd: ExecuteSubcommand) -> Result { hex::encode(key.unprotected_as_bytes()) ); - // Encrypt -> MAC Ok(()) } diff --git a/src/subcommand/launch.rs b/src/subcommand/launch.rs index 61fabf8..527a2a9 100644 --- a/src/subcommand/launch.rs +++ b/src/subcommand/launch.rs @@ -16,10 +16,19 @@ pub enum Error { Utf8Error(FromUtf8Error), } -pub async fn run(cmd: LaunchSubcommand) -> Result<(), Error> { - let remote_command = format!("{} listen --daemon --host 0.0.0.0", cmd.remote_program); +pub fn run(cmd: LaunchSubcommand) -> Result<(), Error> { + let rt = tokio::runtime::Runtime::new()?; + rt.block_on(async { run_async(cmd).await }) +} + +async fn run_async(cmd: LaunchSubcommand) -> Result<(), Error> { + let remote_command = format!( + "{} listen --daemon --host {}", + cmd.remote_program, cmd.bind_server + ); let ssh_command = format!( - "ssh -o StrictHostKeyChecking=no ssh://{}@{}:{} {} {}", + "{} -o StrictHostKeyChecking=no ssh://{}@{}:{} {} {}", + cmd.ssh_program, cmd.username, cmd.host, cmd.port, diff --git a/src/subcommand/listen.rs b/src/subcommand/listen.rs index bdcb109..0db0032 100644 --- a/src/subcommand/listen.rs +++ b/src/subcommand/listen.rs @@ -1,6 +1,7 @@ use crate::opt::ListenSubcommand; use derive_more::{Display, Error, From}; -use orion::aead; +use fork::{daemon, Fork}; +use orion::aead::SecretKey; use std::string::FromUtf8Error; use tokio::io; @@ -8,23 +9,60 @@ pub type Result = std::result::Result<(), Error>; #[derive(Debug, Display, Error, From)] pub enum Error { + ForkError, IoError(io::Error), Utf8Error(FromUtf8Error), } -pub async fn run(cmd: ListenSubcommand) -> Result { +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 = aead::SecretKey::default(); + 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 + 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 })?; + } + + // MAC -> Decrypt + Ok(()) +} + +async fn run_async(_cmd: ListenSubcommand) -> Result { + // TODO: Implement server logic + Ok(()) +} + +fn publish_data(port: u16, key: &SecretKey) { // TODO: We have to share the key in some manner (maybe use k256 to arrive at the same key?) // For now, we do what mosh does and print out the key knowing that this is shared over // ssh, which should provide security - print!( + println!( "DISTANT DATA {} {}", port, hex::encode(key.unprotected_as_bytes()) ); - - // MAC -> Decrypt - Ok(()) }