330 lines
11 KiB
Rust
330 lines
11 KiB
Rust
//! 进程探测与流式执行辅助（内部）
|
|
//!
|
|
//! 所有命令经 `agentdock-exec::validate_argv` 校验后执行，禁止 shell 拼接。
|
|
//! 探测类命令（`--version`）隐藏控制台窗口（Windows CREATE_NO_WINDOW）。
|
|
//! 流式执行（`run_streaming` / `run_streaming_cancellable`）用两条读取线程
|
|
//! **并发**消费 stdout 与 stderr，再经 mpsc 通道按到达顺序回调——保证实时滚动，
|
|
//! 避免「先读 stdout 再读 stderr」造成的 stderr 延迟到进程结束时才出现的旧行为。
|
|
|
|
use std::io::BufRead;
|
|
use std::io::Write;
|
|
use std::path::{Path, PathBuf};
|
|
use std::process::{Child, Command, Stdio};
|
|
use std::sync::atomic::{AtomicBool, Ordering};
|
|
use std::sync::{Arc, Mutex, mpsc};
|
|
|
|
/// 按平台分隔符拆分 PATH。
|
|
fn path_entries() -> Vec<String> {
|
|
std::env::var("PATH")
|
|
.unwrap_or_default()
|
|
.split(if cfg!(windows) { ';' } else { ':' })
|
|
.filter(|s| !s.is_empty())
|
|
.map(|s| s.to_string())
|
|
.collect()
|
|
}
|
|
|
|
/// 在 PATH 中查找可执行文件（返回全部命中，用于版本冲突检测）。
|
|
/// Windows 优先 .exe/.cmd/.bat，最后回退无扩展名。
|
|
pub fn which_all(program: &str) -> Vec<PathBuf> {
|
|
let direct = Path::new(program);
|
|
if direct.is_absolute() && direct.is_file() {
|
|
return vec![direct.to_path_buf()];
|
|
}
|
|
let exts: &[&str] = if cfg!(windows) { &[".exe", ".cmd", ".bat", ""] } else { &[""] };
|
|
let mut found = Vec::new();
|
|
for dir in path_entries() {
|
|
for ext in exts {
|
|
let cand = Path::new(&dir).join(format!("{program}{ext}"));
|
|
if cand.is_file() && !found.contains(&cand) {
|
|
found.push(cand);
|
|
}
|
|
}
|
|
}
|
|
found
|
|
}
|
|
|
|
/// 在 PATH 中解析可执行文件为可运行路径（找不到时回退原 prog，让 spawn 报错）。
|
|
pub fn resolve_exe(program: &str) -> String {
|
|
which_all(program)
|
|
.first()
|
|
.map(|p| p.to_string_lossy().to_string())
|
|
.unwrap_or_else(|| program.to_string())
|
|
}
|
|
|
|
/// 运行只读探测命令并返回 (exit_ok, stdout+stderr 合并文本)。
|
|
pub fn run_capture(exe: &Path, args: &[String]) -> std::io::Result<(bool, String)> {
|
|
let mut cmd = Command::new(exe);
|
|
cmd.args(args);
|
|
#[cfg(windows)]
|
|
{
|
|
use std::os::windows::process::CommandExt;
|
|
cmd.creation_flags(0x0800_0000); // CREATE_NO_WINDOW
|
|
}
|
|
let out = cmd.output()?;
|
|
let mut text = String::from_utf8_lossy(&out.stdout).to_string();
|
|
if text.trim().is_empty() {
|
|
text = String::from_utf8_lossy(&out.stderr).to_string();
|
|
}
|
|
Ok((out.status.success(), text))
|
|
}
|
|
|
|
/// 运行命令并把一段文本写入其 stdin（API Key 经 stdin 注入官方登录命令用），
|
|
/// 返回 (exit_ok, stdout+stderr 合并文本)。输入内容绝不明文落日志。
|
|
pub fn run_with_stdin(exe: &Path, args: &[String], stdin_text: &str) -> std::io::Result<(bool, String)> {
|
|
let mut cmd = Command::new(exe);
|
|
cmd.args(args);
|
|
cmd.stdin(Stdio::piped());
|
|
cmd.stdout(Stdio::piped());
|
|
cmd.stderr(Stdio::piped());
|
|
#[cfg(windows)]
|
|
{
|
|
use std::os::windows::process::CommandExt;
|
|
cmd.creation_flags(0x0800_0000); // CREATE_NO_WINDOW
|
|
}
|
|
let mut child = cmd.spawn()?;
|
|
if let Some(mut stdin) = child.stdin.take() {
|
|
let _ = stdin.write_all(stdin_text.as_bytes());
|
|
// 关闭 stdin，通知子进程输入结束
|
|
}
|
|
let out = child.wait_with_output()?;
|
|
let mut text = String::from_utf8_lossy(&out.stdout).to_string();
|
|
if text.trim().is_empty() {
|
|
text = String::from_utf8_lossy(&out.stderr).to_string();
|
|
}
|
|
Ok((out.status.success(), text))
|
|
}
|
|
|
|
/// 可取消的进程句柄：`cancel` 置位后流式读取循环会终止子进程；
|
|
/// 授权等交互式流程需要取消时调用 `request_cancel`。
|
|
#[derive(Default)]
|
|
pub struct RunningProcess {
|
|
cancel: Arc<AtomicBool>,
|
|
child: Arc<Mutex<Option<Child>>>,
|
|
}
|
|
|
|
impl RunningProcess {
|
|
pub fn new() -> Self {
|
|
RunningProcess::default()
|
|
}
|
|
|
|
/// 请求取消（安全幂等）。
|
|
pub fn request_cancel(&self) {
|
|
self.cancel.store(true, Ordering::SeqCst);
|
|
}
|
|
|
|
/// 取取消标志的共享句柄（供其它线程定时/条件取消）。
|
|
pub fn cancel_flag(&self) -> Arc<AtomicBool> {
|
|
self.cancel.clone()
|
|
}
|
|
|
|
pub fn is_cancelled(&self) -> bool {
|
|
self.cancel.load(Ordering::SeqCst)
|
|
}
|
|
|
|
fn kill(&self) {
|
|
if let Some(mut child) = self.child.lock().unwrap().take() {
|
|
kill_child_tree(&mut child);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// 流式执行命令，逐行回调（stdout 与 stderr 并发、按到达顺序实时输出）。
|
|
/// 返回是否成功。
|
|
pub fn run_streaming<F>(exe: &str, args: &[String], on_line: F) -> std::io::Result<bool>
|
|
where
|
|
F: FnMut(bool, &str),
|
|
{
|
|
let rp = RunningProcess::new();
|
|
run_streaming_cancellable(&rp, exe, args, on_line)
|
|
}
|
|
|
|
/// 可取消的流式执行：stdout/stderr 并发读取，主线程按行回调；
|
|
/// 取消置位时终止子进程并结束。
|
|
pub fn run_streaming_cancellable<F>(
|
|
rp: &RunningProcess,
|
|
exe: &str,
|
|
args: &[String],
|
|
mut on_line: F,
|
|
) -> std::io::Result<bool>
|
|
where
|
|
F: FnMut(bool, &str),
|
|
{
|
|
let mut cmd = Command::new(exe);
|
|
cmd.args(args);
|
|
cmd.stdout(Stdio::piped());
|
|
cmd.stderr(Stdio::piped());
|
|
#[cfg(windows)]
|
|
{
|
|
use std::os::windows::process::CommandExt;
|
|
cmd.creation_flags(0x0800_0000);
|
|
}
|
|
let mut child = cmd.spawn()?;
|
|
|
|
let stdout = child.stdout.take();
|
|
let stderr = child.stderr.take();
|
|
*rp.child.lock().unwrap() = Some(child);
|
|
|
|
let (tx, rx) = mpsc::channel::<(bool, String)>();
|
|
|
|
// stdout 读取线程
|
|
if let Some(out) = stdout {
|
|
let tx = tx.clone();
|
|
std::thread::spawn(move || {
|
|
let reader = std::io::BufReader::new(out);
|
|
for line in reader.lines().map_while(Result::ok) {
|
|
if tx.send((false, line)).is_err() {
|
|
break;
|
|
}
|
|
}
|
|
});
|
|
}
|
|
// stderr 读取线程
|
|
if let Some(err) = stderr {
|
|
let tx = tx.clone();
|
|
std::thread::spawn(move || {
|
|
let reader = std::io::BufReader::new(err);
|
|
for line in reader.lines().map_while(Result::ok) {
|
|
if tx.send((true, line)).is_err() {
|
|
break;
|
|
}
|
|
}
|
|
});
|
|
}
|
|
drop(tx); // 主线程持有的发送端关闭，两线程结束后通道自然关闭
|
|
|
|
// 主线程按到达顺序实时回调（recv_timeout 兜底：静默期也能响应取消）
|
|
loop {
|
|
match rx.recv_timeout(std::time::Duration::from_millis(200)) {
|
|
Ok((is_err, line)) => {
|
|
on_line(is_err, &line);
|
|
if rp.is_cancelled() {
|
|
rp.kill();
|
|
break;
|
|
}
|
|
}
|
|
Err(mpsc::RecvTimeoutError::Timeout) => {
|
|
if rp.is_cancelled() {
|
|
rp.kill();
|
|
break;
|
|
}
|
|
}
|
|
Err(mpsc::RecvTimeoutError::Disconnected) => break,
|
|
}
|
|
}
|
|
|
|
// 等待子进程退出并取状态
|
|
let status = rp
|
|
.child
|
|
.lock()
|
|
.unwrap()
|
|
.take()
|
|
.and_then(|mut c| c.wait().ok())
|
|
.map(|s| s.success())
|
|
.unwrap_or(false);
|
|
Ok(status)
|
|
}
|
|
|
|
/// 终止子进程（Windows 用 taskkill /T 杀掉整棵进程树，避免 .cmd 包装的 node 等子进程变孤儿）。
|
|
fn kill_child_tree(child: &mut Child) {
|
|
#[cfg(windows)]
|
|
{
|
|
use std::os::windows::process::CommandExt;
|
|
let pid = child.id();
|
|
let _ = Command::new("taskkill")
|
|
.args(["/PID", &pid.to_string(), "/T", "/F"])
|
|
.creation_flags(0x0800_0000)
|
|
.status();
|
|
let _ = child.kill();
|
|
let _ = child.wait();
|
|
}
|
|
#[cfg(not(windows))]
|
|
{
|
|
let _ = child.kill();
|
|
let _ = child.wait();
|
|
}
|
|
}
|
|
|
|
/// 在**独立控制台窗口**中运行命令（一次性本机终端，PRD FR-06 设计），
|
|
/// 结束后窗口自动关闭。仅回传退出成功与否，不回传输出。
|
|
/// 用于必须在终端里交互的登录流程（local_tui 等）。
|
|
#[cfg(windows)]
|
|
pub fn run_terminal(exe: &str, args: &[String]) -> std::io::Result<bool> {
|
|
use std::os::windows::process::CommandExt;
|
|
let mut cmd = Command::new(exe);
|
|
cmd.args(args);
|
|
cmd.creation_flags(0x0000_0010); // CREATE_NEW_CONSOLE
|
|
let status = cmd.status()?;
|
|
Ok(status.success())
|
|
}
|
|
|
|
/// 非 Windows 平台：终端窗口回退为普通前台执行。
|
|
#[cfg(not(windows))]
|
|
pub fn run_terminal(exe: &str, args: &[String]) -> std::io::Result<bool> {
|
|
let status = Command::new(exe).args(args).status()?;
|
|
Ok(status.success())
|
|
}
|
|
|
|
/// 打开本机默认程序/浏览器（文件用默认处理器打开、URL 用默认浏览器打开）。
|
|
/// 仅接受已通过来源白名单校验的目标，不做 shell 拼接。
|
|
#[cfg(windows)]
|
|
pub fn open_with_shell(target: &str) -> std::io::Result<()> {
|
|
// explorer.exe 同时能打开文件（默认处理器）与 URL（默认浏览器）
|
|
let status = Command::new("explorer.exe").arg(target).status()?;
|
|
if status.success() {
|
|
Ok(())
|
|
} else {
|
|
Err(std::io::Error::new(
|
|
std::io::ErrorKind::Other,
|
|
"explorer 打开目标失败",
|
|
))
|
|
}
|
|
}
|
|
|
|
#[cfg(not(windows))]
|
|
pub fn open_with_shell(target: &str) -> std::io::Result<()> {
|
|
let status = Command::new("xdg-open").arg(target).status()?;
|
|
if status.success() {
|
|
Ok(())
|
|
} else {
|
|
Err(std::io::Error::new(std::io::ErrorKind::Other, "xdg-open 打开目标失败"))
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn run_streaming_reports_stdout_in_order() {
|
|
// 用 echo 输出多行，验证 stdout 被逐行实时捕获
|
|
#[cfg(windows)]
|
|
let (exe, args) = (
|
|
"cmd.exe".to_string(),
|
|
vec!["/C".to_string(), "echo line1 & echo line2 & echo line3".to_string()],
|
|
);
|
|
#[cfg(not(windows))]
|
|
let (exe, args) = ("sh".to_string(), vec!["-c".to_string(), "echo line1; echo line2; echo line3".to_string()]);
|
|
|
|
let mut lines = Vec::new();
|
|
let ok = run_streaming(&exe, &args, |is_err, l| {
|
|
if !is_err {
|
|
lines.push(l.to_string());
|
|
}
|
|
})
|
|
.unwrap();
|
|
assert!(ok);
|
|
assert!(lines.iter().any(|l| l.contains("line1")));
|
|
assert!(lines.iter().any(|l| l.contains("line3")));
|
|
}
|
|
|
|
#[test]
|
|
fn running_process_cancel_is_idempotent() {
|
|
let rp = RunningProcess::new();
|
|
assert!(!rp.is_cancelled());
|
|
rp.request_cancel();
|
|
assert!(rp.is_cancelled());
|
|
rp.request_cancel();
|
|
assert!(rp.is_cancelled());
|
|
}
|
|
}
|