From 39f1c72755c83e87b4f5271309cb8c7942d16991 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=AC=E6=98=8E?= <83812544+Ed1s0nZ@users.noreply.github.com> Date: Mon, 29 Jun 2026 14:35:52 +0800 Subject: [PATCH] Add files via upload --- internal/config/config.go | 97 +++++++++++++++++------ internal/config/tools_reload_test.go | 111 +++++++++++++++++++++++++++ internal/handler/config.go | 11 +++ 3 files changed, 195 insertions(+), 24 deletions(-) create mode 100644 internal/config/tools_reload_test.go diff --git a/internal/config/config.go b/internal/config/config.go index 6de8704d..e4ca310b 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -828,33 +828,13 @@ func Load(path string) (*Config, error) { // 如果配置了工具目录,从目录加载工具配置 if cfg.Security.ToolsDir != "" { - configDir := filepath.Dir(path) - toolsDir := cfg.Security.ToolsDir - - // 如果是相对路径,相对于配置文件所在目录 - if !filepath.IsAbs(toolsDir) { - toolsDir = filepath.Join(configDir, toolsDir) - } - - tools, err := LoadToolsFromDir(toolsDir) + inlineTools := append([]ToolConfig(nil), cfg.Security.Tools...) + toolsDir := ResolveToolsDir(cfg.Security.ToolsDir, path) + merged, err := MergeToolsFromDir(toolsDir, inlineTools) if err != nil { return nil, fmt.Errorf("从工具目录加载工具配置失败: %w", err) } - - // 合并工具配置:目录中的工具优先,主配置中的工具作为补充 - existingTools := make(map[string]bool) - for _, tool := range tools { - existingTools[tool.Name] = true - } - - // 添加主配置中不存在于目录中的工具(向后兼容) - for _, tool := range cfg.Security.Tools { - if !existingTools[tool.Name] { - tools = append(tools, tool) - } - } - - cfg.Security.Tools = tools + cfg.Security.Tools = merged } // 外部 MCP:迁移 + 环境变量展开 @@ -1126,6 +1106,75 @@ func PrintMCPConfigJSON(mcp MCPConfig) { fmt.Println("----------------------------------------------------------------") } +// ResolveToolsDir 将 tools_dir 解析为绝对路径(相对路径相对于 configPath 所在目录)。 +func ResolveToolsDir(toolsDir, configPath string) string { + toolsDir = strings.TrimSpace(toolsDir) + if toolsDir == "" { + return "" + } + if filepath.IsAbs(toolsDir) { + return toolsDir + } + return filepath.Join(filepath.Dir(configPath), toolsDir) +} + +// MergeToolsFromDir 从目录加载工具并与 inline 列表合并:目录中的工具优先,主配置中的工具作为补充。 +func MergeToolsFromDir(toolsDir string, inlineTools []ToolConfig) ([]ToolConfig, error) { + dirTools, err := LoadToolsFromDir(toolsDir) + if err != nil { + return nil, err + } + existing := make(map[string]bool, len(dirTools)) + for _, tool := range dirTools { + existing[tool.Name] = true + } + merged := append([]ToolConfig(nil), dirTools...) + for _, tool := range inlineTools { + if !existing[tool.Name] { + merged = append(merged, tool) + } + } + return merged, nil +} + +// loadInlineSecurityToolsFromYAML 读取 config.yaml 中 security.tools(不含 tools_dir 扫描结果)。 +func loadInlineSecurityToolsFromYAML(configPath string) ([]ToolConfig, error) { + data, err := os.ReadFile(configPath) + if err != nil { + return nil, fmt.Errorf("读取配置文件失败: %w", err) + } + var partial struct { + Security struct { + Tools []ToolConfig `yaml:"tools"` + } `yaml:"security"` + } + if err := yaml.Unmarshal(data, &partial); err != nil { + return nil, fmt.Errorf("解析配置文件失败: %w", err) + } + if partial.Security.Tools == nil { + return []ToolConfig{}, nil + } + return partial.Security.Tools, nil +} + +// ReloadSecurityToolsFromDir 从 tools_dir 重新加载工具并更新 cfg.Security.Tools(ApplyConfig 热重载用)。 +func ReloadSecurityToolsFromDir(cfg *Config, configPath string) error { + if cfg == nil || strings.TrimSpace(cfg.Security.ToolsDir) == "" { + return nil + } + inlineTools, err := loadInlineSecurityToolsFromYAML(configPath) + if err != nil { + return err + } + toolsDir := ResolveToolsDir(cfg.Security.ToolsDir, configPath) + merged, err := MergeToolsFromDir(toolsDir, inlineTools) + if err != nil { + return fmt.Errorf("从工具目录加载工具配置失败: %w", err) + } + cfg.Security.Tools = merged + return nil +} + // LoadToolsFromDir 从目录加载所有工具配置文件 func LoadToolsFromDir(dir string) ([]ToolConfig, error) { var tools []ToolConfig diff --git a/internal/config/tools_reload_test.go b/internal/config/tools_reload_test.go new file mode 100644 index 00000000..0fbbd074 --- /dev/null +++ b/internal/config/tools_reload_test.go @@ -0,0 +1,111 @@ +package config + +import ( + "os" + "path/filepath" + "testing" +) + +func TestReloadSecurityToolsFromDir(t *testing.T) { + root := t.TempDir() + toolsDir := filepath.Join(root, "tools") + if err := os.MkdirAll(toolsDir, 0755); err != nil { + t.Fatal(err) + } + + configPath := filepath.Join(root, "config.yaml") + if err := os.WriteFile(configPath, []byte(`security: + tools_dir: tools + tools: + - name: inline-only + command: inline-cmd + enabled: true + description: inline tool +`), 0644); err != nil { + t.Fatal(err) + } + + writeTool := func(name, command string) { + t.Helper() + content := "name: " + name + "\ncommand: " + command + "\nenabled: true\ndescription: test\n" + if err := os.WriteFile(filepath.Join(toolsDir, name+".yaml"), []byte(content), 0644); err != nil { + t.Fatal(err) + } + } + + writeTool("alpha", "alpha-cmd") + + cfg := &Config{ + Security: SecurityConfig{ + ToolsDir: "tools", + Tools: []ToolConfig{ + {Name: "stale", Command: "stale-cmd", Enabled: true, Description: "should be removed"}, + }, + }, + } + + if err := ReloadSecurityToolsFromDir(cfg, configPath); err != nil { + t.Fatalf("reload: %v", err) + } + if len(cfg.Security.Tools) != 2 { + t.Fatalf("expected 2 tools, got %d", len(cfg.Security.Tools)) + } + + names := map[string]string{} + for _, tool := range cfg.Security.Tools { + names[tool.Name] = tool.Command + } + if names["alpha"] != "alpha-cmd" { + t.Fatalf("alpha tool missing or wrong command: %#v", names) + } + if names["inline-only"] != "inline-cmd" { + t.Fatalf("inline-only tool missing: %#v", names) + } + if _, ok := names["stale"]; ok { + t.Fatal("stale in-memory tool should not survive reload") + } + + writeTool("beta", "beta-cmd") + if err := ReloadSecurityToolsFromDir(cfg, configPath); err != nil { + t.Fatalf("second reload: %v", err) + } + if len(cfg.Security.Tools) != 3 { + t.Fatalf("expected 3 tools after add, got %d", len(cfg.Security.Tools)) + } + foundBeta := false + for _, tool := range cfg.Security.Tools { + if tool.Name == "beta" { + foundBeta = true + break + } + } + if !foundBeta { + t.Fatal("beta tool not found after second reload") + } +} + +func TestMergeToolsFromDir_DirOverridesInline(t *testing.T) { + root := t.TempDir() + toolsDir := filepath.Join(root, "tools") + if err := os.MkdirAll(toolsDir, 0755); err != nil { + t.Fatal(err) + } + content := "name: shared\ncommand: dir-cmd\nenabled: true\ndescription: from dir\n" + if err := os.WriteFile(filepath.Join(toolsDir, "shared.yaml"), []byte(content), 0644); err != nil { + t.Fatal(err) + } + + inline := []ToolConfig{ + {Name: "shared", Command: "inline-cmd", Enabled: true, Description: "from inline"}, + } + merged, err := MergeToolsFromDir(toolsDir, inline) + if err != nil { + t.Fatal(err) + } + if len(merged) != 1 { + t.Fatalf("expected 1 tool, got %d", len(merged)) + } + if merged[0].Command != "dir-cmd" { + t.Fatalf("dir tool should win, got command %q", merged[0].Command) + } +} diff --git a/internal/handler/config.go b/internal/handler/config.go index 93022bfb..3dd4d804 100644 --- a/internal/handler/config.go +++ b/internal/handler/config.go @@ -1333,6 +1333,17 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) { h.logger.Info("已更新嵌入模型配置记录") } + // 从 tools 目录重新加载工具配置(新增/修改/删除 yaml 后无需重启) + if err := config.ReloadSecurityToolsFromDir(h.config, h.configPath); err != nil { + h.logger.Error("重新加载工具配置失败", zap.Error(err)) + if h.audit != nil { + h.audit.RecordFail(c, "config", "apply", "应用配置失败:重新加载工具", map[string]interface{}{"error": err.Error()}) + } + c.JSON(http.StatusInternalServerError, gin.H{"error": "重新加载工具配置失败: " + err.Error()}) + return + } + h.logger.Info("已从 tools 目录重新加载工具配置", zap.Int("tools_count", len(h.config.Security.Tools))) + // 重新注册工具(根据新的启用状态) h.logger.Info("重新注册工具")