Add files via upload

README.md

@@ -10,6 +10,7 @@
   ![Preview](./img/mcp-stdio2.png)
 
 ## Changelog
+- 2025.11.14 Performance optimizations: optimized tool lookup from O(n) to O(1) using index map, added automatic cleanup mechanism for execution records to prevent memory leaks, and added pagination support for database queries
 - 2025.11.13 Added authentication for the web mode, including automatic password generation and in-app password change
 - 2025.11.13 Added `Settings` feature in the frontend
 - 2025.11.13 Added MCP Stdio mode support, now seamlessly integrated and usable in code editors, CLI, and automation scripts

README_CN.md

@@ -7,6 +7,7 @@
   ![详情预览](./img/mcp-stdio2.png) 
   
 ## 更新日志
+- 2025.11.14 性能优化：工具查找从 O(n) 优化为 O(1)（使用索引映射），添加执行记录自动清理机制防止内存泄漏，数据库查询支持分页加载
 - 2025.11.13 Web 端新增统一鉴权，支持自动生成强密码与前端修改密码；
 - 2025.11.13 在前端新增`设置`功能；
 - 2025.11.13 新增 MCP Stdio 模式支持，现可在代码编辑器、CLI 及自动化脚本等多种场景下，无缝集成并使用全套安全工具；

internal/database/monitor.go

@@ -69,16 +69,30 @@ func (db *DB) SaveToolExecution(exec *mcp.ToolExecution) error {
 	return nil
 }
 
