use crate::{logs::LogEntry, remote::RemoteMachine, Error, Result};
use std::path::{Path, PathBuf};
use std::process::{Child, Command, Stdio};
pub(crate) trait PeerProcess {
fn is_running(&self) -> bool;
fn kill(&mut self) -> Result<()>;
fn log_path(&self) -> PathBuf;
fn read_logs(&self) -> Result<Vec<LogEntry>>;
}
pub(crate) struct LocalProcess {
pub(crate) process: Option<Child>,
pub(crate) data_dir: PathBuf,
}
impl PeerProcess for LocalProcess {
fn is_running(&self) -> bool {
if let Some(process) = &self.process {
let pid = process.id();
let mut system = sysinfo::System::new();
system.refresh_processes(sysinfo::ProcessesToUpdate::All, true);
system.process(sysinfo::Pid::from_u32(pid as u32)).is_some()
} else {
false
}
}
fn kill(&mut self) -> Result<()> {
if let Some(mut process) = self.process.take() {
process.kill()?;
process.wait()?;
}
Ok(())
}
fn log_path(&self) -> PathBuf {
self.data_dir.join("peer.log")
}
fn read_logs(&self) -> Result<Vec<LogEntry>> {
crate::logs::read_log_file(&self.log_path())
}
}
impl Drop for LocalProcess {
fn drop(&mut self) {
if let Some(mut process) = self.process.take() {
let _ = process.kill();
let _ = process.wait();
}
}
}
pub(crate) struct SshProcess {
pub(crate) remote: RemoteMachine,
pub(crate) remote_pid: Option<u32>,
pub(crate) remote_data_dir: PathBuf,
pub(crate) local_cache_dir: PathBuf, }
impl PeerProcess for SshProcess {
fn is_running(&self) -> bool {
if let Some(pid) = self.remote_pid {
let cmd = format!("kill -0 {}", pid);
self.remote.exec(&cmd).is_ok()
} else {
false
}
}
fn kill(&mut self) -> Result<()> {
if let Some(pid) = self.remote_pid.take() {
let cmd = format!("kill {}", pid);
let _ = self.remote.exec(&cmd);
std::thread::sleep(std::time::Duration::from_millis(500));
if self.remote.exec(&format!("kill -0 {}", pid)).is_ok() {
let _ = self.remote.exec(&format!("kill -9 {}", pid));
}
}
Ok(())
}
fn log_path(&self) -> PathBuf {
self.local_cache_dir.join("peer.log")
}
fn read_logs(&self) -> Result<Vec<LogEntry>> {
let remote_log = self.remote_data_dir.join("peer.log");
let local_log = self.local_cache_dir.join("peer.log");
if let Some(parent) = local_log.parent() {
std::fs::create_dir_all(parent)?;
}
self.remote.scp_download(
remote_log
.to_str()
.ok_or_else(|| Error::PeerStartupFailed("Invalid remote log path".to_string()))?,
&local_log,
)?;
crate::logs::read_log_file(&local_log)
}
}
impl Drop for SshProcess {
fn drop(&mut self) {
let _ = self.kill();
}
}
pub(crate) fn spawn_local_peer(
binary_path: &Path,
args: &[String],
data_dir: &Path,
env_vars: &[(String, String)],
) -> Result<LocalProcess> {
let log_path = data_dir.join("peer.log");
let log_file = std::fs::File::create(&log_path)?;
let mut cmd = Command::new(binary_path);
let rust_log = std::env::var("RUST_LOG").unwrap_or_else(|_| "info".to_string());
cmd.args(args)
.env("RUST_LOG", rust_log)
.env("RUST_BACKTRACE", "1")
.stdout(Stdio::from(log_file.try_clone()?))
.stderr(Stdio::from(log_file));
for (key, value) in env_vars {
cmd.env(key, value);
}
let process = cmd
.spawn()
.map_err(|e| Error::PeerStartupFailed(format!("Failed to spawn local process: {}", e)))?;
Ok(LocalProcess {
process: Some(process),
data_dir: data_dir.to_path_buf(),
})
}
pub(crate) async fn spawn_remote_peer(
local_binary_path: &Path,
args: &[String],
remote: &RemoteMachine,
remote_data_dir: &Path,
local_cache_dir: &Path,
env_vars: &[(String, String)],
) -> Result<SshProcess> {
let remote_binary_path = if let Some(binary) = &remote.freenet_binary {
binary.clone()
} else {
let remote_bin_dir = remote.remote_work_dir().join("bin");
let binary_name = local_binary_path
.file_name()
.ok_or_else(|| Error::PeerStartupFailed("Invalid binary path".to_string()))?;
let remote_binary = remote_bin_dir.join(binary_name);
let mkdir_cmd = format!("mkdir -p {}", remote_bin_dir.display());
remote.exec(&mkdir_cmd)?;
tracing::debug!(
"Deploying binary to {}:{}",
remote.host,
remote_binary.display()
);
remote.scp_upload(
local_binary_path,
remote_binary.to_str().ok_or_else(|| {
Error::PeerStartupFailed("Invalid remote binary path".to_string())
})?,
)?;
remote_binary
};
let mkdir_cmd = format!("mkdir -p {}", remote_data_dir.display());
remote.exec(&mkdir_cmd)?;
let mut cmd_parts = vec![remote_binary_path.to_string_lossy().to_string()];
cmd_parts.extend(args.iter().map(|s| shell_escape(s)));
let mut env_prefix = String::new();
env_prefix.push_str("RUST_LOG=info RUST_BACKTRACE=1 ");
for (key, value) in env_vars {
env_prefix.push_str(&format!("{}={} ", shell_escape(key), shell_escape(value)));
}
let log_path = remote_data_dir.join("peer.log");
let full_command = format!(
"nohup {} {} > {} 2>&1 & echo $!",
env_prefix,
cmd_parts.join(" "),
log_path.display()
);
tracing::debug!("Executing remote command: {}", full_command);
let output = remote.exec(&full_command)?;
let pid: u32 = output
.trim()
.parse()
.map_err(|e| Error::PeerStartupFailed(format!("Failed to parse remote PID: {}", e)))?;
tracing::debug!("Remote peer started with PID {}", pid);
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
Ok(SshProcess {
remote: remote.clone(),
remote_pid: Some(pid),
remote_data_dir: remote_data_dir.to_path_buf(),
local_cache_dir: local_cache_dir.to_path_buf(),
})
}
fn shell_escape(s: &str) -> String {
if s.chars()
.all(|c| c.is_alphanumeric() || c == '_' || c == '-' || c == '/' || c == '.')
{
s.to_string()
} else {
format!("'{}'", s.replace('\'', r"'\''"))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_shell_escape() {
assert_eq!(shell_escape("simple"), "simple");
assert_eq!(shell_escape("/path/to/file"), "/path/to/file");
assert_eq!(shell_escape("hello world"), "'hello world'");
assert_eq!(shell_escape("it's"), "'it'\\''s'");
}
}