Allow custom enc key and write password

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

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

crates/sshx/src/controller.rs

@@ -34,8 +34,9 @@ pub struct Controller {
 
     name: String,
     token: String,
+    link: String,
     url: String,
-    write_url: Option<String>,
+    write_password: Option<String>,
 
     /// Channels with backpressure routing messages to each shell task.
     shells_tx: HashMap<Sid, mpsc::Sender<ShellData>>,
@@ -50,19 +51,29 @@ impl Controller {
     pub async fn new(
         origin: &str,
         name: &str,
+        encryption_key: &str,
         runner: Runner,
         enable_readers: bool,
+        write_password: &str,
     ) -> Result<Self> {
         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();
             task::spawn_blocking(move || Encrypt::new(&encryption_key))
         };
 
         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.

crates/sshx/src/main.rs

@@ -27,45 +27,70 @@ struct Args {
     #[clap(long)]
     name: Option<String>,
 
+    /// 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) {
     let version_str = match option_env!("CARGO_PKG_VERSION") {
         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 {
         println!(
             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());
         }