From fb7e9ea271970421ad3fa3cfa2c86193d2f1f2c8 Mon Sep 17 00:00:00 2001 From: Jae Chen Date: Tue, 16 Jun 2026 07:26:45 +0800 Subject: [PATCH 1/4] feat(mcp): add MCP server skeleton with Bearer Token auth and get_proxy_status tool Implements Slice 1 of the MCP plugin: Streamable HTTP endpoint at POST /mcp with JSON-RPC 2.0 protocol, Bearer Token authentication (auto-generated random token persisted to ~/.vibeguard/mcp_token), and the get_proxy_status tool that returns proxy runtime stats and configuration. - JSON-RPC 2.0: initialize, tools/list, tools/call with proper error codes - Auth middleware: 401 for missing/invalid Bearer Token - Route wiring in proxy: POST /mcp intercepted before proxy fallback - 12 tests covering protocol, auth, and tool behaviors --- internal/admin/admin.go | 8 + internal/mcp/auth.go | 105 +++++++++++ internal/mcp/server.go | 207 +++++++++++++++++++++ internal/mcp/server_test.go | 359 ++++++++++++++++++++++++++++++++++++ internal/mcp/tools.go | 57 ++++++ internal/proxy/proxy.go | 9 + 6 files changed, 745 insertions(+) create mode 100644 internal/mcp/auth.go create mode 100644 internal/mcp/server.go create mode 100644 internal/mcp/server_test.go create mode 100644 internal/mcp/tools.go diff --git a/internal/admin/admin.go b/internal/admin/admin.go index 92e8a2c..8c650e7 100644 --- a/internal/admin/admin.go +++ b/internal/admin/admin.go @@ -72,6 +72,14 @@ func (a *Admin) SetStartTime(unix int64) { a.started.Store(unix) } +// StartedTime returns the Unix timestamp when the proxy started +func (a *Admin) StartedTime() int64 { + if a == nil { + return 0 + } + return a.started.Load() +} + // RecordAudit records one audit event about whether redaction rules were hit. func (a *Admin) RecordAudit(ev AuditEvent) AuditEvent { if a == nil || a.audit == nil { diff --git a/internal/mcp/auth.go b/internal/mcp/auth.go new file mode 100644 index 0000000..cfa0fee --- /dev/null +++ b/internal/mcp/auth.go @@ -0,0 +1,105 @@ +package mcp + +import ( + "crypto/rand" + "encoding/base64" + "net/http" + "os" + "path/filepath" + "strings" + "sync" + + "github.com/inkdust2021/vibeguard/internal/config" +) + +const mcpTokenFile = "mcp_token" + +var ( + mcpTokenOnce sync.Once + mcpTokenVal string + mcpTokenErr error +) + +// AuthMiddleware 包装 MCP handler，进行 Bearer Token 认证 +func AuthMiddleware(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + token, err := getOrCreateMcpToken() + if err != nil { + http.Error(w, "MCP auth not available", http.StatusServiceUnavailable) + return + } + + bearer := extractBearerToken(r) + if bearer == "" || bearer != token { + http.Error(w, "Unauthorized", http.StatusUnauthorized) + return + } + + next.ServeHTTP(w, r) + }) +} + +func extractBearerToken(r *http.Request) string { + h := strings.TrimSpace(r.Header.Get("Authorization")) + const prefix = "Bearer " + if !strings.HasPrefix(h, prefix) { + return "" + } + return strings.TrimSpace(h[len(prefix):]) +} + +// getOrCreateMcpToken 返回 MCP token，如果不存在则自动生成 +func getOrCreateMcpToken() (string, error) { + mcpTokenOnce.Do(func() { + mcpTokenVal, mcpTokenErr = loadOrGenerateMcpToken() + }) + return mcpTokenVal, mcpTokenErr +} + +func mcpTokenPath() string { + return filepath.Join(config.GetConfigDir(), mcpTokenFile) +} + +func loadOrGenerateMcpToken() (string, error) { + path := mcpTokenPath() + + // 尝试读取已有 token + data, err := os.ReadFile(path) + if err == nil { + tok := strings.TrimSpace(string(data)) + if tok != "" { + return tok, nil + } + } + + // 不存在则生成新的随机 token + tok, err := generateRandomToken(32) + if err != nil { + return "", err + } + + // 写入文件（权限 0600） + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + return "", err + } + if err := os.WriteFile(path, []byte(tok+"\n"), 0o600); err != nil { + return "", err + } + + return tok, nil +} + +func generateRandomToken(n int) (string, error) { + b := make([]byte, n) + if _, err := rand.Read(b); err != nil { + return "", err + } + return base64.RawURLEncoding.EncodeToString(b), nil +} + +// ResetMcpTokenForTesting 重置 token（仅用于测试） +func ResetMcpTokenForTesting() { + mcpTokenOnce = sync.Once{} + mcpTokenVal = "" + mcpTokenErr = nil +} diff --git a/internal/mcp/server.go b/internal/mcp/server.go new file mode 100644 index 0000000..a62915e --- /dev/null +++ b/internal/mcp/server.go @@ -0,0 +1,207 @@ +package mcp + +import ( + "encoding/json" + "fmt" + "log/slog" + "net/http" + "sync" + + "github.com/inkdust2021/vibeguard/internal/admin" + "github.com/inkdust2021/vibeguard/internal/config" +) + +// JSONRPCRequest 表示 MCP JSON-RPC 请求 +type JSONRPCRequest struct { + JSONRPC string `json:"jsonrpc"` + Method string `json:"method"` + Params json.RawMessage `json:"params,omitempty"` + ID any `json:"id,omitempty"` +} + +// JSONRPCResponse 表示 MCP JSON-RPC 响应 +type JSONRPCResponse struct { + JSONRPC string `json:"jsonrpc"` + ID any `json:"id,omitempty"` + Result json.RawMessage `json:"result,omitempty"` + Error *JSONRPCError `json:"error,omitempty"` +} + +// JSONRPCError 表示 JSON-RPC 错误 +type JSONRPCError struct { + Code int `json:"code"` + Message string `json:"message"` + Data any `json:"data,omitempty"` +} + +func (e *JSONRPCError) Error() string { + return fmt.Sprintf("JSON-RPC error %d: %s", e.Code, e.Message) +} + +// 标准 JSON-RPC 错误码 +const ( + ErrParseError = -32700 + ErrInvalidRequest = -32600 + ErrMethodNotFound = -32601 + ErrInvalidParams = -32602 + ErrInternalError = -32603 +) + +// Server 是 VibeGuard MCP Server +type Server struct { + mu sync.RWMutex + version string + name string + cfg *config.Manager + admin *admin.Admin + tools map[string]ToolHandler +} + +// ToolHandler 是一个 MCP 工具处理函数 +type ToolHandler func(params json.RawMessage) (any, error) + +// NewServer 创建 MCP Server +func NewServer(cfg *config.Manager, adm *admin.Admin, version string) *Server { + s := &Server{ + version: version, + name: "vibeguard", + cfg: cfg, + admin: adm, + tools: make(map[string]ToolHandler), + } + s.registerTools() + return s +} + +// Handler 返回 HTTP handler +func (s *Server) Handler() http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + + var req JSONRPCRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + s.writeError(w, nil, ErrParseError, "Parse error", err.Error()) + return + } + + switch req.Method { + case "initialize": + s.handleInitialize(w, &req) + case "tools/list": + s.handleToolsList(w, &req) + case "tools/call": + s.handleToolsCall(w, &req) + default: + s.writeError(w, req.ID, ErrMethodNotFound, "Method not found", req.Method) + } + }) +} + +func (s *Server) handleInitialize(w http.ResponseWriter, req *JSONRPCRequest) { + result := map[string]any{ + "protocolVersion": "2024-11-05", + "capabilities": map[string]any{ + "tools": map[string]any{}, + }, + "serverInfo": map[string]any{ + "name": s.name, + "version": s.version, + }, + } + s.writeResult(w, req.ID, result) +} + +func (s *Server) handleToolsList(w http.ResponseWriter, req *JSONRPCRequest) { + tools := s.listTools() + result := map[string]any{ + "tools": tools, + } + s.writeResult(w, req.ID, result) +} + +func (s *Server) handleToolsCall(w http.ResponseWriter, req *JSONRPCRequest) { + var call struct { + Name string `json:"name"` + Arguments json.RawMessage `json:"arguments"` + } + if err := json.Unmarshal(req.Params, &call); err != nil { + s.writeError(w, req.ID, ErrInvalidParams, "Invalid params", err.Error()) + return + } + if call.Name == "" { + s.writeError(w, req.ID, ErrInvalidParams, "Invalid params", "missing tool name") + return + } + + s.mu.RLock() + handler, ok := s.tools[call.Name] + s.mu.RUnlock() + + if !ok { + s.writeError(w, req.ID, ErrMethodNotFound, "Tool not found", call.Name) + return + } + + result, err := handler(call.Arguments) + if err != nil { + if jerr, ok := err.(*JSONRPCError); ok { + s.writeError(w, req.ID, jerr.Code, jerr.Message, jerr.Data) + return + } + s.writeError(w, req.ID, ErrInternalError, "Internal error", err.Error()) + return + } + + s.writeResult(w, req.ID, map[string]any{ + "content": []map[string]any{ + { + "type": "text", + "text": mustJSON(result), + }, + }, + }) +} + +func (s *Server) writeResult(w http.ResponseWriter, id any, result any) { + w.Header().Set("Content-Type", "application/json") + resp := JSONRPCResponse{ + JSONRPC: "2.0", + ID: id, + } + data, err := json.Marshal(result) + if err != nil { + s.writeError(w, id, ErrInternalError, "Failed to marshal result", err.Error()) + return + } + resp.Result = data + if err := json.NewEncoder(w).Encode(resp); err != nil { + slog.Warn("MCP response write failed", "error", err) + } +} + +func (s *Server) writeError(w http.ResponseWriter, id any, code int, message string, data any) { + w.Header().Set("Content-Type", "application/json") + resp := JSONRPCResponse{ + JSONRPC: "2.0", + ID: id, + Error: &JSONRPCError{ + Code: code, + Message: message, + Data: data, + }, + } + if err := json.NewEncoder(w).Encode(resp); err != nil { + slog.Warn("MCP error response write failed", "error", err) + } +} + +func mustJSON(v any) string { + b, err := json.Marshal(v) + if err != nil { + return fmt.Sprintf(`{"error":%q}`, err.Error()) + } + return string(b) +} diff --git a/internal/mcp/server_test.go b/internal/mcp/server_test.go new file mode 100644 index 0000000..4b1023e --- /dev/null +++ b/internal/mcp/server_test.go @@ -0,0 +1,359 @@ +package mcp + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/inkdust2021/vibeguard/internal/admin" + "github.com/inkdust2021/vibeguard/internal/cert" + "github.com/inkdust2021/vibeguard/internal/config" + "github.com/inkdust2021/vibeguard/internal/session" +) + +func TestMcpServer_Initialize(t *testing.T) { + ResetMcpTokenForTesting() + srv := newTestServer(t) + token, _ := getOrCreateMcpToken() + + reqBody := map[string]any{ + "jsonrpc": "2.0", + "method": "initialize", + "id": 1, + } + resp := postJSON(t, srv, reqBody, token) + + if resp["jsonrpc"] != "2.0" { + t.Errorf("expected jsonrpc=2.0, got %v", resp["jsonrpc"]) + } + if resp["id"] != float64(1) { + t.Errorf("expected id=1, got %v", resp["id"]) + } + result, ok := resp["result"].(map[string]any) + if !ok { + t.Fatalf("expected result object, got %T", resp["result"]) + } + if result["protocolVersion"] != "2024-11-05" { + t.Errorf("expected protocolVersion=2024-11-05, got %v", result["protocolVersion"]) + } + capabilities, ok := result["capabilities"].(map[string]any) + if !ok { + t.Fatalf("expected capabilities object") + } + if _, hasTools := capabilities["tools"]; !hasTools { + t.Error("expected capabilities to include tools") + } + serverInfo, ok := result["serverInfo"].(map[string]any) + if !ok { + t.Fatalf("expected serverInfo object") + } + if serverInfo["name"] != "vibeguard" { + t.Errorf("expected name=vibeguard, got %v", serverInfo["name"]) + } + if serverInfo["version"] != "test-version" { + t.Errorf("expected version=test-version, got %v", serverInfo["version"]) + } +} + +func TestMcpServer_ToolsList(t *testing.T) { + ResetMcpTokenForTesting() + srv := newTestServer(t) + token, _ := getOrCreateMcpToken() + + reqBody := map[string]any{ + "jsonrpc": "2.0", + "method": "tools/list", + "id": 2, + } + resp := postJSON(t, srv, reqBody, token) + + result, ok := resp["result"].(map[string]any) + if !ok { + t.Fatalf("expected result object") + } + tools, ok := result["tools"].([]any) + if !ok { + t.Fatalf("expected tools array") + } + if len(tools) == 0 { + t.Error("expected at least one tool") + } + + first := tools[0].(map[string]any) + if first["name"] != "get_proxy_status" { + t.Errorf("expected first tool=get_proxy_status, got %v", first["name"]) + } +} + +func TestMcpServer_GetProxyStatus(t *testing.T) { + ResetMcpTokenForTesting() + srv := newTestServer(t) + token, _ := getOrCreateMcpToken() + + reqBody := map[string]any{ + "jsonrpc": "2.0", + "method": "tools/call", + "params": map[string]any{ + "name": "get_proxy_status", + "arguments": map[string]any{}, + }, + "id": 3, + } + resp := postJSON(t, srv, reqBody, token) + + if _, hasError := resp["error"]; hasError { + t.Fatalf("unexpected error: %v", resp["error"]) + } + result, ok := resp["result"].(map[string]any) + if !ok { + t.Fatalf("expected result object") + } + content, ok := result["content"].([]any) + if !ok || len(content) == 0 { + t.Fatalf("expected content array") + } + textItem := content[0].(map[string]any) + if textItem["type"] != "text" { + t.Errorf("expected type=text, got %v", textItem["type"]) + } + + // 验证 text 内容中的 JSON 结构 + var status map[string]any + if err := json.Unmarshal([]byte(textItem["text"].(string)), &status); err != nil { + t.Fatalf("failed to parse status JSON: %v", err) + } + requiredFields := []string{"version", "started_at", "uptime", "total_requests", "redacted_count", "restored_count", "errors", "intercept_mode", "listen_address"} + for _, f := range requiredFields { + if _, ok := status[f]; !ok { + t.Errorf("missing field: %s", f) + } + } + if status["version"] != "test-version" { + t.Errorf("expected version=test-version, got %v", status["version"]) + } +} + +func TestMcpServer_InvalidJSON(t *testing.T) { + ResetMcpTokenForTesting() + srv := newTestServer(t) + token, _ := getOrCreateMcpToken() + + r := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader([]byte("not json"))) + r.Header.Set("Content-Type", "application/json") + r.Header.Set("Authorization", "Bearer "+token) + w := httptest.NewRecorder() + AuthMiddleware(srv.Handler()).ServeHTTP(w, r) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } + var resp map[string]any + json.Unmarshal(w.Body.Bytes(), &resp) + errObj, _ := resp["error"].(map[string]any) + if errObj["code"] != float64(-32700) { + t.Errorf("expected error code -32700, got %v", errObj["code"]) + } +} + +func TestMcpServer_MethodNotAllowed(t *testing.T) { + ResetMcpTokenForTesting() + srv := newTestServer(t) + + r := httptest.NewRequest(http.MethodGet, "/mcp", nil) + w := httptest.NewRecorder() + srv.Handler().ServeHTTP(w, r) + + if w.Code != http.StatusMethodNotAllowed { + t.Errorf("expected 405 for GET, got %d", w.Code) + } +} + +func TestMcpServer_EmptyMethod(t *testing.T) { + ResetMcpTokenForTesting() + srv := newTestServer(t) + token, _ := getOrCreateMcpToken() + + reqBody := map[string]any{ + "jsonrpc": "2.0", + "id": 1, + } + resp := postJSON(t, srv, reqBody, token) + + errObj, _ := resp["error"].(map[string]any) + if errObj["code"] != float64(-32601) { + t.Errorf("expected error code -32601 (method not found), got %v", errObj["code"]) + } +} + +func TestMcpServer_UnknownMethod(t *testing.T) { + ResetMcpTokenForTesting() + srv := newTestServer(t) + token, _ := getOrCreateMcpToken() + + reqBody := map[string]any{ + "jsonrpc": "2.0", + "method": "nonexistent/method", + "id": 1, + } + resp := postJSON(t, srv, reqBody, token) + + errObj, _ := resp["error"].(map[string]any) + if errObj["code"] != float64(-32601) { + t.Errorf("expected error code -32601, got %v", errObj["code"]) + } +} + +func TestMcpServer_ToolNotFound(t *testing.T) { + ResetMcpTokenForTesting() + srv := newTestServer(t) + token, _ := getOrCreateMcpToken() + + reqBody := map[string]any{ + "jsonrpc": "2.0", + "method": "tools/call", + "params": map[string]any{ + "name": "nonexistent_tool", + "arguments": map[string]any{}, + }, + "id": 1, + } + resp := postJSON(t, srv, reqBody, token) + + errObj, _ := resp["error"].(map[string]any) + if errObj["code"] != float64(-32601) { + t.Errorf("expected error code -32601 (tool not found), got %v", errObj["code"]) + } +} + +func TestMcpServer_ToolCallMissingName(t *testing.T) { + ResetMcpTokenForTesting() + srv := newTestServer(t) + token, _ := getOrCreateMcpToken() + + reqBody := map[string]any{ + "jsonrpc": "2.0", + "method": "tools/call", + "params": map[string]any{ + "arguments": map[string]any{}, + }, + "id": 1, + } + resp := postJSON(t, srv, reqBody, token) + + errObj, _ := resp["error"].(map[string]any) + if errObj["code"] != float64(-32602) { + t.Errorf("expected error code -32602 (invalid params), got %v", errObj["code"]) + } +} + +func TestMcpServer_ValidTokenMultipleRequests(t *testing.T) { + ResetMcpTokenForTesting() + srv := newTestServer(t) + token, _ := getOrCreateMcpToken() + + reqBody := map[string]any{ + "jsonrpc": "2.0", + "method": "initialize", + "id": 1, + } + + for i := 0; i < 3; i++ { + resp := postJSON(t, srv, reqBody, token) + if _, hasError := resp["error"]; hasError { + t.Errorf("request %d: unexpected error: %v", i, resp["error"]) + } + if resp["jsonrpc"] != "2.0" { + t.Errorf("request %d: expected jsonrpc=2.0, got %v", i, resp["jsonrpc"]) + } + } +} + +func TestMcpServer_AuthRequired(t *testing.T) { + ResetMcpTokenForTesting() + srv := newTestServer(t) + + reqBody := map[string]any{ + "jsonrpc": "2.0", + "method": "initialize", + "id": 1, + } + + body, _ := json.Marshal(reqBody) + r := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(body)) + r.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + AuthMiddleware(srv.Handler()).ServeHTTP(w, r) + if w.Code != http.StatusUnauthorized { + t.Errorf("expected 401 without token, got %d", w.Code) + } +} + +func TestMcpServer_InvalidToken(t *testing.T) { + ResetMcpTokenForTesting() + srv := newTestServer(t) + + reqBody := map[string]any{ + "jsonrpc": "2.0", + "method": "initialize", + "id": 1, + } + + body, _ := json.Marshal(reqBody) + r := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(body)) + r.Header.Set("Content-Type", "application/json") + r.Header.Set("Authorization", "Bearer invalid-token") + w := httptest.NewRecorder() + + AuthMiddleware(srv.Handler()).ServeHTTP(w, r) + if w.Code != http.StatusUnauthorized { + t.Errorf("expected 401 with invalid token, got %d", w.Code) + } +} + +func newTestServer(t *testing.T) *Server { + t.Helper() + dir := t.TempDir() + t.Setenv("HOME", dir) + + cfg := config.NewManager() + caCertPath := dir + "/ca-cert.pem" + caKeyPath := dir + "/ca-key.pem" + ca, err := cert.LoadOrGenerateCA(caCertPath, caKeyPath) + if err != nil { + t.Fatalf("failed to generate test CA: %v", err) + } + sess := session.NewManager(0, 1000) + adm := admin.New(cfg, sess, ca, "", "") + + return NewServer(cfg, adm, "test-version") +} + +func postJSON(t *testing.T, srv *Server, reqBody map[string]any, token string) map[string]any { + t.Helper() + body, _ := json.Marshal(reqBody) + r := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(body)) + r.Header.Set("Content-Type", "application/json") + if token != "" { + r.Header.Set("Authorization", "Bearer "+token) + } + w := httptest.NewRecorder() + + handler := AuthMiddleware(srv.Handler()) + handler.ServeHTTP(w, r) + + if w.Code != http.StatusOK { + if w.Code == http.StatusUnauthorized || w.Code == http.StatusForbidden { + return map[string]any{"_status": w.Code} + } + t.Fatalf("unexpected status: %d, body: %s", w.Code, w.Body.String()) + } + + var resp map[string]any + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatalf("failed to decode response: %v, body: %s", err, w.Body.String()) + } + return resp +} diff --git a/internal/mcp/tools.go b/internal/mcp/tools.go new file mode 100644 index 0000000..dba9533 --- /dev/null +++ b/internal/mcp/tools.go @@ -0,0 +1,57 @@ +package mcp + +import ( + "encoding/json" + "time" +) + +// Tool 表示一个 MCP 工具定义 +type Tool struct { + Name string `json:"name"` + Description string `json:"description"` + InputSchema any `json:"inputSchema"` +} + +func (s *Server) registerTools() { + s.mu.Lock() + defer s.mu.Unlock() + + s.tools["get_proxy_status"] = s.handleGetProxyStatus +} + +func (s *Server) listTools() []Tool { + return []Tool{ + { + Name: "get_proxy_status", + Description: "获取 VibeGuard 代理的运行状态，包括启动时间、请求统计、拦截模式等", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{}, + }, + }, + } +} + +func (s *Server) handleGetProxyStatus(_ json.RawMessage) (any, error) { + stats := s.admin.GetStats() + cfg := s.cfg.Get() + + started := s.admin.StartedTime() + uptime := "" + if started > 0 { + uptime = time.Since(time.Unix(started, 0)).Round(time.Second).String() + } + + return map[string]any{ + "version": s.version, + "started_at": time.Unix(started, 0).Format(time.RFC3339), + "uptime": uptime, + "total_requests": stats.TotalRequests.Load(), + "redacted_count": stats.RedactedRequests.Load(), + "restored_count": stats.RestoredRequests.Load(), + "errors": stats.Errors.Load(), + "intercept_mode": cfg.Proxy.InterceptMode, + "websocket_beta": cfg.Proxy.WebSocketRedactionBeta, + "listen_address": cfg.Proxy.Listen, + }, nil +} diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index a5f8f04..f1fe2ee 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -38,7 +38,9 @@ import ( "github.com/inkdust2021/vibeguard/internal/rulelists" "github.com/inkdust2021/vibeguard/internal/secretsources" "github.com/inkdust2021/vibeguard/internal/session" + "github.com/inkdust2021/vibeguard/internal/mcp" "github.com/inkdust2021/vibeguard/internal/stream" + "github.com/inkdust2021/vibeguard/internal/version" "github.com/inkdust2021/vibeguard/internal/wsproxy" "github.com/inkdust2021/vibeguard/internal/zstd" ) @@ -168,6 +170,9 @@ func (s *Server) Start() error { // ServeMux may return a 301 redirect, breaking HTTPS proxy traffic (clients keep reconnecting). // Use a custom router here: only /manager/ goes to the admin UI; everything else goes to the proxy (including CONNECT). adminHandler := s.admin.Handler() + mcpServer := mcp.NewServer(s.config, s.admin, version.Version) + mcpHandler := mcp.AuthMiddleware(mcpServer.Handler()) + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r != nil { if r.URL.Path == "/manager" { @@ -178,6 +183,10 @@ func (s *Server) Start() error { adminHandler.ServeHTTP(w, r) return } + if r.URL.Path == "/mcp" { + mcpHandler.ServeHTTP(w, r) + return + } } s.proxy.ServeHTTP(w, r) }) From 8dae65d69f7418689000ffc8eab212002b29e35e Mon Sep 17 00:00:00 2001 From: Jae Chen Date: Tue, 16 Jun 2026 23:40:51 +0800 Subject: [PATCH 2/4] feat(mcp): add Detect, rule query tools, and CLI mcp setup wizard MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 实现 Slice 2/3/4，覆盖 jayeliu/VibeGuard#2, #3, #4： Slice 2 (Engine.Detect 只读检测 + detect_sensitive / preview_redacted): - 新增 Engine.Detect 方法：扫描敏感数据但不生成占位符、不注册到 session - detect_sensitive 工具：默认返回分类/位置，include_originals=true 时返回原始值 - preview_redacted 工具：返回带临时占位符的脱敏文本预览 - proxy.Server 暴露 GetRedactEngine() 供 MCP 访问 Slice 3 (规则查询工具): - list_keywords 工具：返回当前加载的关键词规则（user/builtin/engine 来源） - list_rule_lists 工具：返回规则列表元数据（id/name/enabled/source） - Engine 新增 ListKeywords() 和 ListRegexCategories() 方法 Slice 4 (CLI mcp setup 向导): - vibeguard mcp setup 命令：生成 Claude Code 兼容的 MCP 配置 - --format=json (默认) 输出完整 mcpServers JSON 配置 - --format=env 输出可 source 的环境变量 - token 不存在时给出友好提示 测试： - internal/redact/engine_test.go: 6 个测试覆盖 Detect 只读性 - internal/mcp/server_test.go: 新增 7 个测试覆盖 5 个新工具 - 全量回归测试通过 --- cmd/vibeguard/main.go | 78 ++++++++++ internal/mcp/server.go | 17 ++- internal/mcp/server_test.go | 272 ++++++++++++++++++++++++++++++++- internal/mcp/tools.go | 201 ++++++++++++++++++++++++ internal/proxy/proxy.go | 16 +- internal/redact/engine.go | 96 ++++++++++++ internal/redact/engine_test.go | 94 ++++++++++++ 7 files changed, 764 insertions(+), 10 deletions(-) create mode 100644 internal/redact/engine_test.go diff --git a/cmd/vibeguard/main.go b/cmd/vibeguard/main.go index 57102c1..9efccf3 100644 --- a/cmd/vibeguard/main.go +++ b/cmd/vibeguard/main.go @@ -2,6 +2,7 @@ package main import ( "bufio" + "encoding/json" "errors" "fmt" "io" @@ -187,6 +188,22 @@ Pattern is treated as a keyword (exact substring match).`, RunE: runTest, } +var mcpFormat string + +var mcpCmd = &cobra.Command{ + Use: "mcp", + Short: "MCP server configuration", + Long: `Manage MCP server configuration for Claude Code integration.`, +} + +var mcpSetupCmd = &cobra.Command{ + Use: "setup", + Short: "Generate MCP configuration for Claude Code", + Long: `Generate MCP server configuration for Claude Code integration. +Reads the MCP token from the running VibeGuard instance.`, + RunE: runMcpSetup, +} + var versionCmd = &cobra.Command{ Use: "version", Short: "Print version information", @@ -212,6 +229,9 @@ func init() { rootCmd.AddCommand(initCmd) rootCmd.AddCommand(trustCmd) rootCmd.AddCommand(testCmd) + rootCmd.AddCommand(mcpCmd) + mcpCmd.AddCommand(mcpSetupCmd) + mcpSetupCmd.Flags().StringVar(&mcpFormat, "format", "json", "output format: json or env") rootCmd.AddCommand(versionCmd) trustCmd.Flags().StringVar(&trustMode, "mode", string(cert.TrustInstallModeSystem), "trust store mode: system|user|auto") @@ -1221,3 +1241,61 @@ func runTest(cmd *cobra.Command, args []string) error { return nil } + + +func runMcpSetup(cmd *cobra.Command, args []string) error { + lang := uiLang() + + // 读取 MCP token + token, err := readMcpToken() + if err != nil { + fmt.Fprintln(os.Stderr, uiText(lang, "错误: 无法读取 MCP token，请确保 VibeGuard 已启动", "Error: cannot read MCP token, make sure VibeGuard is running")) + return fmt.Errorf("read MCP token: %w", err) + } + + // 读取代理监听地址 + cfg := config.NewManager() + if err := cfg.Load(); err != nil { + // Load 失败时使用默认值 + _ = err + } + c := cfg.Get() + listenAddr := c.Proxy.Listen + if listenAddr == "" { + listenAddr = "127.0.0.1:28657" + } + + mcpURL := fmt.Sprintf("http://%s/mcp", listenAddr) + + switch mcpFormat { + case "env": + fmt.Printf("VG_MCP_URL=%s\n", mcpURL) + fmt.Printf("VG_MCP_TOKEN=%s\n", token) + case "json": + fallthrough + default: + config := map[string]any{ + "mcpServers": map[string]any{ + "vibeguard": map[string]any{ + "url": mcpURL, + "headers": map[string]any{ + "Authorization": fmt.Sprintf("Bearer %s", token), + }, + }, + }, + } + data, _ := json.MarshalIndent(config, "", " ") + fmt.Println(string(data)) + } + + return nil +} + +func readMcpToken() (string, error) { + path := filepath.Join(config.GetConfigDir(), "mcp_token") + data, err := os.ReadFile(path) + if err != nil { + return "", err + } + return strings.TrimSpace(string(data)), nil +} diff --git a/internal/mcp/server.go b/internal/mcp/server.go index a62915e..36335ba 100644 --- a/internal/mcp/server.go +++ b/internal/mcp/server.go @@ -9,6 +9,7 @@ import ( "github.com/inkdust2021/vibeguard/internal/admin" "github.com/inkdust2021/vibeguard/internal/config" + "github.com/inkdust2021/vibeguard/internal/redact" ) // JSONRPCRequest 表示 MCP JSON-RPC 请求 @@ -49,24 +50,26 @@ const ( // Server 是 VibeGuard MCP Server type Server struct { - mu sync.RWMutex - version string - name string - cfg *config.Manager - admin *admin.Admin - tools map[string]ToolHandler + mu sync.RWMutex + version string + name string + cfg *config.Manager + admin *admin.Admin + engine *redact.Engine + tools map[string]ToolHandler } // ToolHandler 是一个 MCP 工具处理函数 type ToolHandler func(params json.RawMessage) (any, error) // NewServer 创建 MCP Server -func NewServer(cfg *config.Manager, adm *admin.Admin, version string) *Server { +func NewServer(cfg *config.Manager, adm *admin.Admin, eng *redact.Engine, version string) *Server { s := &Server{ version: version, name: "vibeguard", cfg: cfg, admin: adm, + engine: eng, tools: make(map[string]ToolHandler), } s.registerTools() diff --git a/internal/mcp/server_test.go b/internal/mcp/server_test.go index 4b1023e..34edb62 100644 --- a/internal/mcp/server_test.go +++ b/internal/mcp/server_test.go @@ -10,6 +10,7 @@ import ( "github.com/inkdust2021/vibeguard/internal/admin" "github.com/inkdust2021/vibeguard/internal/cert" "github.com/inkdust2021/vibeguard/internal/config" + "github.com/inkdust2021/vibeguard/internal/redact" "github.com/inkdust2021/vibeguard/internal/session" ) @@ -248,6 +249,265 @@ func TestMcpServer_ToolCallMissingName(t *testing.T) { } } +func TestMcpServer_DetectSensitive(t *testing.T) { + ResetMcpTokenForTesting() + srv := newTestServerWithKeywords(t, map[string]string{"secret": "TEXT"}) + token, _ := getOrCreateMcpToken() + + reqBody := map[string]any{ + "jsonrpc": "2.0", + "method": "tools/call", + "params": map[string]any{ + "name": "detect_sensitive", + "arguments": map[string]any{ + "text": "this is a secret message", + }, + }, + "id": 1, + } + resp := postJSON(t, srv, reqBody, token) + + if _, hasError := resp["error"]; hasError { + t.Fatalf("unexpected error: %v", resp["error"]) + } + result := resp["result"].(map[string]any) + content := result["content"].([]any) + textItem := content[0].(map[string]any) + + var detectResult map[string]any + json.Unmarshal([]byte(textItem["text"].(string)), &detectResult) + + if detectResult["has_sensitive"] != true { + t.Errorf("expected has_sensitive=true, got %v", detectResult["has_sensitive"]) + } + matches := detectResult["matches"].([]any) + if len(matches) == 0 { + t.Fatal("expected at least one match") + } + hit := matches[0].(map[string]any) + if hit["category"] != "TEXT" { + t.Errorf("expected category=TEXT, got %v", hit["category"]) + } + if _, hasOriginal := hit["original"]; hasOriginal { + t.Error("expected no original field by default") + } +} + +func TestMcpServer_DetectSensitive_WithOriginals(t *testing.T) { + ResetMcpTokenForTesting() + srv := newTestServerWithKeywords(t, map[string]string{"secret": "TEXT"}) + token, _ := getOrCreateMcpToken() + + reqBody := map[string]any{ + "jsonrpc": "2.0", + "method": "tools/call", + "params": map[string]any{ + "name": "detect_sensitive", + "arguments": map[string]any{ + "text": "this is a secret message", + "include_originals": true, + }, + }, + "id": 1, + } + resp := postJSON(t, srv, reqBody, token) + + result := resp["result"].(map[string]any) + content := result["content"].([]any) + textItem := content[0].(map[string]any) + + var detectResult map[string]any + json.Unmarshal([]byte(textItem["text"].(string)), &detectResult) + + matches := detectResult["matches"].([]any) + hit := matches[0].(map[string]any) + if hit["original"] != "secret" { + t.Errorf("expected original=secret, got %v", hit["original"]) + } +} + +func TestMcpServer_DetectSensitive_NoMatch(t *testing.T) { + ResetMcpTokenForTesting() + srv := newTestServerWithKeywords(t, map[string]string{"secret": "TEXT"}) + token, _ := getOrCreateMcpToken() + + reqBody := map[string]any{ + "jsonrpc": "2.0", + "method": "tools/call", + "params": map[string]any{ + "name": "detect_sensitive", + "arguments": map[string]any{ + "text": "this is a clean message", + }, + }, + "id": 1, + } + resp := postJSON(t, srv, reqBody, token) + + result := resp["result"].(map[string]any) + content := result["content"].([]any) + textItem := content[0].(map[string]any) + + var detectResult map[string]any + json.Unmarshal([]byte(textItem["text"].(string)), &detectResult) + + if detectResult["has_sensitive"] != false { + t.Errorf("expected has_sensitive=false, got %v", detectResult["has_sensitive"]) + } +} + +func TestMcpServer_DetectSensitive_MissingText(t *testing.T) { + ResetMcpTokenForTesting() + srv := newTestServerWithKeywords(t, nil) + token, _ := getOrCreateMcpToken() + + reqBody := map[string]any{ + "jsonrpc": "2.0", + "method": "tools/call", + "params": map[string]any{ + "name": "detect_sensitive", + "arguments": map[string]any{}, + }, + "id": 1, + } + resp := postJSON(t, srv, reqBody, token) + + errObj, _ := resp["error"].(map[string]any) + if errObj["code"] != float64(-32602) { + t.Errorf("expected error code -32602, got %v", errObj["code"]) + } +} + +func TestMcpServer_PreviewRedacted(t *testing.T) { + ResetMcpTokenForTesting() + srv := newTestServerWithKeywords(t, map[string]string{"secret": "TEXT"}) + token, _ := getOrCreateMcpToken() + + reqBody := map[string]any{ + "jsonrpc": "2.0", + "method": "tools/call", + "params": map[string]any{ + "name": "preview_redacted", + "arguments": map[string]any{ + "text": "this is a secret message", + }, + }, + "id": 1, + } + resp := postJSON(t, srv, reqBody, token) + + if _, hasError := resp["error"]; hasError { + t.Fatalf("unexpected error: %v", resp["error"]) + } + result := resp["result"].(map[string]any) + content := result["content"].([]any) + textItem := content[0].(map[string]any) + + var previewResult map[string]any + json.Unmarshal([]byte(textItem["text"].(string)), &previewResult) + + if previewResult["match_count"] != float64(1) { + t.Errorf("expected match_count=1, got %v", previewResult["match_count"]) + } + redactedText := previewResult["redacted_text"].(string) + if redactedText == "this is a secret message" { + t.Error("expected text to be redacted, got original") + } + if !contains(redactedText, "__VG_TEXT_PREVIEW__") { + t.Errorf("expected placeholder in redacted text, got %s", redactedText) + } +} + +func TestMcpServer_ListKeywords(t *testing.T) { + ResetMcpTokenForTesting() + srv := newTestServer(t) + token, _ := getOrCreateMcpToken() + + reqBody := map[string]any{ + "jsonrpc": "2.0", + "method": "tools/call", + "params": map[string]any{ + "name": "list_keywords", + "arguments": map[string]any{}, + }, + "id": 1, + } + resp := postJSON(t, srv, reqBody, token) + + if _, hasError := resp["error"]; hasError { + t.Fatalf("unexpected error: %v", resp["error"]) + } + result := resp["result"].(map[string]any) + content := result["content"].([]any) + textItem := content[0].(map[string]any) + + var listResult map[string]any + json.Unmarshal([]byte(textItem["text"].(string)), &listResult) + + keywords := listResult["keywords"].([]any) + if len(keywords) == 0 { + t.Error("expected at least one keyword") + } + first := keywords[0].(map[string]any) + if _, ok := first["keyword"]; !ok { + t.Error("expected keyword field") + } + if _, ok := first["category"]; !ok { + t.Error("expected category field") + } + if _, ok := first["source"]; !ok { + t.Error("expected source field") + } +} + +func TestMcpServer_ListRuleLists(t *testing.T) { + ResetMcpTokenForTesting() + srv := newTestServer(t) + token, _ := getOrCreateMcpToken() + + reqBody := map[string]any{ + "jsonrpc": "2.0", + "method": "tools/call", + "params": map[string]any{ + "name": "list_rule_lists", + "arguments": map[string]any{}, + }, + "id": 1, + } + resp := postJSON(t, srv, reqBody, token) + + if _, hasError := resp["error"]; hasError { + t.Fatalf("unexpected error: %v", resp["error"]) + } + result := resp["result"].(map[string]any) + content := result["content"].([]any) + textItem := content[0].(map[string]any) + + var listResult map[string]any + json.Unmarshal([]byte(textItem["text"].(string)), &listResult) + + // 默认配置没有 rule lists，但字段应该存在 + if _, ok := listResult["rule_lists"]; !ok { + t.Error("expected rule_lists field") + } + if _, ok := listResult["count"]; !ok { + t.Error("expected count field") + } +} + +func contains(s, substr string) bool { + return len(s) >= len(substr) && (s == substr || len(s) > 0 && containsSubstr(s, substr)) +} + +func containsSubstr(s, substr string) bool { + for i := 0; i <= len(s)-len(substr); i++ { + if s[i:i+len(substr)] == substr { + return true + } + } + return false +} + func TestMcpServer_ValidTokenMultipleRequests(t *testing.T) { ResetMcpTokenForTesting() srv := newTestServer(t) @@ -314,6 +574,11 @@ func TestMcpServer_InvalidToken(t *testing.T) { } func newTestServer(t *testing.T) *Server { + t.Helper() + return newTestServerWithKeywords(t, map[string]string{"secret": "TEXT"}) +} + +func newTestServerWithKeywords(t *testing.T, keywords map[string]string) *Server { t.Helper() dir := t.TempDir() t.Setenv("HOME", dir) @@ -328,7 +593,12 @@ func newTestServer(t *testing.T) *Server { sess := session.NewManager(0, 1000) adm := admin.New(cfg, sess, ca, "", "") - return NewServer(cfg, adm, "test-version") + eng := redact.NewEngine(sess, "__VG_") + for kw, cat := range keywords { + eng.AddKeyword(kw, cat) + } + + return NewServer(cfg, adm, eng, "test-version") } func postJSON(t *testing.T, srv *Server, reqBody map[string]any, token string) map[string]any { diff --git a/internal/mcp/tools.go b/internal/mcp/tools.go index dba9533..70dfae4 100644 --- a/internal/mcp/tools.go +++ b/internal/mcp/tools.go @@ -2,7 +2,12 @@ package mcp import ( "encoding/json" + "fmt" "time" + + "github.com/inkdust2021/vibeguard/internal/redact" + + "github.com/inkdust2021/vibeguard/internal/config" ) // Tool 表示一个 MCP 工具定义 @@ -17,6 +22,10 @@ func (s *Server) registerTools() { defer s.mu.Unlock() s.tools["get_proxy_status"] = s.handleGetProxyStatus + s.tools["detect_sensitive"] = s.handleDetectSensitive + s.tools["preview_redacted"] = s.handlePreviewRedacted + s.tools["list_keywords"] = s.handleListKeywords + s.tools["list_rule_lists"] = s.handleListRuleLists } func (s *Server) listTools() []Tool { @@ -29,6 +38,54 @@ func (s *Server) listTools() []Tool { "properties": map[string]any{}, }, }, + { + Name: "detect_sensitive", + Description: "检测文本中是否包含敏感信息，返回命中规则列表。默认不暴露原始值，设置 include_originals=true 可查看", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "text": map[string]any{ + "type": "string", + "description": "待检测的文本", + }, + "include_originals": map[string]any{ + "type": "boolean", + "description": "是否在结果中包含原始敏感值（默认 false）", + }, + }, + "required": []string{"text"}, + }, + }, + { + Name: "preview_redacted", + Description: "预览文本脱敏效果，返回带临时占位符的文本（占位符不持久化，不会注册到 session）", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "text": map[string]any{ + "type": "string", + "description": "待预览脱敏的文本", + }, + }, + "required": []string{"text"}, + }, + }, + { + Name: "list_keywords", + Description: "查询当前加载的关键词规则列表，每条规则包含匹配文本、分类和来源", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{}, + }, + }, + { + Name: "list_rule_lists", + Description: "查询当前启用的规则列表（rule lists）及其元数据", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{}, + }, + }, } } @@ -55,3 +112,147 @@ func (s *Server) handleGetProxyStatus(_ json.RawMessage) (any, error) { "listen_address": cfg.Proxy.Listen, }, nil } + +func (s *Server) handleDetectSensitive(params json.RawMessage) (any, error) { + var args struct { + Text string `json:"text"` + IncludeOriginals bool `json:"include_originals"` + } + if err := json.Unmarshal(params, &args); err != nil { + return nil, &JSONRPCError{Code: ErrInvalidParams, Message: "Invalid params", Data: err.Error()} + } + if args.Text == "" { + return nil, &JSONRPCError{Code: ErrInvalidParams, Message: "Invalid params", Data: "text is required"} + } + + if s.engine == nil { + return nil, &JSONRPCError{Code: -32001, Message: "Detection unavailable", Data: "NER/pipeline mode does not support Detect"} + } + + matches := s.engine.Detect([]byte(args.Text)) + return buildDetectResult(matches, args.IncludeOriginals), nil +} + +func (s *Server) handlePreviewRedacted(params json.RawMessage) (any, error) { + var args struct { + Text string `json:"text"` + } + if err := json.Unmarshal(params, &args); err != nil { + return nil, &JSONRPCError{Code: ErrInvalidParams, Message: "Invalid params", Data: err.Error()} + } + if args.Text == "" { + return nil, &JSONRPCError{Code: ErrInvalidParams, Message: "Invalid params", Data: "text is required"} + } + + if s.engine == nil { + return nil, &JSONRPCError{Code: -32001, Message: "Detection unavailable", Data: "NER/pipeline mode does not support Detect"} + } + + matches := s.engine.Detect([]byte(args.Text)) + redacted := applyTemporaryPlaceholders(args.Text, matches) + return map[string]any{ + "redacted_text": redacted, + "match_count": len(matches), + }, nil +} + +func buildDetectResult(matches []redact.Match, includeOriginals bool) map[string]any { + hits := make([]map[string]any, len(matches)) + for i, m := range matches { + hit := map[string]any{ + "category": m.Category, + "start": m.Start, + "end": m.End, + } + if includeOriginals { + hit["original"] = m.Original + } + hits[i] = hit + } + return map[string]any{ + "has_sensitive": len(matches) > 0, + "match_count": len(matches), + "matches": hits, + } +} + +func applyTemporaryPlaceholders(text string, matches []redact.Match) string { + if len(matches) == 0 { + return text + } + + result := []byte(text) + // 从后往前替换，避免偏移 + for i := len(matches) - 1; i >= 0; i-- { + m := matches[i] + placeholder := fmt.Sprintf("__VG_%s_PREVIEW__", m.Category) + result = append(result[:m.Start], append([]byte(placeholder), result[m.End:]...)...) + } + return string(result) +} + +func (s *Server) handleListKeywords(_ json.RawMessage) (any, error) { + cfg := s.cfg.Get() + + var keywords []map[string]any + + // 从 config 中读取用户定义的关键词 + for _, kw := range cfg.Patterns.Keywords { + keywords = append(keywords, map[string]any{ + "keyword": kw.Value, + "category": kw.Category, + "source": "user", + }) + } + + // 从 config 中读取内置规则 + for _, b := range cfg.Patterns.Builtin { + keywords = append(keywords, map[string]any{ + "keyword": b, + "category": b, + "source": "builtin", + }) + } + + // 从 engine 中读取已加载的关键词（含 rule list 注入的） + if s.engine != nil { + for kw, cat := range s.engine.ListKeywords() { + keywords = append(keywords, map[string]any{ + "keyword": kw, + "category": cat, + "source": "engine", + }) + } + } + + return map[string]any{ + "count": len(keywords), + "keywords": keywords, + }, nil +} + +func (s *Server) handleListRuleLists(_ json.RawMessage) (any, error) { + cfg := s.cfg.Get() + + var lists []map[string]any + for _, rl := range cfg.Patterns.RuleLists { + lists = append(lists, map[string]any{ + "id": rl.ID, + "name": rl.Name, + "enabled": rl.Enabled, + "source": ruleListSource(rl), + }) + } + + return map[string]any{ + "count": len(lists), + "rule_lists": lists, + }, nil +} + +func ruleListSource(rl config.RuleListConfig) string { + if rl.URL != "" { + return "subscribed" + } + return "local" +} diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index f1fe2ee..d4b243a 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -57,6 +57,7 @@ type runtimeConfig struct { interceptMode string targets map[string]bool redactEng redact.Redactor + redactEngine *redact.Engine // 用于 MCP Detect（仅非 NER 模式） restoreEng *restore.Engine websocketRedactionBeta bool } @@ -170,7 +171,7 @@ func (s *Server) Start() error { // ServeMux may return a 301 redirect, breaking HTTPS proxy traffic (clients keep reconnecting). // Use a custom router here: only /manager/ goes to the admin UI; everything else goes to the proxy (including CONNECT). adminHandler := s.admin.Handler() - mcpServer := mcp.NewServer(s.config, s.admin, version.Version) + mcpServer := mcp.NewServer(s.config, s.admin, s.GetRedactEngine(), version.Version) mcpHandler := mcp.AuthMiddleware(mcpServer.Handler()) handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -215,6 +216,15 @@ func (s *Server) runtimeSnapshot() runtimeConfig { return v.(runtimeConfig) } +// GetRedactEngine returns the current redact engine (nil if NER/pipeline mode is active) +func (s *Server) GetRedactEngine() *redact.Engine { + if s == nil { + return nil + } + rt := s.runtimeSnapshot() + return rt.redactEngine +} + func (s *Server) shouldIntercept(host string) bool { rt := s.runtimeSnapshot() if rt.interceptMode == "global" { @@ -1271,6 +1281,7 @@ func (s *Server) applyConfig(c config.Config) { } var redactor redact.Redactor + var redactEng *redact.Engine if len(ruleRecs) > 0 || c.Patterns.NER.Enabled { var merged []piirec.Recognizer // Keywords: merge into one recognizer to avoid duplicate scans. @@ -1298,7 +1309,7 @@ func (s *Server) applyConfig(c config.Config) { p.SetExclude(exclude) redactor = p } else { - redactEng := redact.NewEngine(s.session, prefix) + redactEng = redact.NewEngine(s.session, prefix) for _, kw := range kws { redactEng.AddKeyword(kw.Text, kw.Category) } @@ -1325,6 +1336,7 @@ func (s *Server) applyConfig(c config.Config) { interceptMode: interceptMode, targets: targets, redactEng: redactor, + redactEngine: redactEng, restoreEng: restore.NewEngine(s.session, prefix), websocketRedactionBeta: c.Proxy.WebSocketRedactionBeta, }) diff --git a/internal/redact/engine.go b/internal/redact/engine.go index cb13319..eac1d89 100644 --- a/internal/redact/engine.go +++ b/internal/redact/engine.go @@ -53,6 +53,26 @@ func (e *Engine) AddKeyword(keyword, category string) { e.keywords[keyword] = category } +// ListKeywords returns all loaded keyword rules +func (e *Engine) ListKeywords() map[string]string { + if e == nil { + return nil + } + result := make(map[string]string, len(e.keywords)) + for k, v := range e.keywords { + result[k] = v + } + return result +} + +// ListRegexCategories returns category for each loaded regex +func (e *Engine) ListRegexCategories() []string { + if e == nil { + return nil + } + return append([]string(nil), e.regexCats...) +} + // AddRegex adds a regex pattern func (e *Engine) AddRegex(pattern, category string) error { re, err := regexp.Compile(pattern) @@ -69,6 +89,82 @@ func (e *Engine) AddExclude(pattern string) { e.exclude[pattern] = true } +// Detect 只读检测：扫描输入中的敏感数据，返回匹配列表，但不生成占位符、不注册到 session。 +func (e *Engine) Detect(input []byte) []Match { + e.ensureKeywordMatcher() + spans := textsafe.RedactableSpans(input) + + var matches []Match + if e.kwAC != nil { + matches = make([]Match, 0, min(len(e.kwCats), 64)) + } + + for _, span := range spans { + segment := input[span.Start:span.End] + if len(segment) == 0 { + continue + } + + if e.kwAC != nil { + scratchAny := e.kwScratch.Get() + lastEnd, _ := scratchAny.([]int) + + e.kwAC.EachMatchNonOverlappingPerPattern(segment, lastEnd, func(id, start, end int) bool { + cat := "" + if id >= 0 && id < len(e.kwCats) { + cat = e.kwCats[id] + } + globalStart := span.Start + start + globalEnd := span.Start + end + orig := string(input[globalStart:globalEnd]) + if e.isExcluded(orig) { + return true + } + matches = append(matches, Match{ + Start: globalStart, + End: globalEnd, + Original: orig, + Category: cat, + }) + return true + }) + + if lastEnd != nil { + e.kwScratch.Put(lastEnd) + } + } + + for i, re := range e.regex { + locs := re.FindAllSubmatchIndex(segment, -1) + for _, loc := range locs { + if len(loc) < 2 { + continue + } + start, end := loc[0], loc[1] + if len(loc) >= 4 && loc[2] >= 0 && loc[3] >= 0 { + start, end = loc[2], loc[3] + } + globalStart := span.Start + start + globalEnd := span.Start + end + if globalStart < 0 || globalEnd < 0 || globalStart >= globalEnd || globalEnd > len(input) { + continue + } + original := string(input[globalStart:globalEnd]) + if !e.isExcluded(original) { + matches = append(matches, Match{ + Start: globalStart, + End: globalEnd, + Original: original, + Category: e.regexCats[i], + }) + } + } + } + } + + return matches +} + // Redact scans and redacts sensitive data from the input func (e *Engine) Redact(input []byte) ([]byte, int) { out, matches := e.RedactWithMatches(input) diff --git a/internal/redact/engine_test.go b/internal/redact/engine_test.go new file mode 100644 index 0000000..221150d --- /dev/null +++ b/internal/redact/engine_test.go @@ -0,0 +1,94 @@ +package redact + +import ( + "testing" + + "github.com/inkdust2021/vibeguard/internal/session" +) + +func TestEngine_Detect_FindsKeyword(t *testing.T) { + sess := session.NewManager(0, 1000) + eng := NewEngine(sess, "__VG_") + eng.AddKeyword("secret123", "TEXT") + + matches := eng.Detect([]byte("hello secret123 world")) + if len(matches) != 1 { + t.Fatalf("expected 1 match, got %d", len(matches)) + } + if matches[0].Original != "secret123" { + t.Errorf("expected original=secret123, got %s", matches[0].Original) + } + if matches[0].Category != "TEXT" { + t.Errorf("expected category=TEXT, got %s", matches[0].Category) + } +} + +func TestEngine_Detect_FindsRegex(t *testing.T) { + sess := session.NewManager(0, 1000) + eng := NewEngine(sess, "__VG_") + eng.AddRegex(`\b\d{11}\b`, "CHINA_PHONE") + + input := []byte("call 13912345678 now") + matches := eng.Detect(input) + if len(matches) != 1 { + t.Fatalf("expected 1 match, got %d", len(matches)) + } + if matches[0].Original != "13912345678" { + t.Errorf("expected original=13912345678, got %s", matches[0].Original) + } +} + +func TestEngine_Detect_NoSideEffects(t *testing.T) { + sess := session.NewManager(0, 1000) + eng := NewEngine(sess, "__VG_") + eng.AddKeyword("password", "TEXT") + + // 检测前 session 应为空 + if count := sess.Size(); count != 0 { + t.Fatalf("expected 0 session entries before detect, got %d", count) + } + + eng.Detect([]byte("my password is abc")) + + // 检测后 session 仍应为空 + if count := sess.Size(); count != 0 { + t.Errorf("expected 0 session entries after detect (read-only), got %d", count) + } +} + +func TestEngine_Detect_NoPlaceholder(t *testing.T) { + sess := session.NewManager(0, 1000) + eng := NewEngine(sess, "__VG_") + eng.AddKeyword("secret", "TEXT") + + matches := eng.Detect([]byte("top secret info")) + if len(matches) != 1 { + t.Fatalf("expected 1 match, got %d", len(matches)) + } + if matches[0].Placeholder != "" { + t.Errorf("expected empty placeholder in detect mode, got %s", matches[0].Placeholder) + } +} + +func TestEngine_Detect_EmptyInput(t *testing.T) { + sess := session.NewManager(0, 1000) + eng := NewEngine(sess, "__VG_") + eng.AddKeyword("secret", "TEXT") + + matches := eng.Detect([]byte("")) + if len(matches) != 0 { + t.Errorf("expected 0 matches for empty input, got %d", len(matches)) + } +} + +func TestEngine_Detect_Exclude(t *testing.T) { + sess := session.NewManager(0, 1000) + eng := NewEngine(sess, "__VG_") + eng.AddKeyword("safe_value", "TEXT") + eng.AddExclude("safe_value") + + matches := eng.Detect([]byte("contains safe_value here")) + if len(matches) != 0 { + t.Errorf("expected 0 matches for excluded keyword, got %d", len(matches)) + } +} From 0167954687e3ad3769635efbbae5d1d16b39bf41 Mon Sep 17 00:00:00 2001 From: Jae Chen Date: Tue, 16 Jun 2026 23:56:59 +0800 Subject: [PATCH 3/4] test(mcp): add E2E tests for MCP server via real vibeguard process MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 启动编译后的二进制进程，验证 MCP 端到端通信路径。 覆盖行为： - 进程启动后 /mcp 端点可达，返回有效 JSON-RPC 响应 - token 自动生成机制在真实文件系统上工作 - HTTP 路由正确分流（/mcp 不被 admin 或 proxy 拦截） - 5 个工具的端到端响应（initialize, tools/list, get_proxy_status, detect_sensitive, preview_redacted, list_keywords, list_rule_lists） - 认证：无 token 和无效 token 均返回 401 - detect_sensitive 默认不暴露原始值，include_originals=true 时返回 技术要点： - 二进制只编译一次（sync.Once 共享） - 随机端口避免冲突 - 临时 HOME 隔离每个测试的 .vibeguard 目录 - 显式禁用 rule_lists 让代理走 raw Engine 模式（detect 工具可用） --- cmd/vibeguard/e2e_test.go | 546 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 546 insertions(+) create mode 100644 cmd/vibeguard/e2e_test.go diff --git a/cmd/vibeguard/e2e_test.go b/cmd/vibeguard/e2e_test.go new file mode 100644 index 0000000..59afcaf --- /dev/null +++ b/cmd/vibeguard/e2e_test.go @@ -0,0 +1,546 @@ +package main_test + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "net" + "net/http" + "os" + "os/exec" + "path/filepath" + "strings" + "sync" + "testing" + "time" +) + +var ( + binPathOnce sync.Once + binPath string + binBuildErr error +) + +// buildVibeguardOnce 编译 vibeguard 二进制（仅一次，所有 E2E 测试共享）。 +func buildVibeguardOnce(t *testing.T) string { + t.Helper() + binPathOnce.Do(func() { + dir, err := os.MkdirTemp("", "vibeguard-e2e-bin-*") + if err != nil { + binBuildErr = err + return + } + bin := filepath.Join(dir, "vibeguard") + wd, err := os.Getwd() + if err != nil { + binBuildErr = err + return + } + build := exec.Command("go", "build", "-o", bin, ".") + build.Dir = wd + if out, err := build.CombinedOutput(); err != nil { + binBuildErr = fmt.Errorf("build failed: %v\n%s", err, out) + return + } + binPath = bin + }) + if binBuildErr != nil { + t.Fatalf("build vibeguard: %v", binBuildErr) + } + return binPath +} + +// startVibeguard 启动一个独立的 vibeguard 进程，返回基础 URL、token 和清理函数。 +// +// 使用临时目录作为 HOME，OS 分配随机端口（先抢占再写 config）。 +func startVibeguard(t *testing.T) (baseURL, token string, cleanup func()) { + return startVibeguardWithKeywords(t) +} + +// startVibeguardWithKeywords 启动 vibeguard，并注入若干关键词规则（category 统一为 TEST）。 +func startVibeguardWithKeywords(t *testing.T, keywords ...string) (baseURL, token string, cleanup func()) { + t.Helper() + + bin := buildVibeguardOnce(t) + + // 1. 准备临时 HOME 和 config dir + tmpHome := t.TempDir() + configDir := filepath.Join(tmpHome, ".vibeguard") + if err := os.MkdirAll(configDir, 0o700); err != nil { + t.Fatalf("mkdir config dir: %v", err) + } + + // 2. 抢占一个随机端口 + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("reserve port: %v", err) + } + port := ln.Addr().(*net.TCPAddr).Port + ln.Close() + + listenAddr := fmt.Sprintf("127.0.0.1:%d", port) + + // 3. 写最小 config，覆盖默认监听端口 + // 显式提供空的 patterns 配置，避免加载默认 rule_lists（会让代理走 pipeline 模式，detect 不可用） + configPath := filepath.Join(configDir, "config.yaml") + + var keywordsYAML strings.Builder + if len(keywords) == 0 { + keywordsYAML.WriteString("[]") + } else { + keywordsYAML.WriteString("\n") + for _, kw := range keywords { + fmt.Fprintf(&keywordsYAML, ` - value: %q + category: "TEST" +`, kw) + } + } + + configYAML := fmt.Sprintf(`proxy: + listen: "%s" + intercept_mode: "global" +patterns: + keywords: %s + regex: [] + builtin: [] + rule_lists: [] + ner: + enabled: false +session: + ttl: "1h" + max_mappings: 1000 +log: + level: "error" +`, listenAddr, keywordsYAML.String()) + if err := os.WriteFile(configPath, []byte(configYAML), 0o600); err != nil { + t.Fatalf("write config: %v", err) + } + + // 4. 启动进程 + cmd := exec.Command(bin, "start", "--foreground", "--config", configPath) + cmd.Env = append(os.Environ(), + "HOME="+tmpHome, + "VIBEGUARD_LANG=en", + ) + stderr := &bytes.Buffer{} + cmd.Stderr = stderr + cmd.Stdout = stderr + + if err := cmd.Start(); err != nil { + t.Fatalf("start vibeguard: %v", err) + } + + // 5. 等待端口就绪 + baseURL = "http://" + listenAddr + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + conn, err := net.DialTimeout("tcp", listenAddr, 100*time.Millisecond) + if err == nil { + conn.Close() + break + } + time.Sleep(50 * time.Millisecond) + } + + // 6. 触发 token 生成 + probe, _ := http.NewRequest(http.MethodPost, baseURL+"/mcp", bytes.NewReader([]byte("{}"))) + probe.Header.Set("Content-Type", "application/json") + if r, err := http.DefaultClient.Do(probe); err == nil { + r.Body.Close() + } + + // 7. 读取 token 文件 + tokenPath := filepath.Join(configDir, "mcp_token") + tokenDeadline := time.Now().Add(3 * time.Second) + for time.Now().Before(tokenDeadline) { + data, err := os.ReadFile(tokenPath) + if err == nil && len(data) > 0 { + token = strings.TrimSpace(string(data)) + break + } + time.Sleep(50 * time.Millisecond) + } + + cleanup = func() { + _ = cmd.Process.Kill() + _, _ = cmd.Process.Wait() + if t.Failed() { + t.Logf("vibeguard stderr:\n%s", stderr.String()) + } + } + + if token == "" { + cleanup() + t.Fatalf("token not generated within timeout (config dir: %s)", configDir) + } + + return baseURL, token, cleanup +} + +// postMCP 发送一个 MCP JSON-RPC 请求，返回响应 map。 +func postMCP(t *testing.T, baseURL, token string, body map[string]any) map[string]any { + t.Helper() + data, _ := json.Marshal(body) + req, _ := http.NewRequest(http.MethodPost, baseURL+"/mcp", bytes.NewReader(data)) + req.Header.Set("Content-Type", "application/json") + if token != "" { + req.Header.Set("Authorization", "Bearer "+token) + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("post /mcp: %v", err) + } + defer resp.Body.Close() + respBody, _ := io.ReadAll(resp.Body) + if resp.StatusCode != http.StatusOK { + t.Fatalf("status %d: %s", resp.StatusCode, string(respBody)) + } + var result map[string]any + if err := json.Unmarshal(respBody, &result); err != nil { + t.Fatalf("parse response: %v\nbody: %s", err, string(respBody)) + } + return result +} + +func TestE2E_Initialize(t *testing.T) { + if testing.Short() { + t.Skip("skip E2E in short mode") + } + + baseURL, token, cleanup := startVibeguard(t) + defer cleanup() + + resp := postMCP(t, baseURL, token, map[string]any{ + "jsonrpc": "2.0", + "method": "initialize", + "id": 1, + }) + + if resp["jsonrpc"] != "2.0" { + t.Errorf("expected jsonrpc=2.0, got %v", resp["jsonrpc"]) + } + result, ok := resp["result"].(map[string]any) + if !ok { + t.Fatalf("expected result object, got %v", resp) + } + serverInfo := result["serverInfo"].(map[string]any) + if serverInfo["name"] != "vibeguard" { + t.Errorf("expected name=vibeguard, got %v", serverInfo["name"]) + } +} + +func TestE2E_ToolsList(t *testing.T) { + if testing.Short() { + t.Skip("skip E2E in short mode") + } + + baseURL, token, cleanup := startVibeguard(t) + defer cleanup() + + resp := postMCP(t, baseURL, token, map[string]any{ + "jsonrpc": "2.0", + "method": "tools/list", + "id": 1, + }) + + result := resp["result"].(map[string]any) + tools := result["tools"].([]any) + + // 5 个工具：get_proxy_status, detect_sensitive, preview_redacted, list_keywords, list_rule_lists + if len(tools) != 5 { + t.Errorf("expected 5 tools, got %d", len(tools)) + } + + names := make(map[string]bool) + for _, tool := range tools { + t := tool.(map[string]any) + names[t["name"].(string)] = true + } + expected := []string{"get_proxy_status", "detect_sensitive", "preview_redacted", "list_keywords", "list_rule_lists"} + for _, name := range expected { + if !names[name] { + t.Errorf("expected tool %s in list", name) + } + } +} + +func TestE2E_GetProxyStatus(t *testing.T) { + if testing.Short() { + t.Skip("skip E2E in short mode") + } + + baseURL, token, cleanup := startVibeguard(t) + defer cleanup() + + resp := postMCP(t, baseURL, token, map[string]any{ + "jsonrpc": "2.0", + "method": "tools/call", + "params": map[string]any{ + "name": "get_proxy_status", + "arguments": map[string]any{}, + }, + "id": 1, + }) + + result := resp["result"].(map[string]any) + content := result["content"].([]any) + textItem := content[0].(map[string]any) + + var status map[string]any + if err := json.Unmarshal([]byte(textItem["text"].(string)), &status); err != nil { + t.Fatalf("parse status: %v", err) + } + + for _, field := range []string{"version", "started_at", "total_requests", "intercept_mode", "listen_address"} { + if _, ok := status[field]; !ok { + t.Errorf("missing field: %s", field) + } + } +} + +func TestE2E_DetectSensitive(t *testing.T) { + if testing.Short() { + t.Skip("skip E2E in short mode") + } + + baseURL, token, cleanup := startVibeguardWithKeywords(t, "topsecret123") + defer cleanup() + + resp := postMCP(t, baseURL, token, map[string]any{ + "jsonrpc": "2.0", + "method": "tools/call", + "params": map[string]any{ + "name": "detect_sensitive", + "arguments": map[string]any{ + "text": "this contains topsecret123 keyword", + }, + }, + "id": 1, + }) + + if errObj, ok := resp["error"]; ok { + t.Fatalf("unexpected error: %v", errObj) + } + result := resp["result"].(map[string]any) + content := result["content"].([]any) + textItem := content[0].(map[string]any) + + var detect map[string]any + if err := json.Unmarshal([]byte(textItem["text"].(string)), &detect); err != nil { + t.Fatalf("parse detect: %v", err) + } + + if detect["has_sensitive"] != true { + t.Errorf("expected has_sensitive=true, got %v", detect["has_sensitive"]) + } + matches := detect["matches"].([]any) + if len(matches) == 0 { + t.Fatal("expected at least one match") + } + hit := matches[0].(map[string]any) + if hit["category"] != "TEST" { + t.Errorf("expected category=TEST, got %v", hit["category"]) + } + // 默认不返回 original + if _, ok := hit["original"]; ok { + t.Error("expected no original field by default") + } +} + +func TestE2E_DetectSensitive_WithOriginals(t *testing.T) { + if testing.Short() { + t.Skip("skip E2E in short mode") + } + + baseURL, token, cleanup := startVibeguardWithKeywords(t, "topsecret123") + defer cleanup() + + resp := postMCP(t, baseURL, token, map[string]any{ + "jsonrpc": "2.0", + "method": "tools/call", + "params": map[string]any{ + "name": "detect_sensitive", + "arguments": map[string]any{ + "text": "this contains topsecret123 keyword", + "include_originals": true, + }, + }, + "id": 1, + }) + + result := resp["result"].(map[string]any) + content := result["content"].([]any) + textItem := content[0].(map[string]any) + + var detect map[string]any + json.Unmarshal([]byte(textItem["text"].(string)), &detect) + + matches := detect["matches"].([]any) + hit := matches[0].(map[string]any) + if hit["original"] != "topsecret123" { + t.Errorf("expected original=topsecret123, got %v", hit["original"]) + } +} + +func TestE2E_PreviewRedacted(t *testing.T) { + if testing.Short() { + t.Skip("skip E2E in short mode") + } + + baseURL, token, cleanup := startVibeguardWithKeywords(t, "mysecret") + defer cleanup() + + resp := postMCP(t, baseURL, token, map[string]any{ + "jsonrpc": "2.0", + "method": "tools/call", + "params": map[string]any{ + "name": "preview_redacted", + "arguments": map[string]any{ + "text": "leak: mysecret here", + }, + }, + "id": 1, + }) + + if errObj, ok := resp["error"]; ok { + t.Fatalf("unexpected error: %v", errObj) + } + result := resp["result"].(map[string]any) + content := result["content"].([]any) + textItem := content[0].(map[string]any) + + var preview map[string]any + json.Unmarshal([]byte(textItem["text"].(string)), &preview) + + if preview["match_count"] != float64(1) { + t.Errorf("expected match_count=1, got %v", preview["match_count"]) + } + redacted := preview["redacted_text"].(string) + if strings.Contains(redacted, "mysecret") { + t.Errorf("expected mysecret to be redacted, got: %s", redacted) + } + if !strings.Contains(redacted, "__VG_TEST_PREVIEW__") { + t.Errorf("expected placeholder in redacted text, got: %s", redacted) + } +} + +func TestE2E_ListKeywords(t *testing.T) { + if testing.Short() { + t.Skip("skip E2E in short mode") + } + + baseURL, token, cleanup := startVibeguard(t) + defer cleanup() + + resp := postMCP(t, baseURL, token, map[string]any{ + "jsonrpc": "2.0", + "method": "tools/call", + "params": map[string]any{ + "name": "list_keywords", + "arguments": map[string]any{}, + }, + "id": 1, + }) + + if errObj, ok := resp["error"]; ok { + t.Fatalf("unexpected error: %v", errObj) + } + result := resp["result"].(map[string]any) + content := result["content"].([]any) + textItem := content[0].(map[string]any) + + var listResult map[string]any + if err := json.Unmarshal([]byte(textItem["text"].(string)), &listResult); err != nil { + t.Fatalf("parse list_keywords: %v", err) + } + + if _, ok := listResult["keywords"]; !ok { + t.Error("missing keywords field") + } + if _, ok := listResult["count"]; !ok { + t.Error("missing count field") + } +} + +func TestE2E_ListRuleLists(t *testing.T) { + if testing.Short() { + t.Skip("skip E2E in short mode") + } + + baseURL, token, cleanup := startVibeguard(t) + defer cleanup() + + resp := postMCP(t, baseURL, token, map[string]any{ + "jsonrpc": "2.0", + "method": "tools/call", + "params": map[string]any{ + "name": "list_rule_lists", + "arguments": map[string]any{}, + }, + "id": 1, + }) + + if errObj, ok := resp["error"]; ok { + t.Fatalf("unexpected error: %v", errObj) + } + result := resp["result"].(map[string]any) + content := result["content"].([]any) + textItem := content[0].(map[string]any) + + var listResult map[string]any + if err := json.Unmarshal([]byte(textItem["text"].(string)), &listResult); err != nil { + t.Fatalf("parse list_rule_lists: %v", err) + } + + if _, ok := listResult["rule_lists"]; !ok { + t.Error("missing rule_lists field") + } + if _, ok := listResult["count"]; !ok { + t.Error("missing count field") + } +} + +func TestE2E_AuthRequired(t *testing.T) { + if testing.Short() { + t.Skip("skip E2E in short mode") + } + + baseURL, _, cleanup := startVibeguard(t) + defer cleanup() + + // 不带 token + req, _ := http.NewRequest(http.MethodPost, baseURL+"/mcp", bytes.NewReader([]byte(`{"jsonrpc":"2.0","method":"initialize","id":1}`))) + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("request: %v", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusUnauthorized { + t.Errorf("expected 401 without token, got %d", resp.StatusCode) + } +} + +func TestE2E_InvalidToken(t *testing.T) { + if testing.Short() { + t.Skip("skip E2E in short mode") + } + + baseURL, _, cleanup := startVibeguard(t) + defer cleanup() + + req, _ := http.NewRequest(http.MethodPost, baseURL+"/mcp", bytes.NewReader([]byte(`{"jsonrpc":"2.0","method":"initialize","id":1}`))) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer wrong-token") + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("request: %v", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusUnauthorized { + t.Errorf("expected 401 with invalid token, got %d", resp.StatusCode) + } +} From f44b1b587d6fe21bcbb1c346933ce206265abe94 Mon Sep 17 00:00:00 2001 From: Jae Chen Date: Wed, 17 Jun 2026 00:10:34 +0800 Subject: [PATCH 4/4] docs(mcp): add MCP server usage guide MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 新增 docs/MCP.md，覆盖协议端点、Bearer Token 认证、CLI 一键配置、5 个工具的完整参数与返回字段、错误码、curl 示例、 Claude Code 集成步骤和故障排查。 --- docs/MCP.md | 313 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 313 insertions(+) create mode 100644 docs/MCP.md diff --git a/docs/MCP.md b/docs/MCP.md new file mode 100644 index 0000000..498c409 --- /dev/null +++ b/docs/MCP.md @@ -0,0 +1,313 @@ +# VibeGuard MCP Server 使用指南 + +VibeGuard 内置 MCP（Model Context Protocol）Server，让 Claude Code 等 AI 客户端能够直接查询代理状态、检测敏感信息、预览脱敏效果以及管理规则。 + +> 本文档面向使用者；想了解开发细节请阅读 `internal/mcp/` 源码与 `cmd/vibeguard/e2e_test.go`。 + +--- + +## 1. 协议与端点 + +- **协议**：JSON-RPC 2.0（[Streamable HTTP](https://modelcontextprotocol.io/) 风格） +- **路径**：`POST /mcp`（与代理监听同一端口，默认 `127.0.0.1:28657`） +- **完整 URL**：`http://127.0.0.1:28657/mcp` +- **方法限制**：仅支持 `POST`，其他方法返回 `405 Method Not Allowed` + +--- + +## 2. 认证 + +- 强制 **Bearer Token** 认证。 +- VibeGuard 启动时自动生成随机 token，写入 `~/.vibeguard/mcp_token`（权限 `0600`）。 +- 请求需要带头： + + ```http + Authorization: Bearer + ``` + +- 缺失或错误的 token 一律返回 `401 Unauthorized`。 +- token 持久化在用户目录中，重启 VibeGuard 不会变化；如需轮换可手动删除该文件后重启。 + +--- + +## 3. 一键生成客户端配置 + +VibeGuard 提供 CLI 子命令快速输出 Claude Code 可用的配置： + +```bash +# 输出 JSON（默认） +vibeguard mcp setup + +# 输出环境变量格式（适合 export 注入） +vibeguard mcp setup --format env +``` + +### 3.1 JSON 输出示例 + +```json +{ + "mcpServers": { + "vibeguard": { + "url": "http://127.0.0.1:28657/mcp", + "headers": { + "Authorization": "Bearer <自动生成的 token>" + } + } + } +} +``` + +将该 JSON 片段合并到 Claude Code 的 MCP 配置文件即可（位置以官方文档为准）。 + +### 3.2 env 输出示例 + +```bash +VG_MCP_URL=http://127.0.0.1:28657/mcp +VG_MCP_TOKEN=<自动生成的 token> +``` + +> 前提：`vibeguard start` 已经成功运行至少一次（用以生成 token）。 + +--- + +## 4. 工具清单（Phase 1） + +当前共注册 **5 个工具**，全部为只读，不会修改运行时状态。 + +| 工具名 | 用途 | +|--------|------| +| `get_proxy_status` | 查询代理运行状态、版本、统计计数、监听地址等 | +| `detect_sensitive` | 只读检测：返回命中规则列表，**不会** 生成占位符或注册到 session/WAL | +| `preview_redacted` | 预览脱敏效果，返回 `__VG__PREVIEW__` 临时占位符（不持久化） | +| `list_keywords` | 查询当前加载的关键词规则（来源：用户配置 / 内置 / engine） | +| `list_rule_lists` | 查询启用的规则列表（rule lists）及其元数据 | + +### 4.1 `get_proxy_status` + +**参数**：无。 + +**返回字段**（JSON 文本）： + +| 字段 | 含义 | +|------|------| +| `version` | VibeGuard 版本号 | +| `started_at` | 启动时间（RFC3339） | +| `uptime` | 运行时长 | +| `total_requests` | 处理总请求数 | +| `redacted_count` | 已脱敏请求数 | +| `restored_count` | 已还原响应数 | +| `errors` | 错误计数 | +| `intercept_mode` | 拦截模式（`global` / `selective`） | +| `websocket_beta` | WebSocket 脱敏 Beta 是否启用 | +| `listen_address` | 代理监听地址 | + +### 4.2 `detect_sensitive` + +**参数**： + +| 参数 | 类型 | 必需 | 说明 | +|------|------|------|------| +| `text` | string | 是 | 待检测文本 | +| `include_originals` | bool | 否 | 是否在结果中暴露原始命中值（默认 `false`） | + +**返回**： + +```json +{ + "has_sensitive": true, + "match_count": 1, + "matches": [ + { + "category": "TEST", + "start": 14, + "end": 26 + // include_originals=true 时才出现 "original": "topsecret123" + } + ] +} +``` + +> ⚠️ `detect_sensitive` 与 `preview_redacted` 仅在代理走 **raw Engine 模式** 时可用。如果启用了 NER 或 pipeline 模式，会返回 `-32001 Detection unavailable` 错误。 + +### 4.3 `preview_redacted` + +**参数**： + +| 参数 | 类型 | 必需 | 说明 | +|------|------|------|------| +| `text` | string | 是 | 待预览脱敏的文本 | + +**返回**： + +```json +{ + "redacted_text": "leak: __VG_TEST_PREVIEW__ here", + "match_count": 1 +} +``` + +占位符只用于本次预览，**不会** 写入 session 或 WAL。 + +### 4.4 `list_keywords` + +返回所有已加载关键词，附带来源（`user` / `builtin` / `engine`）： + +```json +{ + "count": 2, + "keywords": [ + {"keyword": "topsecret", "category": "TEST", "source": "user"}, + {"keyword": "email", "category": "EMAIL","source": "engine"} + ] +} +``` + +### 4.5 `list_rule_lists` + +```json +{ + "count": 1, + "rule_lists": [ + {"id": "vibeguard-default", "name": "VibeGuard Default", "enabled": true, "source": "subscribed"} + ] +} +``` + +`source` 取值： + +- `subscribed`：远程订阅 +- `local`：本地文件 + +--- + +## 5. 错误码 + +遵循 JSON-RPC 2.0 规范，业务错误使用 `-32000` 起的私有区段： + +| 错误码 | 含义 | +|--------|------| +| `-32700` | Parse error（请求体不是合法 JSON） | +| `-32600` | Invalid Request | +| `-32601` | Method not found / Tool not found | +| `-32602` | Invalid params（如缺少必填字段、tool name 为空） | +| `-32603` | Internal error | +| `-32001` | Detection unavailable（代理处于 NER/pipeline 模式，无法 Detect） | + +错误响应示例： + +```json +{ + "jsonrpc": "2.0", + "error": { + "code": -32602, + "message": "Invalid params", + "data": "text is required" + }, + "id": 1 +} +``` + +--- + +## 6. 端到端调用示例 + +### 6.1 初始化握手 + +```bash +TOKEN=$(cat ~/.vibeguard/mcp_token) + +curl -sS -X POST http://127.0.0.1:28657/mcp \ + -H "Authorization: Bearer $TOKEN" \ + -H "Content-Type: application/json" \ + -d '{"jsonrpc":"2.0","method":"initialize","id":1}' | jq +``` + +返回： + +```json +{ + "jsonrpc": "2.0", + "id": 1, + "result": { + "protocolVersion": "2024-11-05", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "vibeguard", "version": "..."} + } +} +``` + +### 6.2 列出工具 + +```bash +curl -sS -X POST http://127.0.0.1:28657/mcp \ + -H "Authorization: Bearer $TOKEN" \ + -H "Content-Type: application/json" \ + -d '{"jsonrpc":"2.0","method":"tools/list","id":2}' | jq +``` + +### 6.3 调用工具 + +```bash +curl -sS -X POST http://127.0.0.1:28657/mcp \ + -H "Authorization: Bearer $TOKEN" \ + -H "Content-Type: application/json" \ + -d '{ + "jsonrpc":"2.0", + "method":"tools/call", + "params":{ + "name":"detect_sensitive", + "arguments":{"text":"this contains topsecret123 keyword"} + }, + "id":3 + }' | jq +``` + +工具响应统一通过 `result.content[0].text` 返回 JSON 字符串，需要再 `JSON.parse` 一次。 + +--- + +## 7. Claude Code 集成步骤 + +1. 启动 VibeGuard：`vibeguard start --foreground`（或后台启动）。 +2. 生成配置：`vibeguard mcp setup > vibeguard.mcp.json`。 +3. 把 `mcpServers.vibeguard` 段合并到 Claude Code 的 MCP 配置文件。 +4. 重启 Claude Code，输入 `/mcp` 应能看到 `vibeguard` 服务器及其 5 个工具。 +5. 直接在对话中要求 Claude 调用工具，例如： + + > "用 vibeguard 检测下这段日志里有没有敏感信息：…" + +--- + +## 8. 注意事项 + +- **detect / preview 需要 raw Engine**：若你启用了 `patterns.ner.enabled: true` 或大量 rule_lists 走 pipeline 模式，这两个工具会拒绝服务。可暂时关闭 NER 或减少 rule lists。 +- **token 安全**：`~/.vibeguard/mcp_token` 是明文，请勿提交到 git；CI / 容器场景建议挂载只读 secret。 +- **端口冲突**：MCP 复用代理监听端口；修改 `proxy.listen` 后记得重新生成 `vibeguard mcp setup` 输出。 +- **跨主机访问**：默认绑定 `127.0.0.1`。如需跨机访问，请自行处理 TLS / 反向代理 / 防火墙，**不要** 把 token 直接暴露到公网。 +- **Phase 2 工具**：管理/控制类工具（如重载配置、新增关键词、注销 session）尚未实现，Roadmap 中规划。 + +--- + +## 9. 相关代码位置 + +| 模块 | 路径 | +|------|------| +| MCP Server 核心 | `internal/mcp/server.go` | +| 工具实现 | `internal/mcp/tools.go` | +| Bearer 认证 | `internal/mcp/auth.go` | +| CLI `mcp setup` | `cmd/vibeguard/main.go`（搜索 `runMcpSetup`） | +| 单元测试 | `internal/mcp/server_test.go` | +| E2E 测试 | `cmd/vibeguard/e2e_test.go` | + +--- + +## 10. 故障排查 + +| 现象 | 可能原因 | 解决方案 | +|------|----------|----------| +| `401 Unauthorized` | token 缺失或不正确 | 检查 `Authorization: Bearer ...` 头；重新读取 `~/.vibeguard/mcp_token` | +| `405 Method Not Allowed` | 用了 GET | 必须 `POST /mcp` | +| `-32700 Parse error` | 请求体非 JSON | 检查 Content-Type 与 body | +| `-32601 Tool not found` | 工具名拼错 | 通过 `tools/list` 校对 | +| `-32001 Detection unavailable` | 代理走 NER/pipeline 模式 | 关闭 NER / 减少 rule_lists，或仅使用 `get_proxy_status` 等不依赖 Engine 的工具 | +| `vibeguard mcp setup` 报无法读取 token | VibeGuard 还没启动过 | 先 `vibeguard start` 一次让 token 生成 |