diff --git a/bin/cli/src/command/mod.rs b/bin/cli/src/command/mod.rs index 08b736d11..49c0429da 100644 --- a/bin/cli/src/command/mod.rs +++ b/bin/cli/src/command/mod.rs @@ -18,7 +18,7 @@ pub mod container; pub mod database; pub mod execute; pub mod list; -pub mod ssh; +pub mod terminal; pub mod update; async fn komodo_client() -> anyhow::Result<&'static KomodoClient> { diff --git a/bin/cli/src/command/ssh.rs b/bin/cli/src/command/terminal.rs similarity index 76% rename from bin/cli/src/command/ssh.rs rename to bin/cli/src/command/terminal.rs index ace225271..0e16bc280 100644 --- a/bin/cli/src/command/ssh.rs +++ b/bin/cli/src/command/terminal.rs @@ -4,19 +4,73 @@ use colored::Colorize; use futures_util::{SinkExt, StreamExt}; use komodo_client::{ api::write::{CreateTerminal, TerminalRecreateMode}, - entities::config::cli::args::ssh::Ssh, + entities::config::cli::args::terminal::{Connect, Exec}, +}; +use tokio::{ + io::{AsyncReadExt as _, AsyncWriteExt as _}, + net::TcpStream, +}; +use tokio_tungstenite::{ + MaybeTlsStream, WebSocketStream, tungstenite, }; -use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; -use tokio_tungstenite::tungstenite; use tokio_util::sync::CancellationToken; -pub async fn handle( - Ssh { +pub async fn handle_connect( + Connect { server, name, command, recreate, - }: &Ssh, + }: &Connect, +) -> anyhow::Result<()> { + handle_terminal_forwarding(async { + let client = super::komodo_client().await?; + // Init the terminal if it doesn't exist already. + client + .write(CreateTerminal { + server: server.to_string(), + name: name.to_string(), + command: command.clone(), + recreate: if *recreate { + TerminalRecreateMode::Always + } else { + TerminalRecreateMode::DifferentCommand + }, + }) + .await?; + client.connect_terminal_websocket(server, name).await + }) + .await +} + +pub async fn handle_exec( + Exec { + server, + container, + shell, + recreate, + }: &Exec, +) -> anyhow::Result<()> { + handle_terminal_forwarding(async { + super::komodo_client() + .await? + .connect_container_websocket( + server, + container, + shell, + recreate.then_some(TerminalRecreateMode::Always), + ) + .await + }) + .await +} + +type WsStream = WebSocketStream>; + +async fn handle_terminal_forwarding< + C: Future>, +>( + connect: C, ) -> anyhow::Result<()> { // Need to forward multiple sources into ws write let (write_tx, mut write_rx) = @@ -84,26 +138,7 @@ pub async fn handle( // CONNECT AND FORWARD // ===================== - let client = super::komodo_client().await?; - - // Init the terminal if it doesn't exist already. - client - .write(CreateTerminal { - server: server.to_string(), - name: name.to_string(), - command: command.clone(), - recreate: if *recreate { - TerminalRecreateMode::Always - } else { - TerminalRecreateMode::DifferentCommand - }, - }) - .await?; - - let (mut ws_write, mut ws_read) = client - .connect_terminal_websocket(server, name) - .await? - .split(); + let (mut ws_write, mut ws_read) = connect.await?.split(); let forward_write = async { while let Some(bytes) = diff --git a/bin/cli/src/main.rs b/bin/cli/src/main.rs index dab10130f..621ca1c52 100644 --- a/bin/cli/src/main.rs +++ b/bin/cli/src/main.rs @@ -54,7 +54,12 @@ async fn app() -> anyhow::Result<()> { args::Command::Update { command } => { command::update::handle(command).await } - args::Command::Ssh(ssh) => command::ssh::handle(ssh).await, + args::Command::Connect(connect) => { + command::terminal::handle_connect(connect).await + } + args::Command::Exec(exec) => { + command::terminal::handle_exec(exec).await + } args::Command::Key { command } => { noise::key::command::handle(command).await } diff --git a/client/core/rs/src/api/write/server.rs b/client/core/rs/src/api/write/server.rs index af1acd30d..f1f6cb856 100644 --- a/client/core/rs/src/api/write/server.rs +++ b/client/core/rs/src/api/write/server.rs @@ -1,6 +1,7 @@ use derive_empty_traits::EmptyTraits; use resolver_api::Resolve; use serde::{Deserialize, Serialize}; +use strum::AsRefStr; use typeshare::typeshare; use crate::entities::{ @@ -134,7 +135,9 @@ pub struct CreateNetwork { /// Configures the behavior of [CreateTerminal] if the /// specified terminal name already exists. #[typeshare] -#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)] +#[derive( + Debug, Clone, Copy, Default, Serialize, Deserialize, AsRefStr, +)] pub enum TerminalRecreateMode { /// Never kill the old terminal if it already exists. /// If the command is different, returns error. diff --git a/client/core/rs/src/entities/config/cli/args/mod.rs b/client/core/rs/src/entities/config/cli/args/mod.rs index 07af7a7b8..b7c90bdc0 100644 --- a/client/core/rs/src/entities/config/cli/args/mod.rs +++ b/client/core/rs/src/entities/config/cli/args/mod.rs @@ -7,7 +7,7 @@ use crate::api::execute::Execution; pub mod container; pub mod database; pub mod list; -pub mod ssh; +pub mod terminal; pub mod update; #[derive(Debug, clap::Parser)] @@ -81,7 +81,12 @@ pub enum Command { command: update::UpdateCommand, }, - Ssh(ssh::Ssh), + /// Connect to Server Terminals. (alias: `ssh`) + #[clap(alias = "ssh")] + Connect(terminal::Connect), + + /// Connect to Container Terminals. + Exec(terminal::Exec), /// Private-Public key utilities. (alias: `k`) #[clap(alias = "k")] diff --git a/client/core/rs/src/entities/config/cli/args/ssh.rs b/client/core/rs/src/entities/config/cli/args/terminal.rs similarity index 52% rename from client/core/rs/src/entities/config/cli/args/ssh.rs rename to client/core/rs/src/entities/config/cli/args/terminal.rs index 840fc695d..ddb6d5b95 100644 --- a/client/core/rs/src/entities/config/cli/args/ssh.rs +++ b/client/core/rs/src/entities/config/cli/args/terminal.rs @@ -1,5 +1,5 @@ #[derive(Debug, Clone, clap::Parser)] -pub struct Ssh { +pub struct Connect { /// The server to connect to. pub server: String, @@ -15,3 +15,17 @@ pub struct Ssh { #[arg(long, short = 'r', default_value_t = false)] pub recreate: bool, } + +#[derive(Debug, Clone, clap::Parser)] +pub struct Exec { + /// Specify Server + pub server: String, + /// The container (name) to connect to. + /// Will error if matches multiple containers but no Server is defined. + pub container: String, + /// The shell, eg `bash`. + pub shell: String, + /// Force fresh terminal to replace existing one. + #[arg(long, short = 'r', default_value_t = false)] + pub recreate: bool, +} diff --git a/client/core/rs/src/ws/mod.rs b/client/core/rs/src/ws/mod.rs index 9187b190c..a293013d1 100644 --- a/client/core/rs/src/ws/mod.rs +++ b/client/core/rs/src/ws/mod.rs @@ -9,7 +9,7 @@ use tokio_tungstenite::{ }; use typeshare::typeshare; -use crate::KomodoClient; +use crate::{KomodoClient, api::write::TerminalRecreateMode}; pub mod update; @@ -47,6 +47,26 @@ impl KomodoClient { .await } + pub async fn connect_container_websocket( + &self, + server: &str, + container: &str, + shell: &str, + recreate: Option, + ) -> anyhow::Result>> { + let mut query = + format!("server={server}&container={container}&shell={shell}"); + if let Some(recreate) = recreate { + let _ = write!(&mut query, "&recreate={}", recreate.as_ref()); + } + self + .connect_login_user_websocket( + "/container/terminal", + Some(&query), + ) + .await + } + async fn connect_login_user_websocket( &self, path: &str,