From 5e6281ae997ee6a4490fa8d296c6d96c4db6e5e8 Mon Sep 17 00:00:00 2001 From: kory Date: Fri, 2 Oct 2026 14:32:40 +0800 Subject: [PATCH] feat: add constrained read-only diagnostics --- cmd/agent/main.go | 11 +++ internal/api/mcp_adapter.go | 27 +++++++ internal/cluster/manager.go | 14 ++++ internal/executor/diagnostic.go | 117 +++++++++++++++++++++++++++ internal/executor/diagnostic_test.go | 24 ++++++ internal/mcp/mcp.go | 24 ++++++ internal/mcp/mcp_test.go | 2 +- internal/security/registry.go | 1 + 8 files changed, 219 insertions(+), 1 deletion(-) create mode 100644 internal/executor/diagnostic.go create mode 100644 internal/executor/diagnostic_test.go diff --git a/cmd/agent/main.go b/cmd/agent/main.go index d885fb6..a2ea86e 100644 --- a/cmd/agent/main.go +++ b/cmd/agent/main.go @@ -152,6 +152,17 @@ func run(hub, name, token string, insecure bool) error { func handle(writeJSON func(cluster.Message) error, exec executor.Local, mcpRuntime *cluster.MCPRuntime, sessions *executor.SessionManager, msg cluster.Message) error { switch msg.Action { + case "diagnostic_command": + var p executor.DiagnosticRequest + if err := decodePayload(msg.Payload, &p); err != nil { + return sendError(writeJSON, msg.ID, err) + } + v, err := executor.Diagnostic(context.Background(), p) + if err != nil { + return sendError(writeJSON, msg.ID, err) + } + return writeJSON(cluster.Message{ID: msg.ID, Status: "success", Payload: v}) + case "mcp_list": return writeJSON(cluster.Message{ID: msg.ID, Status: "success", Payload: mcpRuntime.List()}) case "mcp_test": diff --git a/internal/api/mcp_adapter.go b/internal/api/mcp_adapter.go index e62876b..75bf807 100644 --- a/internal/api/mcp_adapter.go +++ b/internal/api/mcp_adapter.go @@ -98,6 +98,33 @@ func (b *MCPBackend) ExecCommand(ctx context.Context, clientIP, device, command, }, nil } +func (b *MCPBackend) DiagnosticCommand(ctx context.Context, clientIP, device string, req executor.DiagnosticRequest) (any, error) { + if device == "" { + device = b.server.LocalName + } + if err := b.authorize(device, "diagnostic", clientIP); err != nil { + return nil, err + } + if req.Action == "" { + return nil, errors.New("diagnostic action required") + } + if b.server.isRemote(device) { + if b.server.Cluster == nil || !b.server.Cluster.Online(device) { + return nil, fmt.Errorf("device %s is offline", device) + } + res, err := b.server.Cluster.Diagnostic(ctx, device, req) + if err != nil { + return nil, err + } + return res, nil + } + res, err := executor.Diagnostic(ctx, req) + if err != nil { + return nil, err + } + return res, nil +} + func (b *MCPBackend) ReadFile(ctx context.Context, clientIP, device, path string, offset, limit int64) (any, error) { if device == "" { device = b.server.LocalName diff --git a/internal/cluster/manager.go b/internal/cluster/manager.go index a104f0d..707a1e4 100644 --- a/internal/cluster/manager.go +++ b/internal/cluster/manager.go @@ -13,6 +13,7 @@ import ( "time" "github.com/gorilla/websocket" + "github.com/wjyhk/vps-commander/internal/executor" "github.com/wjyhk/vps-commander/internal/security" "github.com/wjyhk/vps-commander/internal/storage" "github.com/wjyhk/vps-commander/internal/sysinfo" @@ -464,6 +465,19 @@ func (m *Manager) Exec(ctx context.Context, name string, p ExecPayload) (Result, return Result{msg.ExitCode, msg.Stdout, msg.Stderr, msg.DurationMS}, nil } +func (m *Manager) Diagnostic(ctx context.Context, name string, req executor.DiagnosticRequest) (executor.DiagnosticResult, error) { + msg, err := m.call(ctx, name, "diagnostic_command", req) + if err != nil { + return executor.DiagnosticResult{}, err + } + b, _ := json.Marshal(msg.Payload) + var out executor.DiagnosticResult + if err := json.Unmarshal(b, &out); err != nil { + return out, err + } + return out, nil +} + func (m *Manager) ReadFile(ctx context.Context, name string, p FileReadPayload) (FileResult, error) { msg, err := m.call(ctx, name, "file_read", p) if err != nil { diff --git a/internal/executor/diagnostic.go b/internal/executor/diagnostic.go new file mode 100644 index 0000000..1faa3a2 --- /dev/null +++ b/internal/executor/diagnostic.go @@ -0,0 +1,117 @@ +package executor + +import ( + "context" + "errors" + "fmt" + "os/exec" + "strconv" + "strings" + "time" +) + +// DiagnosticRequest is a structured, read-only diagnostic operation. It is deliberately +// not a shell command: every action maps to a fixed executable/argument set. +type DiagnosticRequest struct { + Action string `json:"action"` + Target string `json:"target,omitempty"` + Lines int `json:"lines,omitempty"` +} + +type DiagnosticResult struct { + Action string `json:"action"` + Target string `json:"target,omitempty"` + ExitCode int `json:"exit_code"` + Stdout string `json:"stdout"` + Stderr string `json:"stderr"` + DurationMS int64 `json:"duration_ms"` +} + +func Diagnostic(ctx context.Context, req DiagnosticRequest) (DiagnosticResult, error) { + args, err := diagnosticArgs(req) + if err != nil { + return DiagnosticResult{}, err + } + if req.Lines <= 0 { + req.Lines = 100 + } + if req.Lines > 500 { + req.Lines = 500 + } + // docker logs is the only operation where a caller-controlled numeric limit is useful. + if req.Action == "docker_logs" { + args = []string{"logs", "--tail", strconv.Itoa(req.Lines), req.Target} + } + + start := time.Now() + c, cancel := context.WithTimeout(ctx, 15*time.Second) + defer cancel() + cmd := exec.CommandContext(c, args[0], args[1:]...) + out, runErr := cmd.CombinedOutput() + const maxOutput = 256 * 1024 + if len(out) > maxOutput { + out = append(out[:maxOutput], []byte("\n[output truncated at 256 KiB]\n")...) + } + code := 0 + if cmd.ProcessState != nil { + code = cmd.ProcessState.ExitCode() + } + result := DiagnosticResult{Action: req.Action, Target: req.Target, ExitCode: code, Stdout: string(out), DurationMS: time.Since(start).Milliseconds()} + if runErr != nil && code == -1 { + return result, fmt.Errorf("diagnostic failed: %w", runErr) + } + return result, nil +} + +func diagnosticArgs(req DiagnosticRequest) ([]string, error) { + target := strings.TrimSpace(req.Target) + switch req.Action { + case "docker_ps": + if target != "" { + return nil, errors.New("docker_ps does not accept target") + } + return []string{"docker", "ps", "--no-trunc"}, nil + case "docker_logs": + if target == "" || strings.ContainsAny(target, " \t\r\n;|&<>$`") { + return nil, errors.New("invalid docker container target") + } + return []string{"docker", "logs", "--tail", "100", target}, nil + case "docker_inspect": + if target == "" || strings.ContainsAny(target, " \t\r\n;|&<>$`") { + return nil, errors.New("invalid docker target") + } + return []string{"docker", "inspect", target}, nil + case "systemctl_status": + if target == "" || strings.ContainsAny(target, " \t\r\n;|&<>$`") { + return nil, errors.New("invalid systemd unit") + } + return []string{"systemctl", "status", "--no-pager", "--full", target}, nil + case "journalctl": + if target == "" || strings.ContainsAny(target, " \t\r\n;|&<>$`") { + return nil, errors.New("invalid journal unit") + } + return []string{"journalctl", "--no-pager", "-u", target, "-n", "100"}, nil + case "ss_listen": + if target != "" { + return nil, errors.New("ss_listen does not accept target") + } + return []string{"ss", "-lntp"}, nil + case "df": + if target != "" { + return nil, errors.New("df does not accept target") + } + return []string{"df", "-h"}, nil + case "free": + if target != "" { + return nil, errors.New("free does not accept target") + } + return []string{"free", "-h"}, nil + case "uptime": + if target != "" { + return nil, errors.New("uptime does not accept target") + } + return []string{"uptime"}, nil + default: + return nil, fmt.Errorf("unsupported diagnostic action: %s", req.Action) + } +} diff --git a/internal/executor/diagnostic_test.go b/internal/executor/diagnostic_test.go new file mode 100644 index 0000000..99826b8 --- /dev/null +++ b/internal/executor/diagnostic_test.go @@ -0,0 +1,24 @@ +package executor + +import ( + "context" + "strings" + "testing" +) + +func TestDiagnosticRejectsShellInjection(t *testing.T) { + _, err := Diagnostic(context.Background(), DiagnosticRequest{Action: "docker_logs", Target: "edgeever;rm -rf /"}) + if err == nil || !strings.Contains(err.Error(), "invalid docker") { + t.Fatalf("expected target rejection, got %v", err) + } +} + +func TestDiagnosticAllowsReadOnlyUptime(t *testing.T) { + got, err := Diagnostic(context.Background(), DiagnosticRequest{Action: "uptime"}) + if err != nil { + t.Fatal(err) + } + if got.ExitCode != 0 || strings.TrimSpace(got.Stdout) == "" { + t.Fatalf("unexpected uptime result: %#v", got) + } +} diff --git a/internal/mcp/mcp.go b/internal/mcp/mcp.go index fc1682a..bb6b556 100644 --- a/internal/mcp/mcp.go +++ b/internal/mcp/mcp.go @@ -12,6 +12,7 @@ import ( "time" "github.com/wjyhk/vps-commander/internal/auth" + "github.com/wjyhk/vps-commander/internal/executor" ) // MCP Protocol / JSON-RPC 2.0 Types @@ -77,6 +78,10 @@ type ActionBackend interface { EditBlock(ctx context.Context, clientIP, device, path, oldText, newText string) (any, error) } +type DiagnosticBackend interface { + DiagnosticCommand(ctx context.Context, clientIP, device string, req executor.DiagnosticRequest) (any, error) +} + type ExtendedBackend interface { ReadMultipleFiles(ctx context.Context, clientIP, device string, paths []string, limit int64) (any, error) CreateDirectory(ctx context.Context, clientIP, device, path string) (any, error) @@ -280,6 +285,7 @@ func (s *Server) handleRequest(ctx context.Context, r *http.Request, req JSONRPC {Name: "create_directory", Description: "创建目录", InputSchema: InputSchema{Type: "object", Properties: map[string]PropertyDef{"device": {Type: "string", Description: "目标设备"}, "path": {Type: "string", Description: "目录路径"}}, Required: []string{"path"}}}, {Name: "move_file", Description: "移动或重命名文件", InputSchema: InputSchema{Type: "object", Properties: map[string]PropertyDef{"device": {Type: "string", Description: "目标设备"}, "source": {Type: "string", Description: "源路径"}, "destination": {Type: "string", Description: "目标路径"}}, Required: []string{"source", "destination"}}}, {Name: "list_processes", Description: "列出受控节点当前进程", InputSchema: InputSchema{Type: "object", Properties: map[string]PropertyDef{"device": {Type: "string", Description: "目标设备"}}}}, + {Name: "diagnostic_command", Description: "执行严格白名单的只读系统诊断，不接受任意 Shell", InputSchema: InputSchema{Type: "object", Properties: map[string]PropertyDef{"device": {Type: "string", Description: "目标设备"}, "action": {Type: "string", Description: "docker_ps|docker_logs|docker_inspect|systemctl_status|journalctl|ss_listen|df|free|uptime"}, "target": {Type: "string", Description: "容器名或 systemd unit（按 action 使用）"}, "lines": {Type: "integer", Description: "docker_logs 最大行数，1-500"}}, Required: []string{"action"}}}, { Name: "edit_block", Description: "对指定文件执行唯一文本块的原子替换；old_text 必须恰好匹配一次", @@ -430,6 +436,24 @@ func (s *Server) callTool(ctx context.Context, clientIP string, params CallToolP return map[string]string{"error": err.Error()}, true } return res, false + case "diagnostic_command": + var a struct { + Device string `json:"device"` + Action string `json:"action"` + Target string `json:"target"` + Lines int `json:"lines"` + } + _ = json.Unmarshal(params.Arguments, &a) + b, ok := s.Backend.(DiagnosticBackend) + if !ok { + return map[string]string{"error": "diagnostic tools unavailable"}, true + } + res, err := b.DiagnosticCommand(ctx, clientIP, a.Device, executor.DiagnosticRequest{Action: a.Action, Target: a.Target, Lines: a.Lines}) + if err != nil { + return map[string]string{"error": err.Error()}, true + } + return res, false + case "list_processes": var a struct { Device string `json:"device"` diff --git a/internal/mcp/mcp_test.go b/internal/mcp/mcp_test.go index 1a95473..29c51d5 100644 --- a/internal/mcp/mcp_test.go +++ b/internal/mcp/mcp_test.go @@ -79,7 +79,7 @@ func TestMCPInitializeAndListTools(t *testing.T) { t.Fatalf("expected map result") } tools, ok := resMap["tools"].([]any) - if !ok || len(tools) != 11 { + if !ok || len(tools) != 12 { t.Fatalf("expected 11 tools, got %d", len(tools)) } foundWriteFile := false diff --git a/internal/security/registry.go b/internal/security/registry.go index 0945baa..856d987 100644 --- a/internal/security/registry.go +++ b/internal/security/registry.go @@ -15,6 +15,7 @@ var registry = map[string]ToolSpec{ "start_search": {Name: "start_search", RequiredRisk: RiskLow, Category: "search", ReadOnly: true}, "get_more_search_results": {Name: "get_more_search_results", RequiredRisk: RiskLow, Category: "search", ReadOnly: true}, "list_processes": {Name: "list_processes", RequiredRisk: RiskLow, Category: "process", ReadOnly: true}, + "diagnostic_command": {Name: "diagnostic_command", RequiredRisk: RiskLow, Category: "diagnostic", ReadOnly: true}, "list_sessions": {Name: "list_sessions", RequiredRisk: RiskLow, Category: "process", ReadOnly: true}, "read_process_output": {Name: "read_process_output", RequiredRisk: RiskLow, Category: "process", ReadOnly: true}, "edit_block": {Name: "edit_block", RequiredRisk: RiskMedium, Category: "file", ReadOnly: false},