From 508ba38550cf8af6ad18e9d0066ada95e3dae390 Mon Sep 17 00:00:00 2001 From: Ovidiu Ionescu Date: Sat, 11 Jan 2025 01:14:22 +0200 Subject: [PATCH] Added support for sessionId, secret as parameters --- crates/sshx-core/proto/sshx.proto | 1 + crates/sshx-server/src/grpc.rs | 2 +- crates/sshx-server/tests/simple.rs | 1 + crates/sshx-server/tests/snapshot.rs | 3 ++- crates/sshx-server/tests/with_client.rs | 23 +++++++++++++++-------- crates/sshx/src/controller.rs | 9 ++++++++- crates/sshx/src/main.rs | 18 +++++++++++++++++- 7 files changed, 45 insertions(+), 12 deletions(-) diff --git a/crates/sshx-core/proto/sshx.proto b/crates/sshx-core/proto/sshx.proto index 111d0dd..ddcde68 100644 --- a/crates/sshx-core/proto/sshx.proto +++ b/crates/sshx-core/proto/sshx.proto @@ -42,6 +42,7 @@ message OpenRequest { bytes encrypted_zeros = 2; // Encrypted zero block, for client verification. string name = 3; // Name of the session (user@hostname). optional bytes write_password_hash = 4; // Hashed write password, if read-only mode is enabled. + optional string session_id = 5; // Add this field } // Details of a newly-created sshx session. diff --git a/crates/sshx-server/src/grpc.rs b/crates/sshx-server/src/grpc.rs index afe6444..e89f5f5 100644 --- a/crates/sshx-server/src/grpc.rs +++ b/crates/sshx-server/src/grpc.rs @@ -48,7 +48,7 @@ impl SshxService for GrpcServer { if origin.is_empty() { return Err(Status::invalid_argument("origin is empty")); } - let name = rand_alphanumeric(10); + let name = request.session_id.unwrap_or_else(|| rand_alphanumeric(10)); info!(%name, "creating new session"); match self.0.lookup(&name) { diff --git a/crates/sshx-server/tests/simple.rs b/crates/sshx-server/tests/simple.rs index 3549030..ff6947c 100644 --- a/crates/sshx-server/tests/simple.rs +++ b/crates/sshx-server/tests/simple.rs @@ -16,6 +16,7 @@ async fn test_rpc() -> Result<()> { encrypted_zeros: Encrypt::new("").zeros().into(), name: String::new(), write_password_hash: None, + session_id: None, }; let resp = client.open(req).await?; assert!(!resp.into_inner().name.is_empty()); diff --git a/crates/sshx-server/tests/snapshot.rs b/crates/sshx-server/tests/snapshot.rs index dde5f6e..fb1bf6e 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, None, None).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..c9b2c0c 100644 --- a/crates/sshx-server/tests/with_client.rs +++ b/crates/sshx-server/tests/with_client.rs @@ -14,7 +14,8 @@ 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, None, None).await?; controller.close().await?; Ok(()) } @@ -23,7 +24,7 @@ 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, false, None, None).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, None, None).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, None, None).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, None, None).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, None, None).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, None, None).await?; let name = controller.name().to_owned(); let key = controller.encryption_key().to_owned(); tokio::spawn(async move { controller.run().await }); @@ -234,7 +240,8 @@ 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, None, None).await?; let name = controller.name().to_owned(); let key = controller.encryption_key().to_owned(); let write_url = controller diff --git a/crates/sshx/src/controller.rs b/crates/sshx/src/controller.rs index 5eff767..4596428 100644 --- a/crates/sshx/src/controller.rs +++ b/crates/sshx/src/controller.rs @@ -52,9 +52,15 @@ impl Controller { name: &str, runner: Runner, enable_readers: bool, + session_id: Option, + secret: Option, ) -> Result { debug!(%origin, "connecting to server"); - let encryption_key = rand_alphanumeric(14); // 83.3 bits of entropy + + let encryption_key = match secret { + Some(s) => s, + None => rand_alphanumeric(14), // 83.3 bits of entropy + }; let kdf_task = { let encryption_key = encryption_key.clone(); @@ -85,6 +91,7 @@ impl Controller { encrypted_zeros: encrypt.zeros().into(), name: name.into(), write_password_hash, + session_id, }; let mut resp = client.open(req).await?.into_inner(); resp.url = resp.url + "#" + &encryption_key; diff --git a/crates/sshx/src/main.rs b/crates/sshx/src/main.rs index 72faeb7..b1f96a0 100644 --- a/crates/sshx/src/main.rs +++ b/crates/sshx/src/main.rs @@ -31,6 +31,14 @@ struct Args { /// editors. #[clap(long)] enable_readers: bool, + + /// Optional custom session ID. + #[clap(long)] + session_id: Option, + + /// Optional encryption key. + #[clap(long)] + secret: Option, } fn print_greeting(shell: &str, controller: &Controller) { @@ -90,7 +98,15 @@ 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, + runner, + args.enable_readers, + args.session_id, + args.secret, + ) + .await?; if args.quiet { println!("{}", controller.url()); } else {