Add support for localhost.run

2024-09-07 11:04:28 -03:00 · 2024-09-07 11:04:28 -03:00 · 9dc4254647
commit 9dc4254647
parent 00b362621f
4 changed files with 259 additions and 99 deletions
--- a/Cargo.lock
+++ b/Cargo.lock
@ -118,9 +118,9 @@ dependencies = [

 [[package]]
 name = "anyhow"
-version = "1.0.86"
+version = "1.0.87"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b3d1d046238990b9cf5bcde22a3fb3584ee5cf65fb2765f454ed428c7a0063da"
+checksum = "10f00e1f6e58a40e807377c75c6a7f97bf9044fab57816f2414e6f5f4499d7b8"

 [[package]]
 name = "async-trait"
@ -295,9 +295,9 @@ dependencies = [

 [[package]]
 name = "cc"
-version = "1.1.16"
+version = "1.1.18"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e9d013ecb737093c0e86b151a7b837993cf9ec6c502946cfb44bedc392421e0b"
+checksum = "b62ac837cdb5cb22e10a256099b4fc502b1dfe560cb282963a974d7abd80e476"
 dependencies = [
 "shlex",
 ]
@ -383,9 +383,9 @@ checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8"

 [[package]]
 name = "cpufeatures"
-version = "0.2.13"
+version = "0.2.14"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "51e852e6dc9a5bed1fae92dd2375037bf2b768725bf3be87811edee3249d09ad"
+checksum = "608697df725056feaccfa42cffdaeeec3fccc4ffc38358ecd19b243e716a78e0"
 dependencies = [
 "libc",
 ]
@ -1533,6 +1533,7 @@ dependencies = [
 "hyper",
 "hyper-util",
 "russh",
+ "termsize",
 "tokio",
 "tokio-stream",
 "tokio-util",
@ -1678,18 +1679,18 @@ checksum = "61697e0a1c7e512e84a621326239844a24d8207b4669b41bc18b32ea5cbf988b"

 [[package]]
 name = "serde"
-version = "1.0.209"
+version = "1.0.210"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "99fce0ffe7310761ca6bf9faf5115afbc19688edd00171d81b1bb1b116c63e09"
+checksum = "c8e3592472072e6e22e0a54d5904d9febf8508f65fb8552499a1abc7d1078c3a"
 dependencies = [
 "serde_derive",
 ]

 [[package]]
 name = "serde_derive"
-version = "1.0.209"
+version = "1.0.210"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a5831b979fd7b5439637af1752d535ff49f4860c0f341d1baeb6faf0f4242170"
+checksum = "243902eda00fad750862fc144cea25caca5e20d615af0a81bee94ca738f1df1f"
 dependencies = [
 "proc-macro2",
 "quote",
@ -1698,9 +1699,9 @@ dependencies = [

 [[package]]
 name = "serde_json"
-version = "1.0.127"
+version = "1.0.128"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8043c06d9f82bd7271361ed64f415fe5e12a77fdb52e573e7f06a516dea329ad"
+checksum = "6ff5456707a1de34e7e37f2a6fd3d3f808c318259cbd01ab6377795054b483d8"
 dependencies = [
 "itoa",
 "memchr",
@ -1913,6 +1914,16 @@ version = "1.0.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a7065abeca94b6a8a577f9bd45aa0867a2238b74e8eb67cf10d492bc39351394"

+[[package]]
+name = "termsize"
+version = "0.1.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6f11ff5c25c172608d5b85e2fb43ee9a6d683a7f4ab7f96ae07b3d8b590368fd"
+dependencies = [
+ "libc",
+ "winapi",
+]
+
 [[package]]
 name = "thiserror"
 version = "1.0.63"
@ -1974,9 +1985,9 @@ dependencies = [

 [[package]]
 name = "tokio-stream"
-version = "0.1.15"
+version = "0.1.16"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "267ac89e0bec6e691e5813911606935d77c476ff49024f98abcea3e7b15e37af"
+checksum = "4f4e6ce100d0eb49a2734f8c0812bcd324cf357d21810932c5df6b96ef2b86f1"
 dependencies = [
 "futures-core",
 "pin-project-lite",
@ -1986,9 +1997,9 @@ dependencies = [

 [[package]]
 name = "tokio-util"
-version = "0.7.11"
+version = "0.7.12"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9cf6b47b3771c49ac75ad09a6162f53ad4b8088b76ac60e8ec1455b31a189fe1"
+checksum = "61e7c3654c13bcd040d4a03abee2c75b1d14a37b423cf5a813ceae1cc903ec6a"
 dependencies = [
 "bytes",
 "futures-core",
--- a/Cargo.toml
+++ b/Cargo.toml
@ -18,6 +18,7 @@ http = "1.1.0"
 hyper = { version = "1", features = ["full"] }
 hyper-util = { version = "0.1", features = ["full"] }
 russh = "0.45"
+termsize = "0.1.9"
 tokio = { version = "1", features = ["full"] }
 tokio-stream = { version = "0.1.15", features = ["net", "sync"] }
 tokio-util = "0.7.11"
--- a/README.md
+++ b/README.md
@ -2,4 +2,18 @@

 A Rust project demonstrating how to serve Axum's HTTP server on a remote host's port, using SSH tunneling and streaming to avoid opening a socket on the client.

-Tokio, Tower, hyper, and `async` are responsible for gluing everything together. They are pretty awesome! The hardest part to implement was Axum's half; mainly, figuring out how to accept a streaming socket instead of the default TcpListener.
+Tokio, Tower, and hyper are responsible for gluing everything together with async. They are pretty awesome!
+
+## Usage
+
+With [`localhost.run`](https://localhost.run/):
+
+```sh
+cargo run -- localhost.run -i ~/.ssh/id_ed25519 -l username --request-pty ""
+```
+
+With [`sish`](https://github.com/antoniomika/sish):
+
+```sh
+cargo run -- tuns.sh -i ~/.ssh/id_ed25519 -R test
+```
--- a/src/main.rs
+++ b/src/main.rs
@ -8,7 +8,7 @@ use std::{
 use anyhow::{anyhow, Context, Result};
 use async_trait::async_trait;
 use axum::{extract::State, routing::get, Router};
-use clap::Parser;
+use clap::{Args, Parser};
 use hyper::service::service_fn;
 use hyper_util::{
    rt::{TokioExecutor, TokioIo},
@ -17,10 +17,14 @@ use hyper_util::{
 use russh::{
    client::{self, Config, Handle, KeyboardInteractiveAuthResponse, Msg, Session},
    keys::{decode_secret_key, key},
-    Channel, ChannelMsg, Disconnect,
+    Channel, ChannelId, ChannelMsg, Disconnect,
 };
 use tokio::io::AsyncWriteExt;
-use tokio::{fs, io::stdout, time::sleep};
+use tokio::{
+    fs,
+    io::{stderr, stdout},
+    time::sleep,
+};
 use tower::Service;
 use tracing::{debug, debug_span, error, info, trace, warn};
 use tracing_subscriber::{fmt, prelude::*, EnvFilter};
@ -32,24 +36,42 @@ use tracing_subscriber::{fmt, prelude::*, EnvFilter};
 #[command(version, about, long_about = None)]
 struct ClapArgs {
    /// SSH hostname
-    #[arg(short = 'H', long)]
    host: String,

+    /// Identity file containing private key
+    #[arg(short, long, default_value_t = String::from(""))]
+    login_name: String,
+
    /// SSH port
    #[arg(short, long, default_value_t = 22)]
    port: u16,

-    /// Identity file containing private key
-    #[arg(short, long)]
-    identity_file: Option<PathBuf>,
+    #[command(flatten)]
+    auth: Option<Authentication>,

    /// Remote hostname to bind to
-    #[arg(short, long, default_value_t = String::from("localhost"))]
+    #[arg(short = 'R', long, default_value_t = String::from(""))]
    remote_host: String,

    /// Remote port to bind to
-    #[arg(short = 't', long, default_value_t = 80)]
+    #[arg(short = 'P', long, default_value_t = 80)]
    remote_port: u16,
+
+    /// Request a pseudo-terminal to be allocated with the given command.
+    #[arg(long)]
+    request_pty: Option<String>,
+}
+
+#[derive(Args, Debug)]
+#[group(required = false, multiple = false)]
+struct Authentication {
+    /// Identity file containing private key.
+    #[arg(short, long, value_name = "FILE")]
+    identity_file: Option<PathBuf>,
+
+    /// Request keyboard-interactive based SSH authentication.
+    #[arg(long)]
+    keyboard_interactive: bool,
 }

 #[tokio::main]
@ -60,13 +82,21 @@ async fn main() -> Result<()> {
        .init();
    trace!("Tracing is up!");
    let args = ClapArgs::parse();
-    let secret_key = match args.identity_file {
+    let session_auth = match args.auth {
        None => None,
-        Some(file) => {
-            let secret_key = fs::read_to_string(file)
-                .await
-                .with_context(|| "Failed to open secret key")?;
-            Some(decode_secret_key(&secret_key, None).with_context(|| "Invalid secret key")?)
+        Some(auth) => {
+            if auth.keyboard_interactive {
+                Some(SessionAuth::KeyboardInteractive)
+            } else if let Some(file) = auth.identity_file {
+                let secret_key = fs::read_to_string(file)
+                    .await
+                    .with_context(|| "Failed to open secret key")?;
+                Some(SessionAuth::SecretKey(Arc::new(
+                    decode_secret_key(&secret_key, None).with_context(|| "Invalid secret key")?,
+                )))
+            } else {
+                unreachable!();
+            }
        }
    };
    let config = Arc::new(client::Config {
@ -75,18 +105,23 @@ async fn main() -> Result<()> {
    let mut session = TcpForwardSession::connect(
        &args.host,
        args.port,
+        &args.login_name,
        config,
-        secret_key.map(|key| Arc::new(key)),
+        &session_auth,
    )
    .await
    .with_context(|| "Initial connection failed")?;
    loop {
        match session
-            .start_forwarding(&args.remote_host, args.remote_port)
+            .start_forwarding(
+                &args.remote_host,
+                args.remote_port,
+                args.request_pty.as_deref(),
+            )
            .await
        {
            Err(e) => error!(error = ?e, "TCP forward session failed."),
-            _ => info!("Connection closed."),
+            Ok(code) => info!("Connection closed with code {}.", code),
        }
        debug!("Attempting graceful disconnect.");
        if let Err(e) = session.close().await {
@ -98,6 +133,7 @@ async fn main() -> Result<()> {
            .reconnect_with(
                &args.host,
                args.port,
+                &args.login_name,
                iter::from_fn(move || {
                    reconnect_attempt += 1;
                    if reconnect_attempt <= 5 {
@ -141,10 +177,17 @@ async fn hello(State(state): State<AppState>) -> String {

 /* Russh session and client */

+/// Private type to decide on the authentication method.
+#[derive(Clone)]
+enum SessionAuth {
+    SecretKey(Arc<key::KeyPair>),
+    KeyboardInteractive,
+}
+
 /// User-implemented session type as a helper for interfacing with the SSH protocol.
 struct TcpForwardSession {
    config: Arc<Config>,
-    secret_key: Option<Arc<key::KeyPair>>,
+    session_auth: Option<SessionAuth>,
    session: Handle<Client>,
 }

@ -154,8 +197,9 @@ impl TcpForwardSession {
    async fn connect(
        host: &str,
        port: u16,
+        login_name: &str,
        config: Arc<Config>,
-        secret_key: Option<Arc<key::KeyPair>>,
+        session_auth: &Option<SessionAuth>,
    ) -> Result<Self> {
        let span = debug_span!("TcpForwardSession.connect");
        let _enter = span;
@ -164,101 +208,133 @@ impl TcpForwardSession {
        let mut session = client::connect(Arc::clone(&config), (host, port), client)
            .await
            .with_context(|| "Unable to connect to remote host.")?;
-        let authentication_result = match secret_key.as_ref() {
-            None => None,
-            Some(secret_key) => {
+        let session = match session_auth {
+            Some(SessionAuth::SecretKey(ref secret_key)) => {
                if session
-                    .authenticate_publickey("root", Arc::clone(&secret_key))
+                    .authenticate_publickey(login_name, Arc::clone(secret_key))
                    .await
                    .with_context(|| "Error while authenticating with public key.")?
                {
                    debug!("Public key authentication succeeded!");
-                    Some(Ok(()))
+                    Ok(session)
                } else {
-                    Some(Err(anyhow!("Public key authentication failed.")))
+                    Err(anyhow!("Public key authentication failed."))
+                }
+            }
+            Some(SessionAuth::KeyboardInteractive) => {
+                match session
+                    .authenticate_keyboard_interactive_start(login_name, None)
+                    .await
+                    .with_context(|| {
+                        "Error while authenticating with keyboard interactive session."
+                    })? {
+                    KeyboardInteractiveAuthResponse::Success => {
+                        debug!("Keyboard interactive authentication succeeded!");
+                        Ok(session)
+                    }
+                    KeyboardInteractiveAuthResponse::Failure => {
+                        Err(anyhow!("Keyboard interactive authentication failed."))
+                    }
+                    response => Err(anyhow!(
+                        "Unhandled keyboard interactive authentication event {:?}",
+                        response
+                    )),
+                }
+            }
+            None => {
+                if session
+                    .authenticate_none(login_name)
+                    .await
+                    .with_context(|| "Error while authenticating without credentials.")?
+                {
+                    debug!("Authentication without credentials succeeded!");
+                    Ok(session)
+                } else {
+                    Err(anyhow!("Authentication without credentials failed."))
                }
            }
        };
-        if matches!(authentication_result, None | Some(Err(_))) {
-            if authentication_result.is_some() {
-                debug!(
-                    "Public key authentication failed; trying keyboard interactive authentication..."
-                );
-            }
-            match session
-                .authenticate_keyboard_interactive_start("russh-axum-tcpip-forward", None)
-                .await
-                .with_context(|| "Error while authenticating with keyboard interactive session.")?
-            {
-                KeyboardInteractiveAuthResponse::Success => {
-                    debug!("Keyboard interactive authentication succeeded!");
-                }
-                KeyboardInteractiveAuthResponse::Failure => match authentication_result {
-                    None => return Err(anyhow!("Keyboard interactive authentication failed.")),
-                    Some(Err(result)) => {
-                        debug!("Keyboard interactive authentication failed; propagating public key authentication error...");
-                        return Err(result);
-                    }
-                    _ => unreachable!(),
-                },
-                response => match authentication_result {
-                    None => {
-                        return Err(anyhow!(
-                            "Unhandled keyboard interactive authentication event {:?}",
-                            response
-                        ))
-                    }
-                    Some(Err(result)) => {
-                        debug!("Keyboard interactive authentication failed; propagating public key authentication error...");
-                        return Err(result);
-                    }
-                    _ => unreachable!(),
-                },
-            }
+        match session {
+            Ok(session) => Ok(Self {
+                config,
+                session,
+                session_auth: session_auth.clone(),
+            }),
+            Err(e) => Err(e),
        }
-        Ok(Self {
-            config,
-            session,
-            secret_key,
-        })
    }

    /// Sends a port forwarding request and opens a session to receive miscellaneous data.
    /// The function yields when the session is broken (for example, if the connection was lost).
-    async fn start_forwarding(&mut self, remote_host: &str, remote_port: u16) -> Result<u32> {
+    async fn start_forwarding(
+        &mut self,
+        remote_host: &str,
+        remote_port: u16,
+        request_pty: Option<&str>,
+    ) -> Result<u32> {
        let span = debug_span!("TcpForwardSession.start");
        let _enter = span;
        self.session
            .tcpip_forward(remote_host, remote_port.into())
            .await
            .with_context(|| "tcpip_forward error.")?;
+        debug!("Requested tcpip_forward session.");
        let mut channel = self
            .session
            .channel_open_session()
            .await
            .with_context(|| "channel_open_session error.")?;
        debug!("Created open session channel.");
+        // let mut stdin = stdin();
        let mut stdout = stdout();
-        let mut code = 0;
-        loop {
+        let mut stderr = stderr();
+        if let Some(cmd) = request_pty {
+            let size = termsize::get().unwrap();
+            channel
+                .request_pty(
+                    false,
+                    &std::env::var("TERM").unwrap_or("xterm".into()),
+                    size.cols.into(),
+                    size.rows.into(),
+                    0,
+                    0,
+                    &[],
+                )
+                .await
+                .with_context(|| "Unable to request pseudo-terminal.")?;
+            debug!("Requested pseudo-terminal.");
+            channel
+                .exec(true, cmd)
+                .await
+                .with_context(|| "Unable to execute command for pseudo-terminal.")?;
+        };
+        let code = loop {
            let Some(msg) = channel.wait().await else {
                return Err(anyhow!("Unexpected end of channel."));
            };
-            trace!("Got a message!");
+            trace!("Got a message through initial session!");
            match msg {
                ChannelMsg::Data { ref data } => {
                    stdout.write_all(data).await?;
                    stdout.flush().await?;
                }
-                ChannelMsg::Close => break,
+                ChannelMsg::ExtendedData { ref data, ext: 1 } => {
+                    stderr.write_all(data).await?;
+                    stderr.flush().await?;
+                }
+                ChannelMsg::Success => (),
+                ChannelMsg::Close => break 0,
                ChannelMsg::ExitStatus { exit_status } => {
                    debug!("Exited with code {exit_status}");
-                    channel.eof().await?;
-                    code = exit_status;
+                    channel
+                        .eof()
+                        .await
+                        .with_context(|| "Unable to close connection.")?;
+                    break exit_status;
                }
                msg => return Err(anyhow!("Unknown message type {:?}.", msg)),
            }
-        }
+        };
        Ok(code)
    }

@ -271,12 +347,17 @@ impl TcpForwardSession {
        self,
        host: &str,
        port: u16,
+        login_name: &str,
        timer_iterator: impl Iterator<Item = Duration>,
    ) -> Result<Self> {
        let TcpForwardSession {
-            config, secret_key, ..
+            config,
+            session_auth,
+            ..
        } = self;
-        match TcpForwardSession::connect(host, port, config.clone(), secret_key.clone()).await {
+        match TcpForwardSession::connect(host, port, login_name, config.clone(), &session_auth)
+            .await
+        {
            Err(err) => {
                let mut e = err;
                for (i, duration) in timer_iterator.enumerate() {
@ -285,8 +366,9 @@ impl TcpForwardSession {
                    e = match TcpForwardSession::connect(
                        host,
                        port,
+                        login_name,
                        config.clone(),
-                        secret_key.clone(),
+                        &session_auth,
                    )
                    .await
                    {
@ -323,9 +405,10 @@ impl client::Handler for Client {
    type Error = anyhow::Error;

    /// Always accept the SSH server's pubkey. Don't do this in production.
+    #[allow(unused_variables)]
    async fn check_server_key(
        &mut self,
-        _server_public_key: &key::PublicKey,
+        server_public_key: &key::PublicKey,
    ) -> Result<bool, Self::Error> {
        Ok(true)
    }
@ -338,6 +421,7 @@ impl client::Handler for Client {
    /// AsyncRead/Write stream into a `hyper` IO object.
    ///
    /// See also: [axum/examples/serve-with-hyper](https://github.com/tokio-rs/axum/blob/main/examples/serve-with-hyper/src/main.rs)
+    #[allow(unused_variables)]
    async fn server_channel_open_forwarded_tcpip(
        &mut self,
        channel: Channel<Msg>,
@ -347,7 +431,7 @@ impl client::Handler for Client {
        originator_port: u32,
        session: &mut Session,
    ) -> Result<(), Self::Error> {
-        let span = debug_span!("server_channel_open_forwarded_tcpip",);
+        let span = debug_span!("server_channel_open_forwarded_tcpip");
        let _enter = span.enter();
        debug!(
            sshid = %String::from_utf8_lossy(session.remote_sshid()),
@ -357,7 +441,6 @@ impl client::Handler for Client {
            originator_port = originator_port,
            "New connection!"
        );
-        // Get our router from the lazy static.
        let router = &*ROUTER;
        let service = service_fn(move |req| router.clone().call(req));
        let server = Builder::new(TokioExecutor::new());
@ -366,8 +449,59 @@ impl client::Handler for Client {
            server
                .serve_connection_with_upgrades(TokioIo::new(channel.into_stream()), service)
                .await
-                .unwrap();
+                .expect("Invalid request");
        });
        Ok(())
    }
+
+    #[allow(unused_variables)]
+    async fn auth_banner(
+        &mut self,
+        banner: &str,
+        session: &mut Session,
+    ) -> Result<(), Self::Error> {
+        debug!("Received auth banner.");
+        let mut stdout = stdout();
+        stdout.write_all(banner.as_bytes()).await?;
+        stdout.flush().await?;
+        Ok(())
+    }
+
+    #[allow(unused_variables)]
+    async fn exit_status(
+        &mut self,
+        channel: ChannelId,
+        exit_status: u32,
+        session: &mut Session,
+    ) -> Result<(), Self::Error> {
+        debug!(channel = ?channel, "exit_status");
+        if exit_status == 0 {
+            info!("Remote exited with status {}.", exit_status);
+        } else {
+            info!("Remote exited with status {}.", exit_status);
+        }
+        Ok(())
+    }
+
+    #[allow(unused_variables)]
+    async fn channel_open_confirmation(
+        &mut self,
+        channel: ChannelId,
+        max_packet_size: u32,
+        window_size: u32,
+        session: &mut Session,
+    ) -> Result<(), Self::Error> {
+        debug!(channel = ?channel, max_packet_size, window_size, "channel_open_confirmation");
+        Ok(())
+    }
+
+    #[allow(unused_variables)]
+    async fn channel_success(
+        &mut self,
+        channel: ChannelId,
+        session: &mut Session,
+    ) -> Result<(), Self::Error> {
+        debug!(channel = ?channel, "channel_success");
+        Ok(())
+    }
 }