-// LoadToolExecutions 加载所有工具执行记录
+// LoadToolExecutions 加载所有工具执行记录（支持分页）
 func (db *DB) LoadToolExecutions() ([]*mcp.ToolExecution, error) {
+	return db.LoadToolExecutionsWithPagination(0, 1000)
+}
+
+// LoadToolExecutionsWithPagination 分页加载工具执行记录
+// limit: 最大返回记录数，0 表示使用默认值 1000
+// offset: 跳过的记录数，用于分页
+func (db *DB) LoadToolExecutionsWithPagination(offset, limit int) ([]*mcp.ToolExecution, error) {
+	if limit <= 0 {
+		limit = 1000 // 默认限制
+	}
+	if limit > 10000 {
+		limit = 10000 // 最大限制，防止一次性加载过多数据
+	}
+	
 	query := `
 		SELECT id, tool_name, arguments, status, result, error, start_time, end_time, duration_ms
 		FROM tool_executions
 		ORDER BY start_time DESC
-		LIMIT 1000
+		LIMIT ? OFFSET ?
 	`
 
-	rows, err := db.Query(query)
+	rows, err := db.Query(query, limit, offset)
 	if err != nil {
 		return nil, err
 	}

internal/mcp/server.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"net/http"
 	"os"
+	"sort"
 	"strings"
 	"sync"
 	"time"
@@ -27,15 +28,16 @@ type MonitorStorage interface {
 
 // Server MCP服务器
 type Server struct {
-	tools      map[string]ToolHandler
+	tools                 map[string]ToolHandler
-	toolDefs   map[string]Tool // 工具定义
+	toolDefs              map[string]Tool // 工具定义
-	executions map[string]*ToolExecution
+	executions            map[string]*ToolExecution
-	stats      map[string]*ToolStats
+	stats                 map[string]*ToolStats
-	prompts    map[string]*Prompt   // 提示词模板
+	prompts               map[string]*Prompt   // 提示词模板
-	resources  map[string]*Resource // 资源
+	resources             map[string]*Resource // 资源
-	storage    MonitorStorage       // 可选的持久化存储
+	storage               MonitorStorage       // 可选的持久化存储
-	mu         sync.RWMutex
+	mu                    sync.RWMutex
-	logger     *zap.Logger
+	logger                *zap.Logger
+	maxExecutionsInMemory int // 内存中最大执行记录数
 }
 
 // ToolHandler 工具处理函数
@@ -49,14 +51,15 @@ func NewServer(logger *zap.Logger) *Server {
 // NewServerWithStorage 创建新的MCP服务器（带持久化存储）
 func NewServerWithStorage(logger *zap.Logger, storage MonitorStorage) *Server {
 	s := &Server{
-		tools:      make(map[string]ToolHandler),
+		tools:                 make(map[string]ToolHandler),
-		toolDefs:   make(map[string]Tool),
+		toolDefs:              make(map[string]Tool),
-		executions: make(map[string]*ToolExecution),
+		executions:            make(map[string]*ToolExecution),
-		stats:      make(map[string]*ToolStats),
+		stats:                 make(map[string]*ToolStats),
-		prompts:    make(map[string]*Prompt),
+		prompts:               make(map[string]*Prompt),
-		resources:  make(map[string]*Resource),
+		resources:             make(map[string]*Resource),
-		storage:    storage,
+		storage:               storage,
-		logger:     logger,
+		logger:                logger,
+		maxExecutionsInMemory: 1000, // 默认最多在内存中保留1000条执行记录
 	}
 
 	// 初始化默认提示词和资源
@@ -267,6 +270,8 @@ func (s *Server) handleCallTool(msg *Message) *Message {
 
 	s.mu.Lock()
 	s.executions[executionID] = execution
+	// 如果内存中的执行记录超过限制，清理最旧的记录
+	s.cleanupOldExecutions()
 	s.mu.Unlock()
 
 	if s.storage != nil {
@@ -499,9 +504,11 @@ func (s *Server) loadHistoricalData() {
 	} else {
 		s.mu.Lock()
 		for _, exec := range executions {
-			// 只加载最近1000条，避免内存占用过大
+			// 只加载最近 maxExecutionsInMemory 条，避免内存占用过大
-			if len(s.executions) < 1000 {
+			if len(s.executions) < s.maxExecutionsInMemory {
 				s.executions[exec.ID] = exec
+			} else {
+				break
 			}
 		}
 		s.mu.Unlock()
@@ -618,6 +625,8 @@ func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]
 
 	s.mu.Lock()
 	s.executions[executionID] = execution
+	// 如果内存中的执行记录超过限制，清理最旧的记录
+	s.cleanupOldExecutions()
 	s.mu.Unlock()
 
 	if s.storage != nil {
@@ -689,6 +698,43 @@ func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]
 	return finalResult, executionID, nil
 }
 
+// cleanupOldExecutions 清理旧的执行记录，防止内存无限增长
+func (s *Server) cleanupOldExecutions() {
+	if len(s.executions) <= s.maxExecutionsInMemory {
+		return
+	}
+
+	// 按开始时间排序，找出最旧的记录
+	type execWithTime struct {
+		id        string
+		startTime time.Time
+	}
+	execs := make([]execWithTime, 0, len(s.executions))
+	for id, exec := range s.executions {
+		execs = append(execs, execWithTime{
+			id:        id,
+			startTime: exec.StartTime,
+		})
+	}
+
+	// 使用 sort 包进行高效排序（最旧的在前）
+	sort.Slice(execs, func(i, j int) bool {
+		return execs[i].startTime.Before(execs[j].startTime)
+	})
+
+	// 删除最旧的记录，保留 maxExecutionsInMemory 条
+	toDelete := len(s.executions) - s.maxExecutionsInMemory
+	for i := 0; i < toDelete; i++ {
+		delete(s.executions, execs[i].id)
+	}
+
+	s.logger.Debug("清理旧的执行记录",
+		zap.Int("before", len(execs)),
+		zap.Int("after", len(s.executions)),
+		zap.Int("deleted", toDelete),
+	)
+}
+
 // initDefaultPrompts 初始化默认提示词模板
 func (s *Server) initDefaultPrompts() {
 	s.mu.Lock()

internal/security/executor.go

@@ -9,23 +9,43 @@ import (
 
 	"cyberstrike-ai/internal/config"
 	"cyberstrike-ai/internal/mcp"
+
 	"go.uber.org/zap"
 )
 
 // Executor 安全工具执行器
 type Executor struct {
 	config    *config.SecurityConfig
+	toolIndex map[string]*config.ToolConfig // 工具索引，用于 O(1) 查找
 	mcpServer *mcp.Server
 	logger    *zap.Logger
 }
 
 // NewExecutor 创建新的执行器
 func NewExecutor(cfg *config.SecurityConfig, mcpServer *mcp.Server, logger *zap.Logger) *Executor {
-	return &Executor{
+	executor := &Executor{
 		config:    cfg,
+		toolIndex: make(map[string]*config.ToolConfig),
 		mcpServer: mcpServer,
 		logger:    logger,
 	}
+	// 构建工具索引
+	executor.buildToolIndex()
+	return executor
+}
+
+// buildToolIndex 构建工具索引，将 O(n) 查找优化为 O(1)
+func (e *Executor) buildToolIndex() {
+	e.toolIndex = make(map[string]*config.ToolConfig)
+	for i := range e.config.Tools {
+		if e.config.Tools[i].Enabled {
+			e.toolIndex[e.config.Tools[i].Name] = &e.config.Tools[i]
+		}
+	}
+	e.logger.Info("工具索引构建完成",
+		zap.Int("totalTools", len(e.config.Tools)),
+		zap.Int("enabledTools", len(e.toolIndex)),
+	)
 }
 
 // ExecuteTool 执行安全工具
@@ -41,19 +61,13 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st
 		return e.executeSystemCommand(ctx, args)
 	}
 
-	// 查找工具配置
+	// 使用索引查找工具配置（O(1) 查找）
-	var toolConfig *config.ToolConfig
+	toolConfig, exists := e.toolIndex[toolName]
-	for i := range e.config.Tools {
+	if !exists {
-		if e.config.Tools[i].Name == toolName && e.config.Tools[i].Enabled {
-			toolConfig = &e.config.Tools[i]
-			break
-		}
-	}
-
-	if toolConfig == nil {
 		e.logger.Error("工具未找到或未启用",
 			zap.String("toolName", toolName),
 			zap.Int("totalTools", len(e.config.Tools)),
+			zap.Int("enabledTools", len(e.toolIndex)),
 		)
 		return nil, fmt.Errorf("工具 %s 未找到或未启用", toolName)
 	}
@@ -136,8 +150,12 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st
 func (e *Executor) RegisterTools(mcpServer *mcp.Server) {
 	e.logger.Info("开始注册工具",
 		zap.Int("totalTools", len(e.config.Tools)),
+		zap.Int("enabledTools", len(e.toolIndex)),
 	)
 
+	// 重新构建索引（以防配置更新）
+	e.buildToolIndex()
+
 	for i, toolConfig := range e.config.Tools {
 		if !toolConfig.Enabled {
 			e.logger.Debug("跳过未启用的工具",
@@ -638,9 +656,9 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
 // buildInputSchema 构建输入模式
 func (e *Executor) buildInputSchema(toolConfig *config.ToolConfig) map[string]interface{} {
 	schema := map[string]interface{}{
-		"type": "object",
+		"type":       "object",
 		"properties": map[string]interface{}{},
-		"required": []string{},
+		"required":   []string{},
 	}
 
 	// 如果配置中定义了参数，优先使用配置中的参数定义
@@ -750,5 +768,3 @@ func (e *Executor) GetVulnerabilityReport(vulnerabilities []Vulnerability) map[s
 		"generatedAt":     time.Now(),
 	}
 }
-
-