From ea1d9f2306042be93df334df02bf442481ec4dfc Mon Sep 17 00:00:00 2001 From: limitcool Date: Wed, 1 Apr 2026 16:12:29 +0800 Subject: [PATCH] fix: share runtime env injection for reposync and terminal --- internal/controllers/terminal_controller.go | 53 +++++----- internal/services/tasks/executor_service.go | 63 +----------- internal/utils/runtime_env.go | 104 ++++++++++++++++++++ 3 files changed, 135 insertions(+), 85 deletions(-) create mode 100644 internal/utils/runtime_env.go diff --git a/internal/controllers/terminal_controller.go b/internal/controllers/terminal_controller.go index f961237..5f12de4 100644 --- a/internal/controllers/terminal_controller.go +++ b/internal/controllers/terminal_controller.go @@ -1,8 +1,8 @@ package controllers import ( - "encoding/json" "bufio" + "encoding/json" "io" "net/http" "os" @@ -89,7 +89,6 @@ func (tc *TerminalController) HandleWebSocket(c *gin.Context) { if userID == "" { userID = "1" // 兜底 } - if runtime.GOOS == "windows" { tc.handlePipeMode(conn, userID) } else { @@ -108,17 +107,7 @@ func (tc *TerminalController) handlePtyMode(conn *websocket.Conn, userID string) cmd.Dir = absDir } - cmd.Env = append(os.Environ(), "TERM=xterm-256color") - - // 注入 baihu 命令环境变量 - if absBinDir, err := filepath.Abs(filepath.Join(constant.DataDir, "bin")); err == nil { - pathStr := absBinDir + string(os.PathListSeparator) + os.Getenv("PATH") - cmd.Env = append(cmd.Env, "PATH="+pathStr) - } - - // 注入环境变量(支持同名合并) - envVars := tc.envService.GetFormattedEnvVarsByUserID(userID) - cmd.Env = append(cmd.Env, envVars...) + cmd.Env = tc.buildTerminalEnv(userID, "TERM=xterm-256color") ptmx, err := pty.Start(cmd) if err != nil { @@ -195,17 +184,7 @@ func (tc *TerminalController) handlePipeMode(conn *websocket.Conn, userID string } // 注入环境变量 - cmd.Env = os.Environ() - - // 注入 baihu 命令环境变量 - if absBinDir, err := filepath.Abs(filepath.Join(constant.DataDir, "bin")); err == nil { - pathStr := absBinDir + string(os.PathListSeparator) + os.Getenv("PATH") - cmd.Env = append(cmd.Env, "PATH="+pathStr) - } - - // 注入环境变量(支持同名合并) - envVars := tc.envService.GetFormattedEnvVarsByUserID(userID) - cmd.Env = append(cmd.Env, envVars...) + cmd.Env = tc.buildTerminalEnv(userID) stdin, err := cmd.StdinPipe() if err != nil { @@ -303,6 +282,11 @@ func (tc *TerminalController) ExecuteShellCommand(c *gin.Context) { } cmd := utils.NewShellCommandCmd(req.Command) + userID := c.GetString("userID") + if userID == "" { + userID = "1" // 与 WebSocket 终端保持一致,保留原有兜底行为 + } + cmd.Env = tc.buildTerminalEnv(userID) output, err := cmd.CombinedOutput() if err != nil { @@ -318,6 +302,27 @@ func (tc *TerminalController) ExecuteShellCommand(c *gin.Context) { }) } +func (tc *TerminalController) buildTerminalEnv(userID string, extraEnvs ...string) []string { + env := os.Environ() + env = append(env, extraEnvs...) + + // 注入 baihu 命令运行时路径与配置,保证在终端里手动执行 baihu 子命令时 + // 仍然能连接到与主服务一致的数据库,而不是回退到默认 sqlite。 + if absBinDir, err := filepath.Abs(filepath.Join(constant.DataDir, "bin")); err == nil { + pathStr := absBinDir + string(os.PathListSeparator) + os.Getenv("PATH") + env = append(env, "PATH="+pathStr) + } + env = append(env, utils.BuildRuntimeProcessEnv()...) + + // 注入环境变量(支持同名合并) + if tc.envService != nil { + envVars := tc.envService.GetFormattedEnvVarsByUserID(userID) + env = append(env, envVars...) + } + + return env +} + // GetCommands 获取所有可用的 cmd 列表及说明 func (tc *TerminalController) GetCommands(c *gin.Context) { var cmds []map[string]string diff --git a/internal/services/tasks/executor_service.go b/internal/services/tasks/executor_service.go index 5b796e8..85860aa 100644 --- a/internal/services/tasks/executor_service.go +++ b/internal/services/tasks/executor_service.go @@ -1075,68 +1075,9 @@ func (es *ExecutorService) ResolvePath(path string) string { } func buildRepoCommandEnvPrefix() string { - absConfig, err := filepath.Abs(constant.ConfigPath) - if err != nil { - absConfig = constant.ConfigPath - } - - absScriptsDir := resolveAbsScriptsDir() - parts := []string{ - shellEnvAssignment("BH_CONFIG_PATH", absConfig), - shellEnvAssignment("BH_SCRIPTS_DIR", absScriptsDir), - } - - appendIfSet := func(key string, value string) { - if strings.TrimSpace(value) == "" { - return - } - parts = append(parts, shellEnvAssignment(key, value)) - } - - appendIfSet("BH_DB_TYPE", constant.RuntimeDBType) - appendIfSet("BH_DB_HOST", constant.RuntimeDBHost) - if constant.RuntimeDBPort > 0 { - parts = append(parts, shellEnvAssignment("BH_DB_PORT", fmt.Sprintf("%d", constant.RuntimeDBPort))) - } - appendIfSet("BH_DB_USER", constant.RuntimeDBUser) - appendIfSet("BH_DB_PASSWORD", constant.RuntimeDBPassword) - appendIfSet("BH_DB_NAME", constant.RuntimeDBName) - appendIfSet("BH_DB_PATH", constant.RuntimeDBPath) - appendIfSet("BH_DB_DSN", constant.RuntimeDBDSN) - appendIfSet("BH_DB_TABLE_PREFIX", constant.RuntimeDBTablePrefix) - - return strings.Join(parts, " ") + " " -} - -func shellEnvAssignment(key, value string) string { - return key + "='" + strings.ReplaceAll(value, "'", "'\\''") + "'" + return utils.BuildShellEnvPrefix(utils.BuildRuntimeProcessEnv()) } func resolveAbsScriptsDir() string { - if scriptsDir := os.Getenv("BH_SCRIPTS_DIR"); scriptsDir != "" { - if filepath.IsAbs(scriptsDir) { - return filepath.Clean(scriptsDir) - } - if absScriptsDir, err := filepath.Abs(scriptsDir); err == nil { - return absScriptsDir - } - return filepath.Clean(scriptsDir) - } - - if configPath := os.Getenv("BH_CONFIG_PATH"); configPath != "" { - if !filepath.IsAbs(configPath) { - if absConfigPath, err := filepath.Abs(configPath); err == nil { - configPath = absConfigPath - } - } - - projectRoot := filepath.Dir(filepath.Dir(configPath)) - return filepath.Clean(filepath.Join(projectRoot, constant.ScriptsWorkDir)) - } - - if absScriptsDir, err := filepath.Abs(constant.ScriptsWorkDir); err == nil { - return absScriptsDir - } - - return filepath.Clean(constant.ScriptsWorkDir) + return utils.ResolveAbsScriptsDir() } diff --git a/internal/utils/runtime_env.go b/internal/utils/runtime_env.go new file mode 100644 index 0000000..4d7672e --- /dev/null +++ b/internal/utils/runtime_env.go @@ -0,0 +1,104 @@ +package utils + +import ( + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/engigu/baihu-panel/internal/constant" +) + +// BuildRuntimeProcessEnv 构造 Baihu 内部可信子进程需要继承的运行时环境变量。 +// 仅包含 Baihu 自己的路径/数据库配置,不包含用户任务环境变量。 +func BuildRuntimeProcessEnv() []string { + envs := make([]string, 0, 11) + + configPath := constant.ConfigPath + if absConfig, err := filepath.Abs(constant.ConfigPath); err == nil { + configPath = absConfig + } + envs = append(envs, formatEnvVar("BH_CONFIG_PATH", configPath)) + + if scriptsDir := ResolveAbsScriptsDir(); strings.TrimSpace(scriptsDir) != "" { + envs = append(envs, formatEnvVar("BH_SCRIPTS_DIR", scriptsDir)) + } + + appendEnvIfSet(&envs, "BH_DB_TYPE", constant.RuntimeDBType) + appendEnvIfSet(&envs, "BH_DB_HOST", constant.RuntimeDBHost) + if constant.RuntimeDBPort > 0 { + envs = append(envs, formatEnvVar("BH_DB_PORT", fmt.Sprintf("%d", constant.RuntimeDBPort))) + } + appendEnvIfSet(&envs, "BH_DB_USER", constant.RuntimeDBUser) + appendEnvIfSet(&envs, "BH_DB_PASSWORD", constant.RuntimeDBPassword) + appendEnvIfSet(&envs, "BH_DB_NAME", constant.RuntimeDBName) + appendEnvIfSet(&envs, "BH_DB_PATH", constant.RuntimeDBPath) + appendEnvIfSet(&envs, "BH_DB_DSN", constant.RuntimeDBDSN) + appendEnvIfSet(&envs, "BH_DB_TABLE_PREFIX", constant.RuntimeDBTablePrefix) + + return envs +} + +// BuildShellEnvPrefix 将 KEY=VALUE 环境变量切片转换为 shell 前缀。 +func BuildShellEnvPrefix(envs []string) string { + parts := make([]string, 0, len(envs)) + for _, env := range envs { + key, value, ok := strings.Cut(env, "=") + if !ok || strings.TrimSpace(key) == "" { + continue + } + parts = append(parts, ShellEnvAssignment(key, value)) + } + + if len(parts) == 0 { + return "" + } + + return strings.Join(parts, " ") + " " +} + +// ShellEnvAssignment 生成 shell 可安全使用的 KEY='VALUE' 赋值片段。 +func ShellEnvAssignment(key, value string) string { + return key + "='" + strings.ReplaceAll(value, "'", "'\\''") + "'" +} + +// ResolveAbsScriptsDir 解析 Baihu 运行时脚本目录的绝对路径。 +func ResolveAbsScriptsDir() string { + if scriptsDir := os.Getenv("BH_SCRIPTS_DIR"); scriptsDir != "" { + if filepath.IsAbs(scriptsDir) { + return filepath.Clean(scriptsDir) + } + if absScriptsDir, err := filepath.Abs(scriptsDir); err == nil { + return absScriptsDir + } + return filepath.Clean(scriptsDir) + } + + if configPath := os.Getenv("BH_CONFIG_PATH"); configPath != "" { + if !filepath.IsAbs(configPath) { + if absConfigPath, err := filepath.Abs(configPath); err == nil { + configPath = absConfigPath + } + } + + projectRoot := filepath.Dir(filepath.Dir(configPath)) + return filepath.Clean(filepath.Join(projectRoot, constant.ScriptsWorkDir)) + } + + if absScriptsDir, err := filepath.Abs(constant.ScriptsWorkDir); err == nil { + return absScriptsDir + } + + return filepath.Clean(constant.ScriptsWorkDir) +} + +func appendEnvIfSet(envs *[]string, key, value string) { + if strings.TrimSpace(value) == "" { + return + } + *envs = append(*envs, formatEnvVar(key, value)) +} + +func formatEnvVar(key, value string) string { + return fmt.Sprintf("%s=%s", key, value) +}