From 9bb31ee6fa6335c4e10f892d3a226bb5810b7591 Mon Sep 17 00:00:00 2001 From: Andrew Rioux Date: Tue, 5 Sep 2023 19:58:40 -0400 Subject: [PATCH] feat: added the ability to run commands remotely commands can now be sent and their outputs received --- .../src/commands/connect/mod.rs | 2 +- sparse-05/sparse-05-server/src/connection.rs | 11 +- .../src/connection/command.rs | 157 ++++++++++++------ 3 files changed, 115 insertions(+), 55 deletions(-) diff --git a/sparse-05/sparse-05-client/src/commands/connect/mod.rs b/sparse-05/sparse-05-client/src/commands/connect/mod.rs index 3d6139e..70991b3 100644 --- a/sparse-05/sparse-05-client/src/commands/connect/mod.rs +++ b/sparse-05/sparse-05-client/src/commands/connect/mod.rs @@ -67,7 +67,7 @@ impl Connection { bail!("received packet from wrong computer"); } - dbg!(self.decrypt_and_verify(&buffer[..read])) + self.decrypt_and_verify(&buffer[..read]) } } diff --git a/sparse-05/sparse-05-server/src/connection.rs b/sparse-05/sparse-05-server/src/connection.rs index 23d0d36..a30d802 100644 --- a/sparse-05/sparse-05-server/src/connection.rs +++ b/sparse-05/sparse-05-server/src/connection.rs @@ -25,12 +25,13 @@ impl ConnectionHandle { } } +#[derive(Clone)] struct ConnectionInformation { remote_sign_pubkey: ed25519_dalek::PublicKey, remote_enc_pubkey: ecies_ed25519::PublicKey, - local_sign_keypair: Keypair, + local_sign_keypair: Arc, local_enc_pubkey: ecies_ed25519::PublicKey, - local_enc_privkey: ecies_ed25519::SecretKey, + local_enc_privkey: Arc, srcmac: [u8; 6], dstmac: [u8; 6], srcip: Ipv4Addr, @@ -131,8 +132,8 @@ pub fn spawn_connection_handler( ConnectionInformation { remote_sign_pubkey: (*sign_pubkey).clone(), remote_enc_pubkey: (*enc_pubkey).clone(), - local_sign_keypair, - local_enc_privkey, + local_sign_keypair: Arc::new(local_sign_keypair), + local_enc_privkey: Arc::new(local_enc_privkey), local_enc_pubkey, srcmac: *packet.source_address(), dstmac: *packet.destination_address(), @@ -250,7 +251,7 @@ where match data { Command::RunCommand(comm) => { - let handler = match command::spawn_command(&s, comm, &conninfo) { + let handler = match command::spawn_command(&s, comm, conninfo.clone()) { Ok(handler) => handler, Err(e) => { eprintln!("error spawning command: {e:?}"); diff --git a/sparse-05/sparse-05-server/src/connection/command.rs b/sparse-05/sparse-05-server/src/connection/command.rs index 3ac7ddf..ac0c19e 100644 --- a/sparse-05/sparse-05-server/src/connection/command.rs +++ b/sparse-05/sparse-05-server/src/connection/command.rs @@ -2,10 +2,11 @@ use std::{ io::{Read, Write}, process::{Command, Stdio}, sync::{ - atomic::AtomicU64, + atomic::{AtomicBool, AtomicU64, Ordering}, mpsc::{channel, Sender}, + Arc, }, - thread::Scope, + thread::{scope, Scope}, }; use sparse_05_common::messages::Response; @@ -22,7 +23,7 @@ pub(super) struct CommandHandler { pub(super) fn spawn_command<'a, 'b: 'a>( s: &'a Scope<'a, 'b>, command: String, - conninfo: &'a ConnectionInformation, + conninfo: ConnectionInformation, ) -> anyhow::Result { let (data_sender, data_receiver) = channel(); @@ -41,62 +42,120 @@ pub(super) fn spawn_command<'a, 'b: 'a>( let resp = Response::AckRunCommand(id); conninfo.send(conninfo.encrypt_and_sign_resp(resp)?)?; + let data_sender_2 = data_sender.clone(); s.spawn(move || -> anyhow::Result<()> { - if let Some(mut stdout) = command.stdout.take() { - s.spawn(move || { - let mut stdout_buffer = [0u8; 1024]; + let should_die = AtomicBool::new(false); + scope(|s2| -> anyhow::Result<()> { + let stdout_thread = if let Some(mut stdout) = command.stdout.take() { + Some(s2.spawn({ + let conninfo = conninfo.clone(); + let should_die = &should_die; + move || { + let mut stdout_buffer = [0u8; 1024]; - loop { - let Ok(count) = stdout.read(&mut stdout_buffer) else { - continue; - }; + loop { + let Ok(count) = stdout.read(&mut stdout_buffer) else { + continue; + }; - let resp = Response::SendStdout(stdout_buffer[..count].to_vec(), id); - let Ok(packet) = conninfo.encrypt_and_sign_resp(resp) else { - continue; - }; - _ = conninfo.send(packet); - } - }); - } + if should_die.load(Ordering::Relaxed) { + break; + } - if let Some(mut stderr) = command.stderr.take() { - s.spawn(move || { - let mut stderr_buffer = [0u8; 1024]; + if count == 0 { + continue; + } - loop { - let Ok(count) = stderr.read(&mut stderr_buffer) else { - continue; - }; + let resp = Response::SendStdout(stdout_buffer[..count].to_vec(), id); + let Ok(packet) = conninfo.encrypt_and_sign_resp(resp) else { + continue; + }; + _ = conninfo.send(packet); + } + } + })) + } else { + None + }; - let resp = Response::SendStderr(stderr_buffer[..count].to_vec(), id); - let Ok(packet) = conninfo.encrypt_and_sign_resp(resp) else { - continue; - }; - _ = conninfo.send(packet); - } - }); - } + let stderr_thread = if let Some(mut stderr) = command.stderr.take() { + Some(s2.spawn({ + let conninfo = conninfo.clone(); + let should_die = &should_die; + move || { + let mut stderr_buffer = [0u8; 1024]; - if let Some(mut stdin) = command.stdin.take() { - s.spawn(move || loop { - let Ok(input): Result, _> = data_receiver.recv() else { - continue; - }; + loop { + let Ok(count) = stderr.read(&mut stderr_buffer) else { + continue; + }; - _ = stdin.write(&input); - }); - } + if should_die.load(Ordering::Relaxed) { + break; + } - let code = match command.wait() { - Ok(status) => status.code().unwrap_or(-1), - Err(_) => -1, - }; - let resp = Response::CommandDone(id, code); - let resp = conninfo.encrypt_and_sign_resp(resp)?; - conninfo.send(resp)?; - Ok(()) + if count == 0 { + continue; + } + + let resp = Response::SendStderr(stderr_buffer[..count].to_vec(), id); + let Ok(packet) = conninfo.encrypt_and_sign_resp(resp) else { + continue; + }; + _ = conninfo.send(packet); + } + } + })) + } else { + None + }; + + let stdin_thread = if let Some(mut stdin) = command.stdin.take() { + Some(s2.spawn({ + let should_die = &should_die; + move || loop { + let Ok(input): Result, _> = + data_receiver.recv_timeout(std::time::Duration::from_millis(250)) + else { + if should_die.load(Ordering::Relaxed) { + break; + } + + continue; + }; + + if should_die.load(Ordering::Relaxed) { + break; + } + + _ = stdin.write(&input); + } + })) + } else { + None + }; + + let code = match command.wait() { + Ok(status) => status.code().unwrap_or(-1), + Err(_) => -1, + }; + drop(data_sender_2); + should_die.store(true, Ordering::Relaxed); + if let Some(thread) = stdout_thread { + _ = thread.join(); + } + if let Some(thread) = stderr_thread { + _ = thread.join(); + } + let resp = Response::CommandDone(id, code); + let resp = conninfo.encrypt_and_sign_resp(resp)?; + conninfo.send(resp)?; + if let Some(thread) = stdin_thread { + _ = thread.join(); + } + Ok(()) + }) }); Ok(CommandHandler { id, data_sender })