Refactor to create tokio runtimes within commands instead of at main

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")

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();
 }

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<dyn std::error::Error>> {
+    pub fn run(self) -> Result<(), Box<dyn std::error::Error>> {
         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<Self, Self::Err> {
+        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,
 

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(())
 }

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,

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(())
 }