CNWei
e52ba67d7e
- 支持在 Lua 中配置数组格式的环境变量（如 PATH），由 Rust 按 OS 自动使用 `;` 或 `:` 拼接 - 保持普通字符串格式环境变量的直接覆盖逻辑 - 新增 `ShimConfig::to_command()` 用于快速创建 Command 对象 - 新增 error.rs 明确错误类型 - 新增 build.rs 优化编译
284 lines
10 KiB
Rust
284 lines
10 KiB
Rust
use crate::error::ShimError;
|
||
use mlua::{Error as LuaError, FromLua, Lua, StdLib, Table, Value};
|
||
use std::process::Command;
|
||
use std::{
|
||
collections::HashMap,
|
||
env, fs,
|
||
io::{Error, ErrorKind},
|
||
path::{Path, PathBuf},
|
||
};
|
||
|
||
/// 将 Path 转换为适合 Lua 使用的安全字符串路径
|
||
fn normalize_path_for_lua(path: &Path) -> String {
|
||
let path_str = path.to_string_lossy();
|
||
|
||
// 1. 剥离 Windows UNC 规范路径前缀 (\\?\)
|
||
let clean_str = path_str.strip_prefix(r"\\?\").unwrap_or(&path_str);
|
||
|
||
// 2. 将反斜杠转换成正斜杠（在非 UNC 路径下，Windows 和 Lua 均完美支持 /）
|
||
// 这样既避免了 Lua 字符串转义隐患，又不会破坏 Windows 路径
|
||
clean_str.replace('\\', "/")
|
||
}
|
||
|
||
pub struct ShimEnv {
|
||
pub bin_dir: PathBuf,
|
||
pub root_dir: PathBuf,
|
||
pub tools_dir: PathBuf,
|
||
pub target_name: String,
|
||
}
|
||
|
||
impl ShimEnv {
|
||
/// 提取当前代理程序的运行环境信息
|
||
pub fn new(current_exe: PathBuf) -> Result<Self, ShimError> {
|
||
let bin_dir = current_exe
|
||
.parent()
|
||
.ok_or_else(|| {
|
||
ShimError::PathResolutionError(format!(
|
||
"无法获取程序 [{}] 的父级 bin 目录",
|
||
current_exe.display()
|
||
))
|
||
})?
|
||
.to_path_buf();
|
||
println!("bin_dir目录 {}", bin_dir.display());
|
||
|
||
let root_dir = bin_dir
|
||
.parent()
|
||
.ok_or_else(|| {
|
||
ShimError::PathResolutionError(format!(
|
||
"无法获取 bin 目录 [{}] 的父级 root 目录",
|
||
bin_dir.display()
|
||
))
|
||
})?
|
||
.to_path_buf();
|
||
println!("root_dir 目录 {}", root_dir.display());
|
||
|
||
let tools_dir = root_dir.join("tools");
|
||
println!("tools_dir 目录 {}", tools_dir.display());
|
||
|
||
let target_name = current_exe
|
||
.file_stem()
|
||
.and_then(|s| s.to_str())
|
||
.ok_or_else(|| {
|
||
ShimError::PathResolutionError(format!(
|
||
"无法从路径 [{}] 提取有效的程序名称",
|
||
current_exe.display()
|
||
))
|
||
})?
|
||
.to_lowercase();
|
||
println!("程序名称 {}", target_name);
|
||
|
||
Ok(Self {
|
||
bin_dir,
|
||
root_dir,
|
||
tools_dir,
|
||
target_name,
|
||
})
|
||
}
|
||
}
|
||
|
||
#[derive(Debug, Clone)]
|
||
pub struct ShimConfig {
|
||
pub target_path: PathBuf,
|
||
pub args: Option<Vec<String>>,
|
||
pub envs: Option<HashMap<String, String>>,
|
||
}
|
||
|
||
impl ShimConfig {
|
||
/// 根据配置快速构建准备执行的 Command 对象
|
||
pub fn to_command(&self) -> Command {
|
||
let mut cmd = Command::new(&self.target_path);
|
||
|
||
if let Some(args) = &self.args {
|
||
cmd.args(args);
|
||
}
|
||
|
||
if let Some(envs) = &self.envs {
|
||
for (key, val) in envs {
|
||
// 直接应用环境变量（Lua 端已经处理好字符串拼接或列表合并）
|
||
cmd.env(key, val);
|
||
}
|
||
}
|
||
|
||
cmd
|
||
}
|
||
}
|
||
|
||
/// 实现 FromLua Trait，由 mlua 自动处理 Table 转换
|
||
impl FromLua for ShimConfig {
|
||
fn from_lua(value: Value, _lua: &Lua) -> mlua::Result<Self> {
|
||
match value {
|
||
Value::Table(table) => {
|
||
let path_str: String = table.get("path")?;
|
||
let args: Option<Vec<String>> = table.get("args")?;
|
||
// 解析 env Table
|
||
let mut envs_map = HashMap::new();
|
||
if let Ok(env_table) = table.get::<Table>("env") {
|
||
// 获取当前系统的路径分隔符（Windows 为 ";"，Unix 为 ":"）
|
||
#[cfg(windows)]
|
||
let sep = ";";
|
||
#[cfg(not(windows))]
|
||
let sep = ":";
|
||
|
||
for pair in env_table.pairs::<String, Value>() {
|
||
let (k, v) = pair?;
|
||
match v {
|
||
// 情况 1: 普通字符串，如 HOME = "C:/path" -> 直接覆盖
|
||
Value::String(s) => {
|
||
envs_map.insert(k, s.to_str()?.to_string());
|
||
}
|
||
// 情况 2: 数组 Table，如 PATH = { bin_dir, get_env("PATH") }
|
||
Value::Table(arr) => {
|
||
let paths: Vec<String> = arr
|
||
// 将 arr 作为序列（数组）处理，每个元素转为 String
|
||
.sequence_values::<String>()
|
||
.filter_map(|r| r.ok())
|
||
.filter(|s| !s.is_empty()) // 过滤空串，防止生成不必要的连续 ;;
|
||
.collect();
|
||
|
||
let combined = paths.join(sep);
|
||
envs_map.insert(k, combined);
|
||
}
|
||
_ => {}
|
||
}
|
||
}
|
||
}
|
||
|
||
let envs = if envs_map.is_empty() {
|
||
None
|
||
} else {
|
||
Some(envs_map)
|
||
};
|
||
|
||
Ok(ShimConfig {
|
||
target_path: PathBuf::from(path_str),
|
||
args,
|
||
envs,
|
||
})
|
||
}
|
||
_ => Err(mlua::Error::FromLuaConversionError {
|
||
from: value.type_name(),
|
||
to: "ShimConfig".into(),
|
||
message: Some("Expected a Lua table".to_string()),
|
||
}),
|
||
}
|
||
}
|
||
}
|
||
pub struct LuaRuntime {
|
||
lua: Lua,
|
||
}
|
||
|
||
impl LuaRuntime {
|
||
/// 初始化限定权限的 Lua 沙箱环境
|
||
pub fn new(shim_env: &ShimEnv) -> Result<Self, ShimError> {
|
||
// 只加载安全的标准库，剥离 os / io 等风险模块
|
||
let lua = Lua::new_with(
|
||
StdLib::TABLE | StdLib::STRING | StdLib::MATH | StdLib::PACKAGE,
|
||
mlua::LuaOptions::default(),
|
||
)
|
||
.map_err(|e| ShimError::EnvError(format!("初始化 Lua 失败: {}", e)))?;
|
||
|
||
let globals = lua.globals();
|
||
|
||
// 统一使用 POSIX 风格路径规范化路径字符串
|
||
let root_dir_str = normalize_path_for_lua(&shim_env.root_dir);
|
||
let tools_dir_str = normalize_path_for_lua(&shim_env.tools_dir);
|
||
|
||
// 1. 注入锚点变量
|
||
globals
|
||
.set("__SHIM_DIR__", root_dir_str.clone())
|
||
.map_err(|e| ShimError::EnvError(e.to_string()))?;
|
||
|
||
// 2. 安全暴露 get_env 供配置读取环境变量
|
||
let get_env = lua
|
||
.create_function(|_, key: String| -> mlua::Result<String> {
|
||
Ok(env::var(key).unwrap_or_default())
|
||
})
|
||
.map_err(|e| ShimError::EnvError(e.to_string()))?;
|
||
|
||
globals
|
||
.set("get_env", get_env)
|
||
.map_err(|e| ShimError::EnvError(e.to_string()))?;
|
||
|
||
// 3. 配置 package.path，确保 require 行为正常
|
||
if let Ok(package) = globals.get::<Table>("package") {
|
||
if let Ok(path) = package.get::<String>("path") {
|
||
let new_path = format!(
|
||
"{};{}/?.lua;{}/?/init.lua;{}/?.lua;{}/?/init.lua",
|
||
path, root_dir_str, root_dir_str, tools_dir_str, tools_dir_str
|
||
);
|
||
let _ = package.set("path", new_path);
|
||
}
|
||
}
|
||
|
||
Ok(Self { lua })
|
||
}
|
||
|
||
/// 执行指定脚本文件，直接返回完整的 Lua Table
|
||
pub fn evaluate_lua_script(&self, path: &Path) -> Result<Table, ShimError> {
|
||
let code = fs::read_to_string(path)?;
|
||
|
||
self.lua
|
||
.load(&code)
|
||
.set_name(path.to_string_lossy())
|
||
.eval::<Table>()
|
||
.map_err(|e| ShimError::LuaExecutionError {
|
||
file: path.display().to_string(),
|
||
source: e,
|
||
})
|
||
}
|
||
|
||
/// 将 Lua Value 解析转化为 ShimConfig 数据对象
|
||
pub fn parse_config(&self, value: Value) -> Result<ShimConfig, ShimError> {
|
||
ShimConfig::from_lua(value, &self.lua).map_err(|e| ShimError::InvalidConfig(e.to_string()))
|
||
}
|
||
}
|
||
pub struct Shim;
|
||
|
||
impl Shim {
|
||
pub fn init() -> Result<ShimConfig, ShimError> {
|
||
let current_exe = env::current_exe()
|
||
.map_err(|e| Error::new(ErrorKind::Other, format!("获取代理程序路径失败: {}", e)))?;
|
||
println!("当前目录 {}", current_exe.display());
|
||
let shim_env = ShimEnv::new(current_exe)?;
|
||
let runtime = LuaRuntime::new(&shim_env)?;
|
||
|
||
Self::resolve_config(&runtime, &shim_env)
|
||
}
|
||
|
||
fn resolve_config(engine: &LuaRuntime, env: &ShimEnv) -> Result<ShimConfig, ShimError> {
|
||
// 策略 1: 尝试加载全局配置文件 shims.lua
|
||
let global_config = env.root_dir.join("shims.lua");
|
||
if global_config.is_file() {
|
||
let root_table = engine.evaluate_lua_script(&global_config)?;
|
||
|
||
// 检查 shims.lua 中是否存在以 target_name 命名的 Table 节点
|
||
if let Ok(target_val) = root_table.get::<Value>(env.target_name.as_str()) {
|
||
if matches!(target_val, Value::Table(_)) {
|
||
return engine.parse_config(target_val);
|
||
}
|
||
}
|
||
// 穿透：若全局配置文件存在但未包含当前程序的 key，继续向下探查
|
||
}
|
||
|
||
// 策略 2: 降级寻找独立文件 ({exe}.lua)，优先顺序：tools/ > root/
|
||
let target_filename = format!("{}.lua", env.target_name);
|
||
let candidates = [
|
||
env.tools_dir.join(&target_filename),
|
||
env.root_dir.join(&target_filename),
|
||
];
|
||
|
||
for config_path in &candidates {
|
||
if config_path.is_file() {
|
||
let table = engine.evaluate_lua_script(config_path)?;
|
||
return engine.parse_config(Value::Table(table));
|
||
}
|
||
}
|
||
|
||
// 策略 3: 所有查找失败，抛出错误
|
||
Err(ShimError::ConfigNotFound(format!(
|
||
"未找到关于 '{}' 的配置。请检查 shims.lua 或特定的 {}.lua 文件",
|
||
env.target_name, env.target_name
|
||
)))
|
||
}
|
||
}
|