diff --git a/crates/sshx-server/tests/snapshot.rs b/crates/sshx-server/tests/snapshot.rs index dde5f6e..5a18c11 100644 --- a/crates/sshx-server/tests/snapshot.rs +++ b/crates/sshx-server/tests/snapshot.rs @@ -16,7 +16,8 @@ pub mod common; async fn test_basic_restore() -> Result<()> { let server = TestServer::new().await; - let mut controller = Controller::new(&server.endpoint(), "", Runner::Echo, false).await?; + let mut controller = + Controller::new(&server.endpoint(), "", "", Runner::Echo, false, "").await?; let name = controller.name().to_owned(); let key = controller.encryption_key().to_owned(); tokio::spawn(async move { controller.run().await }); diff --git a/crates/sshx-server/tests/with_client.rs b/crates/sshx-server/tests/with_client.rs index d52222f..14b6484 100644 --- a/crates/sshx-server/tests/with_client.rs +++ b/crates/sshx-server/tests/with_client.rs @@ -14,7 +14,7 @@ pub mod common; #[tokio::test] async fn test_handshake() -> Result<()> { let server = TestServer::new().await; - let controller = Controller::new(&server.endpoint(), "", Runner::Echo, false).await?; + let controller = Controller::new(&server.endpoint(), "", "", Runner::Echo, false, "").await?; controller.close().await?; Ok(()) } @@ -23,7 +23,8 @@ async fn test_handshake() -> Result<()> { async fn test_command() -> Result<()> { let server = TestServer::new().await; let runner = Runner::Shell("/bin/bash".into()); - let mut controller = Controller::new(&server.endpoint(), "", runner, false).await?; + let mut controller = + Controller::new(&server.endpoint(), "", "", Runner::Echo, false, "").await?; let session = server .state() @@ -71,7 +72,8 @@ async fn test_ws_missing() -> Result<()> { async fn test_ws_basic() -> Result<()> { let server = TestServer::new().await; - let mut controller = Controller::new(&server.endpoint(), "", Runner::Echo, false).await?; + let mut controller = + Controller::new(&server.endpoint(), "", "", Runner::Echo, false, "").await?; let name = controller.name().to_owned(); let key = controller.encryption_key().to_owned(); tokio::spawn(async move { controller.run().await }); @@ -103,7 +105,8 @@ async fn test_ws_basic() -> Result<()> { async fn test_ws_resize() -> Result<()> { let server = TestServer::new().await; - let mut controller = Controller::new(&server.endpoint(), "", Runner::Echo, false).await?; + let mut controller = + Controller::new(&server.endpoint(), "", "", Runner::Echo, false, "").await?; let name = controller.name().to_owned(); let key = controller.encryption_key().to_owned(); tokio::spawn(async move { controller.run().await }); @@ -147,7 +150,8 @@ async fn test_ws_resize() -> Result<()> { async fn test_users_join() -> Result<()> { let server = TestServer::new().await; - let mut controller = Controller::new(&server.endpoint(), "", Runner::Echo, false).await?; + let mut controller = + Controller::new(&server.endpoint(), "", "", Runner::Echo, false, "").await?; let name = controller.name().to_owned(); let key = controller.encryption_key().to_owned(); tokio::spawn(async move { controller.run().await }); @@ -176,7 +180,8 @@ async fn test_users_join() -> Result<()> { async fn test_users_metadata() -> Result<()> { let server = TestServer::new().await; - let mut controller = Controller::new(&server.endpoint(), "", Runner::Echo, false).await?; + let mut controller = + Controller::new(&server.endpoint(), "", "", Runner::Echo, false, "").await?; let name = controller.name().to_owned(); let key = controller.encryption_key().to_owned(); tokio::spawn(async move { controller.run().await }); @@ -201,7 +206,8 @@ async fn test_users_metadata() -> Result<()> { async fn test_chat_messages() -> Result<()> { let server = TestServer::new().await; - let mut controller = Controller::new(&server.endpoint(), "", Runner::Echo, false).await?; + let mut controller = + Controller::new(&server.endpoint(), "", "", Runner::Echo, false, "").await?; let name = controller.name().to_owned(); let key = controller.encryption_key().to_owned(); tokio::spawn(async move { controller.run().await }); @@ -234,24 +240,20 @@ async fn test_read_write_permissions() -> Result<()> { let server = TestServer::new().await; // create controller with read-only mode enabled - let mut controller = Controller::new(&server.endpoint(), "", Runner::Echo, true).await?; + let mut controller = + Controller::new(&server.endpoint(), "", "", Runner::Echo, true, "").await?; let name = controller.name().to_owned(); let key = controller.encryption_key().to_owned(); - let write_url = controller - .write_url() - .expect("Should have write URL when enable_readers is true") + let write_password = controller + .write_password() + .expect("Should have write password when enable_readers is true") .to_string(); tokio::spawn(async move { controller.run().await }); - let write_password = write_url - .split(',') - .nth(1) - .expect("Write URL should contain password"); - // connect with write access let mut writer = - ClientSocket::connect(&server.ws_endpoint(&name), &key, Some(write_password)).await?; + ClientSocket::connect(&server.ws_endpoint(&name), &key, Some(&write_password)).await?; writer.flush().await; // test write permissions diff --git a/crates/sshx/src/controller.rs b/crates/sshx/src/controller.rs index 5eff767..e7b741e 100644 --- a/crates/sshx/src/controller.rs +++ b/crates/sshx/src/controller.rs @@ -34,8 +34,9 @@ pub struct Controller { name: String, token: String, + link: String, url: String, - write_url: Option, + write_password: Option, /// Channels with backpressure routing messages to each shell task. shells_tx: HashMap>, @@ -50,11 +51,17 @@ impl Controller { pub async fn new( origin: &str, name: &str, + encryption_key: &str, runner: Runner, enable_readers: bool, + write_password: &str, ) -> Result { debug!(%origin, "connecting to server"); - let encryption_key = rand_alphanumeric(14); // 83.3 bits of entropy + let encryption_key = if encryption_key.is_empty() { + rand_alphanumeric(14) // 83.3 bits of entropy + } else { + encryption_key.into() + }; let kdf_task = { let encryption_key = encryption_key.clone(); @@ -62,7 +69,11 @@ impl Controller { }; let (write_password, kdf_write_password_task) = if enable_readers { - let write_password = rand_alphanumeric(14); // 83.3 bits of entropy + let write_password = if write_password.is_empty() { + rand_alphanumeric(14) + } else { + write_password.to_string() + }; // 83.3 bits of entropy let task = { let write_password = write_password.clone(); task::spawn_blocking(move || Encrypt::new(&write_password)) @@ -87,14 +98,9 @@ impl Controller { write_password_hash, }; let mut resp = client.open(req).await?.into_inner(); + let link = resp.url.clone(); resp.url = resp.url + "#" + &encryption_key; - let write_url = if let Some(write_password) = write_password { - Some(resp.url.clone() + "," + &write_password) - } else { - None - }; - let (output_tx, output_rx) = mpsc::channel(64); Ok(Self { origin: origin.into(), @@ -103,8 +109,9 @@ impl Controller { encryption_key, name: resp.name, token: resp.token, + link, url: resp.url, - write_url, + write_password, shells_tx: HashMap::new(), output_tx, output_rx, @@ -125,14 +132,19 @@ impl Controller { &self.name } + /// Returns the URL of the session without auth. + pub fn link(&self) -> &str { + &self.link + } + /// Returns the URL of the session. pub fn url(&self) -> &str { &self.url } - /// Returns the write URL of the session, if it exists. - pub fn write_url(&self) -> Option<&str> { - self.write_url.as_deref() + /// Returns the write token of the session, if it exists. + pub fn write_password(&self) -> Option<&str> { + self.write_password.as_deref() } /// Returns the encryption key for this session, hidden from the server. diff --git a/crates/sshx/src/main.rs b/crates/sshx/src/main.rs index d91f50e..7f0ca78 100644 --- a/crates/sshx/src/main.rs +++ b/crates/sshx/src/main.rs @@ -27,10 +27,18 @@ struct Args { #[clap(long)] name: Option, + /// Encryption key + #[clap(long, default_value = "")] + encryption_key: String, + /// Enable read-only access mode - generates separate URLs for viewers and /// editors. #[clap(long)] enable_readers: bool, + + /// Write permission key + #[clap(long, default_value = "")] + write_password: String, } fn print_greeting(shell: &str, controller: &Controller) { @@ -38,20 +46,31 @@ fn print_greeting(shell: &str, controller: &Controller) { Some(version) => format!("v{version}"), None => String::from("[dev]"), }; - if let Some(write_url) = controller.write_url() { + if let Some(write_password) = controller.write_password() { println!( r#" {sshx} {version} + {arr} link: {link} {arr} Read-only link: {link_v} {arr} Writable link: {link_e} {arr} Shell: {shell_v} + + {eye} View: {encryption_key} + {key} Key: {write_password} "#, sshx = Green.bold().paint("sshx"), version = Green.paint(&version_str), arr = Green.paint("➜"), + eye = "👀", + key = "🔑", + link = Cyan.underline().paint(controller.link()), link_v = Cyan.underline().paint(controller.url()), - link_e = Cyan.underline().paint(write_url), + encryption_key = Cyan.underline().paint(controller.encryption_key()), + write_password = Cyan.underline().paint(write_password), + link_e = Cyan + .underline() + .paint(controller.url().to_owned() + "," + write_password), shell_v = Fixed(8).paint(shell), ); } else { @@ -59,13 +78,19 @@ fn print_greeting(shell: &str, controller: &Controller) { r#" {sshx} {version} - {arr} Link: {link_v} - {arr} Shell: {shell_v} + {arr} Link: {link} + {arr} Link with auth: {link_v} + {arr} Shell: {shell_v} + + {key} Key: {encryption_key} "#, sshx = Green.bold().paint("sshx"), version = Green.paint(&version_str), arr = Green.paint("➜"), + key = "🔑", + link = Cyan.underline().paint(controller.link()), link_v = Cyan.underline().paint(controller.url()), + encryption_key = Cyan.underline().paint(controller.encryption_key()), shell_v = Fixed(8).paint(shell), ); } @@ -90,10 +115,18 @@ async fn start(args: Args) -> Result<()> { }); let runner = Runner::Shell(shell.clone()); - let mut controller = Controller::new(&args.server, &name, runner, args.enable_readers).await?; + let mut controller = Controller::new( + &args.server, + &name, + &args.encryption_key, + runner, + args.enable_readers, + &args.write_password, + ) + .await?; if args.quiet { - if let Some(write_url) = controller.write_url() { - println!("{}", write_url); + if let Some(write_password) = controller.write_password() { + println!("{}", controller.url().to_owned() + "," + write_password); } else { println!("{}", controller.url()); }