Feature/ctrl channel

.gitignore

@@ -22,3 +22,4 @@ dist/
 report.json
 pkg/
 */scripts/*
+.VSCodeCounter/

cli/Cargo.toml

@@ -18,7 +18,7 @@ ron = "0.8.1"
 eyre = "0.6"
 http-body-util = { version = "0.1" }
 hyper = { version = "1.3.1", features = ["client", "http1"] }
-hyper-util = { version = "0.1", features = ["full"] }
+hyper-util = { version = "0.1", features = ["client", "http1", "tokio"] }
 hyperparameter = { version = "0.5.11", default-features = false }
 libloading = "0.8.3"
 log = "0.4"
@@ -35,7 +35,9 @@ strum = "0.26.3"
 tui-tree-widget = "0.21.0"
 serde_json = "1.0.117"
 once_cell = "1.19.0"
-nu-ansi-term = "0.50.0"
+clap_complete = {version ="4.5.8", features = ["unstable-dynamic"]}
+shlex = "1.3.0"
+rustyline = {version="14.0.0", default-features = false}
 
 [dependencies.tokio]
 version = "1.35.1"

cli/src/cli/commands.rs

@@ -1,15 +1,36 @@
-use clap::Subcommand;
+use clap::{Parser, Subcommand};
 
-#[derive(Subcommand)]
+#[derive(Subcommand, Debug)]
 pub enum Commands {
     #[command(visible_aliases = ["in", "i"])]
     Inject(super::inject::InjectCommand),
     #[command(visible_aliases = ["dbg", "d"])]
     Debug(super::debug::DebugCommand),
     #[command(visible_aliases = ["perf", "p"])]
     Performance(super::performance::PerfCommand),
-    #[command(visible_aliases = ["m"])]
+
+    /// Console visualizer
+    #[command(visible_aliases = ["pnl", "console"])]
+    Panel,
+    #[command()]
+    Repl,
+    #[command(hide=true, visible_aliases = ["m"])]
+    Misc(super::misc::MiscCommand),
+}
+
+#[derive(Parser, Debug)]
+#[command(name = "")]
+pub enum ReplCommands {
+    #[command(visible_aliases = ["in", "i"])]
+    Inject(super::inject::InjectCommand),
+    #[command(visible_aliases = ["dbg", "d"])]
+    Debug(super::debug::DebugCommand),
+    #[command(visible_aliases = ["perf", "p"])]
+    Performance(super::performance::PerfCommand),
+    #[command(visible_aliases = ["pnl", "console"])]
+    Panel,
+    #[command()]
+    Repl,
+    #[command(hide=true, visible_aliases = ["m"])]
     Misc(super::misc::MiscCommand),
-    #[command(visible_aliases = ["con"])]
-    Console,
 }

cli/src/cli/ctrl.rs

@@ -2,25 +2,123 @@ use anyhow::Result;
 
 use http_body_util::{BodyExt, Full};
 use hyper_util::rt::TokioIo;
+use hyperparameter::*;
+use nix::{sys::signal, unistd::Pid};
 
-pub async fn request(pid: i32, url: &str, body: Option<String>) -> Result<String> {
+use crate::inject::{Injector, Process};
+
+#[derive(Clone)]
+pub enum CtrlChannel {
+    Ptrace { pid: i32 },
+    Local { pid: i32 },
+    Remote { addr: String },
+}
+
+impl TryFrom<&str> for CtrlChannel {
+    type Error = anyhow::Error;
+
+    fn try_from(value: &str) -> Result<Self> {
+        if let [_, _] = value.split(':').collect::<Vec<_>>()[..] {
+            return Ok(Self::Remote { addr: value.into() });
+        }
+
+        let callback = |pid| -> Result<CtrlChannel> {
+            with_params! {
+                get use_ptrace = probing.cli.ptrace or false;
+
+                Ok(if use_ptrace {Self::Ptrace { pid }} else {Self::Local { pid }})
+            }
+        };
+
+        if let Ok(pid) = value.parse::<i32>() {
+            return callback(pid);
+        }
+
+        let pid = Process::by_cmdline(value).map_err(|err| {
+            anyhow::anyhow!("failed to find process with cmdline pattern {value}: {err}")
+        })?;
+        if let Some(pid) = pid {
+            callback(pid)
+        } else {
+            Err(anyhow::anyhow!("either `pid` or `name` must be specified"))
+        }
+    }
+}
+
+impl TryFrom<String> for CtrlChannel {
+    type Error = anyhow::Error;
+
+    fn try_from(value: String) -> Result<Self, Self::Error> {
+        value.as_str().try_into()
+    }
+}
+
+impl Into<String> for CtrlChannel {
+    fn into(self) -> String {
+        match self {
+            CtrlChannel::Ptrace { pid } | CtrlChannel::Local { pid } => format! {"{pid}"},
+            CtrlChannel::Remote { addr } => addr,
+        }
+    }
+}
+
+impl CtrlChannel {
+    pub fn send_ctrl(&self, cmd: String) -> Result<()> {
+        match self {
+            CtrlChannel::Ptrace { pid } => send_ctrl_via_ptrace(cmd, *pid),
+            ctrl => {
+                let cmd = if cmd.starts_with('[') {
+                    cmd
+                } else {
+                    format!("[{}]", cmd)
+                };
+                tokio::runtime::Builder::new_current_thread()
+                    .enable_all()
+                    .build()
+                    .unwrap()
+                    .block_on(request(ctrl.clone(), "/ctrl", cmd.into()))?;
+
+                Ok(())
+            }
+        }
+    }
+}
+
+pub async fn request(ctrl: CtrlChannel, url: &str, body: Option<String>) -> Result<String> {
     use hyper::body::Bytes;
     use hyper::client::conn;
     use hyper::Request;
 
-    let prefix = "/tmp/probing".to_string();
-    let path = format!("{}/{}", prefix, pid);
-    let path = std::path::Path::new(&path);
-    if !path.exists() {
-        anyhow::bail!("server not found: {}", path.display());
-    }
-    let stream = tokio::net::UnixStream::connect(path).await?;
-    let io = TokioIo::new(stream);
+    let mut sender = match ctrl {
+        CtrlChannel::Ptrace { pid } | CtrlChannel::Local { pid } => {
+            eprintln!("sending ctrl commands via unix socket...");
+            let prefix = "/tmp/probing".to_string();
+            let path = format!("{}/{}", prefix, pid);
+            let path = std::path::Path::new(&path);
+            if !path.exists() {
+                anyhow::bail!("server not found: {}", path.display());
+            }
+            let stream = tokio::net::UnixStream::connect(path).await?;
+            let io = TokioIo::new(stream);
+
+            let (sender, connection) = conn::http1::handshake(io).await?;
+            tokio::spawn(async move {
+                connection.await.unwrap();
+            });
+            sender
+        }
+        CtrlChannel::Remote { addr } => {
+            eprintln!("sending ctrl commands via tcp socket...");
+            let stream = tokio::net::TcpStream::connect(addr).await?;
+            let io = TokioIo::new(stream);
 
-    let (mut sender, connection) = conn::http1::handshake(io).await?;
-    tokio::spawn(async move {
-        connection.await.unwrap();
-    });
+            let (sender, connection) = conn::http1::handshake(io).await?;
+            tokio::spawn(async move {
+                connection.await.unwrap();
+            });
+            sender
+        }
+    };
     let request = if let Some(body) = body {
         Request::builder()
             .method("POST")
@@ -53,3 +151,14 @@ pub async fn request(pid: i32, url: &str, body: Option<String>) -> Result<String
         anyhow::bail!("Error {}: {}", res.status(), body)
     }
 }
+
+fn send_ctrl_via_ptrace(argstr: String, pid: i32) -> Result<()> {
+    eprintln!("sending ctrl commands via ptrace...");
+    let process = Process::get(pid as u32).unwrap();
+    Injector::attach(process)
+        .unwrap()
+        .setenv(Some("PROBING_ARGS"), Some(argstr.as_str()))
+        .map_err(|e| anyhow::anyhow!(e))?;
+    signal::kill(Pid::from_raw(pid), signal::Signal::SIGUSR1)?;
+    Ok(())
+}

cli/src/cli/debug.rs

@@ -1,15 +1,12 @@
-use anyhow::Context;
-use anyhow::Ok;
 use anyhow::Result;
 
 use clap::Args;
-use nix::{sys::signal, unistd::Pid};
 use probing_common::cli::ProbingCommand;
 
-use super::send_ctrl;
+use super::ctrl::CtrlChannel;
 
 /// Debug and Inspection Tool
-#[derive(Args)]
+#[derive(Args, Debug)]
 pub struct DebugCommand {
     /// Dump the calling stack of the target process
     #[arg(short, long, conflicts_with_all=["dap", "pause"])]
@@ -29,24 +26,21 @@ pub struct DebugCommand {
 }
 
 impl DebugCommand {
-    pub fn run(&self, pid: i32) -> Result<()> {
-        if self.dump {
-            signal::kill(Pid::from_raw(pid), signal::Signal::SIGUSR2)
-                .with_context(|| format!("error sending signal to pid {pid}"))
+    pub fn run(&self, ctrl: CtrlChannel) -> Result<()> {
+        let cmd = if self.dump {
+            ProbingCommand::Dump
         } else if self.pause {
-            let cmd = ProbingCommand::Pause {
+            ProbingCommand::Pause {
                 address: self.address.clone(),
-            };
-            let cmd = ron::to_string(&cmd)?;
-            send_ctrl(cmd, pid)
+            }
         } else if self.dap {
-            let cmd = ProbingCommand::Dap {
+            ProbingCommand::Dap {
                 address: self.address.clone(),
-            };
-            let cmd = ron::to_string(&cmd)?;
-            send_ctrl(cmd, pid)
+            }
         } else {
-            Ok(())
-        }
+            ProbingCommand::Nil
+        };
+        let cmd = ron::to_string(&cmd)?;
+        ctrl.send_ctrl(cmd)
     }
 }