feat: add v2 phase2 security policy

2026-09-30 13:17:19 +08:00
parent d75c500a62
commit bd44465eaa
15 changed files with 773 additions and 19 deletions
@@ -407,6 +407,25 @@ func (m *Manager) CallMCP(ctx context.Context, name, server, tool string, args a
 	return out, nil
 }

+func (m *Manager) TestMCP(ctx context.Context, name string, service MCPService) ([]map[string]any, error) {
+	m.mu.RLock()
+	n := m.nodes[name]
+	m.mu.RUnlock()
+	if n == nil || !m.Online(name) {
+		return nil, fmt.Errorf("device %s is offline", name)
+	}
+	msg, err := m.call(ctx, name, "mcp_test", service)
+	if err != nil {
+		return nil, err
+	}
+	var out []map[string]any
+	b, _ := json.Marshal(msg.Payload)
+	if err := json.Unmarshal(b, &out); err != nil {
+		return nil, err
+	}
+	return out, nil
+}
+
 func (m *Manager) heartbeat(n *Node) {
 	t := time.NewTicker(30 * time.Second)
 	defer t.Stop()
@@ -105,6 +105,160 @@ func (r *MCPRuntime) Call(ctx context.Context, server, tool string, args any) (a
 	}
 }

+func ProbeMCP(ctx context.Context, s MCPService) ([]map[string]any, error) {
+	switch strings.ToLower(s.Transport) {
+	case "stdio":
+		return probeStdio(ctx, s)
+	case "http", "sse":
+		return probeHTTP(ctx, s)
+	default:
+		return nil, fmt.Errorf("unsupported MCP transport %q", s.Transport)
+	}
+}
+
+func probeStdio(ctx context.Context, s MCPService) ([]map[string]any, error) {
+	if s.Command == "" {
+		return nil, errors.New("MCP stdio command is empty")
+	}
+	cmd := exec.CommandContext(ctx, s.Command, s.Args...)
+	if len(s.Env) > 0 {
+		cmd.Env = append(os.Environ(), envList(s.Env)...)
+	}
+	stdin, err := cmd.StdinPipe()
+	if err != nil {
+		return nil, err
+	}
+	stdout, err := cmd.StdoutPipe()
+	if err != nil {
+		return nil, err
+	}
+	cmd.Stderr = io.Discard
+	if err := cmd.Start(); err != nil {
+		return nil, err
+	}
+	defer func() { _ = cmd.Process.Kill(); _ = cmd.Wait() }()
+	enc := json.NewEncoder(stdin)
+	dec := bufio.NewReaderSize(stdout, 64*1024)
+	if err := enc.Encode(mcpRPC{JSONRPC: "2.0", ID: int64(1), Method: "initialize", Params: map[string]any{
+		"protocolVersion": "2024-11-05", "capabilities": map[string]any{},
+		"clientInfo": map[string]string{"name": "vps-commander-probe", "version": "1.0"},
+	}}); err != nil {
+		return nil, err
+	}
+	if _, err := readRPC(ctx, dec, 1); err != nil {
+		return nil, err
+	}
+	if err := enc.Encode(mcpRPC{JSONRPC: "2.0", Method: "notifications/initialized", Params: map[string]any{}}); err != nil {
+		return nil, err
+	}
+	if err := enc.Encode(mcpRPC{JSONRPC: "2.0", ID: int64(2), Method: "tools/list", Params: map[string]any{}}); err != nil {
+		return nil, err
+	}
+	out, err := readRPC(ctx, dec, 2)
+	if err != nil {
+		return nil, err
+	}
+	return extractMCPTools(out)
+}
+
+func probeHTTP(ctx context.Context, s MCPService) ([]map[string]any, error) {
+	if s.URL == "" {
+		return nil, errors.New("MCP HTTP URL is empty")
+	}
+	client := &http.Client{Timeout: 0}
+	var session string
+	post := func(req mcpRPC) (any, error) {
+		b, _ := json.Marshal(req)
+		hr, err := http.NewRequestWithContext(ctx, http.MethodPost, s.URL, bytes.NewReader(b))
+		if err != nil {
+			return nil, err
+		}
+		hr.Header.Set("Content-Type", "application/json")
+		hr.Header.Set("Accept", "application/json, text/event-stream")
+		for k, v := range s.Headers {
+			hr.Header.Set(k, v)
+		}
+		if session != "" {
+			hr.Header.Set("Mcp-Session-Id", session)
+		}
+		resp, err := client.Do(hr)
+		if err != nil {
+			return nil, err
+		}
+		defer resp.Body.Close()
+		if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+			return nil, fmt.Errorf("MCP HTTP status %s", resp.Status)
+		}
+		if v := resp.Header.Get("Mcp-Session-Id"); v != "" {
+			session = v
+		}
+		data, err := io.ReadAll(io.LimitReader(resp.Body, 4*1024*1024))
+		if err != nil {
+			return nil, err
+		}
+		if len(bytes.TrimSpace(data)) == 0 {
+			return nil, nil
+		}
+		if bytes.Contains(data, []byte("event:")) {
+			data = parseSSEData(data)
+		}
+		var rr mcpRPCResponse
+		if err := json.Unmarshal(bytes.TrimSpace(data), &rr); err != nil {
+			return nil, err
+		}
+		if rr.Error != nil {
+			return nil, fmt.Errorf("MCP error %d: %s", rr.Error.Code, rr.Error.Message)
+		}
+		var out any
+		if len(rr.Result) > 0 {
+			if err := json.Unmarshal(rr.Result, &out); err != nil {
+				return nil, err
+			}
+		}
+		return out, nil
+	}
+	if strings.EqualFold(s.Transport, "sse") {
+		endpoint, err := discoverSSEEndpoint(ctx, s)
+		if err != nil {
+			return nil, err
+		}
+		s.URL = endpoint
+	}
+	if _, err := post(mcpRPC{JSONRPC: "2.0", ID: int64(1), Method: "initialize", Params: map[string]any{
+		"protocolVersion": "2024-11-05", "capabilities": map[string]any{},
+		"clientInfo": map[string]string{"name": "vps-commander-probe", "version": "1.0"},
+	}}); err != nil {
+		return nil, err
+	}
+	if _, err := post(mcpRPC{JSONRPC: "2.0", Method: "notifications/initialized", Params: map[string]any{}}); err != nil {
+		return nil, err
+	}
+	out, err := post(mcpRPC{JSONRPC: "2.0", ID: int64(2), Method: "tools/list", Params: map[string]any{}})
+	if err != nil {
+		return nil, err
+	}
+	return extractMCPTools(out)
+}
+
+func extractMCPTools(v any) ([]map[string]any, error) {
+	obj, ok := v.(map[string]any)
+	if !ok {
+		return nil, errors.New("invalid tools/list response")
+	}
+	raw, ok := obj["tools"].([]any)
+	if !ok {
+		return []map[string]any{}, nil
+	}
+	out := make([]map[string]any, 0, len(raw))
+	for _, item := range raw {
+		if m, ok := item.(map[string]any); ok {
+			tool := map[string]any{"name": m["name"], "description": m["description"]}
+			out = append(out, tool)
+		}
+	}
+	return out, nil
+}
+
 func callStdio(ctx context.Context, s MCPService, tool string, args any) (any, error) {
 	if s.Command == "" {
 		return nil, errors.New("MCP stdio command is empty")