cc_audit/
mcp_server.rs

1use crate::engine::scanner::{Scanner, ScannerConfig};
2use crate::engine::scanners::SkillScanner;
3use crate::error::Result;
4use crate::fix::AutoFixer;
5use crate::rules::{Finding, RuleEngine, ScanResult, Summary};
6use crate::scoring::RiskScore;
7use serde::{Deserialize, Serialize};
8use serde_json::{Value, json};
9use std::io::{BufRead, BufReader, Write};
10use std::path::PathBuf;
11
12/// MCP Server for cc-audit
13/// Provides security scanning capabilities via MCP protocol
14pub struct McpServer {
15    rule_engine: RuleEngine,
16}
17
18#[derive(Debug, Serialize, Deserialize)]
19struct JsonRpcRequest {
20    jsonrpc: String,
21    id: Option<Value>,
22    method: String,
23    params: Option<Value>,
24}
25
26#[derive(Debug, Serialize)]
27struct JsonRpcResponse {
28    jsonrpc: String,
29    id: Option<Value>,
30    #[serde(skip_serializing_if = "Option::is_none")]
31    result: Option<Value>,
32    #[serde(skip_serializing_if = "Option::is_none")]
33    error: Option<JsonRpcError>,
34}
35
36#[derive(Debug, Serialize)]
37struct JsonRpcError {
38    code: i32,
39    message: String,
40    #[serde(skip_serializing_if = "Option::is_none")]
41    data: Option<Value>,
42}
43
44/// Tool definitions for MCP
45#[derive(Debug, Serialize)]
46struct Tool {
47    name: String,
48    description: String,
49    #[serde(rename = "inputSchema")]
50    input_schema: Value,
51}
52
53impl McpServer {
54    pub fn new() -> Self {
55        Self {
56            rule_engine: RuleEngine::new(),
57        }
58    }
59
60    /// serde_json::to_string_pretty のエラーを JsonRpcError に変換するヘルパー
61    fn json_pretty(value: &impl serde::Serialize) -> std::result::Result<String, JsonRpcError> {
62        serde_json::to_string_pretty(value).map_err(|e| JsonRpcError {
63            code: -32603,
64            message: format!("Internal serialization error: {}", e),
65            data: None,
66        })
67    }
68
69    /// Run the MCP server (JSON-RPC over stdio)
70    pub fn run(&self) -> Result<()> {
71        let stdin = std::io::stdin();
72        let mut stdout = std::io::stdout();
73        let reader = BufReader::new(stdin.lock());
74
75        eprintln!("cc-audit MCP server started");
76
77        for line in reader.lines() {
78            let line = match line {
79                Ok(l) => l,
80                Err(e) => {
81                    eprintln!("Error reading input: {}", e);
82                    continue;
83                }
84            };
85
86            if line.is_empty() {
87                continue;
88            }
89
90            let request: JsonRpcRequest = match serde_json::from_str(&line) {
91                Ok(r) => r,
92                Err(e) => {
93                    let error_response = JsonRpcResponse {
94                        jsonrpc: "2.0".to_string(),
95                        id: None,
96                        result: None,
97                        error: Some(JsonRpcError {
98                            code: -32700,
99                            message: format!("Parse error: {}", e),
100                            data: None,
101                        }),
102                    };
103                    // SAFETY: JsonRpcResponse contains only simple, serializable types.
104                    // This unwrap_or_else provides a fallback for the unlikely case of serialization failure.
105                    let json = serde_json::to_string(&error_response)
106                        .unwrap_or_else(|_| r#"{"jsonrpc":"2.0","error":{"code":-32603,"message":"Internal error"}}"#.to_string());
107                    let _ = writeln!(stdout, "{}", json);
108                    let _ = stdout.flush();
109                    continue;
110                }
111            };
112
113            let response = self.handle_request(request);
114            // SAFETY: JsonRpcResponse contains only simple, serializable types.
115            // This unwrap_or_else provides a fallback for the unlikely case of serialization failure.
116            let json = serde_json::to_string(&response).unwrap_or_else(|_| {
117                r#"{"jsonrpc":"2.0","error":{"code":-32603,"message":"Internal error"}}"#
118                    .to_string()
119            });
120            let _ = writeln!(stdout, "{}", json);
121            let _ = stdout.flush();
122        }
123
124        Ok(())
125    }
126
127    fn handle_request(&self, request: JsonRpcRequest) -> JsonRpcResponse {
128        let result = match request.method.as_str() {
129            "initialize" => self.handle_initialize(&request.params),
130            "tools/list" => self.handle_list_tools(),
131            "tools/call" => self.handle_tool_call(&request.params),
132            "shutdown" => {
133                eprintln!("MCP server shutting down");
134                Ok(json!({}))
135            }
136            _ => Err(JsonRpcError {
137                code: -32601,
138                message: format!("Method not found: {}", request.method),
139                data: None,
140            }),
141        };
142
143        match result {
144            Ok(value) => JsonRpcResponse {
145                jsonrpc: "2.0".to_string(),
146                id: request.id,
147                result: Some(value),
148                error: None,
149            },
150            Err(error) => JsonRpcResponse {
151                jsonrpc: "2.0".to_string(),
152                id: request.id,
153                result: None,
154                error: Some(error),
155            },
156        }
157    }
158
159    fn handle_initialize(
160        &self,
161        _params: &Option<Value>,
162    ) -> std::result::Result<Value, JsonRpcError> {
163        Ok(json!({
164            "protocolVersion": "2024-11-05",
165            "capabilities": {
166                "tools": {}
167            },
168            "serverInfo": {
169                "name": "cc-audit",
170                "version": env!("CARGO_PKG_VERSION")
171            }
172        }))
173    }
174
175    fn handle_list_tools(&self) -> std::result::Result<Value, JsonRpcError> {
176        let tools = vec![
177            Tool {
178                name: "scan".to_string(),
179                description: "Scan a file or directory for security issues".to_string(),
180                input_schema: json!({
181                    "type": "object",
182                    "properties": {
183                        "path": {
184                            "type": "string",
185                            "description": "Path to scan (file or directory)"
186                        }
187                    },
188                    "required": ["path"]
189                }),
190            },
191            Tool {
192                name: "scan_content".to_string(),
193                description: "Scan content string for security issues".to_string(),
194                input_schema: json!({
195                    "type": "object",
196                    "properties": {
197                        "content": {
198                            "type": "string",
199                            "description": "Content to scan"
200                        },
201                        "filename": {
202                            "type": "string",
203                            "description": "Virtual filename for context"
204                        }
205                    },
206                    "required": ["content"]
207                }),
208            },
209            Tool {
210                name: "check_rule".to_string(),
211                description: "Check if content matches a specific rule".to_string(),
212                input_schema: json!({
213                    "type": "object",
214                    "properties": {
215                        "rule_id": {
216                            "type": "string",
217                            "description": "Rule ID to check (e.g., 'OP-001')"
218                        },
219                        "content": {
220                            "type": "string",
221                            "description": "Content to check"
222                        }
223                    },
224                    "required": ["rule_id", "content"]
225                }),
226            },
227            Tool {
228                name: "list_rules".to_string(),
229                description: "List all available security rules".to_string(),
230                input_schema: json!({
231                    "type": "object",
232                    "properties": {
233                        "category": {
234                            "type": "string",
235                            "description": "Filter by category (optional)"
236                        }
237                    }
238                }),
239            },
240            Tool {
241                name: "get_fix_suggestion".to_string(),
242                description: "Get a fix suggestion for a finding".to_string(),
243                input_schema: json!({
244                    "type": "object",
245                    "properties": {
246                        "finding_id": {
247                            "type": "string",
248                            "description": "Finding ID (rule ID)"
249                        },
250                        "code": {
251                            "type": "string",
252                            "description": "The problematic code"
253                        }
254                    },
255                    "required": ["finding_id", "code"]
256                }),
257            },
258        ];
259
260        Ok(json!({ "tools": tools }))
261    }
262
263    fn handle_tool_call(&self, params: &Option<Value>) -> std::result::Result<Value, JsonRpcError> {
264        let params = params.as_ref().ok_or_else(|| JsonRpcError {
265            code: -32602,
266            message: "Missing params".to_string(),
267            data: None,
268        })?;
269
270        let name = params
271            .get("name")
272            .and_then(|v| v.as_str())
273            .ok_or_else(|| JsonRpcError {
274                code: -32602,
275                message: "Missing tool name".to_string(),
276                data: None,
277            })?;
278
279        let arguments = params.get("arguments").cloned().unwrap_or(json!({}));
280
281        match name {
282            "scan" => self.tool_scan(&arguments),
283            "scan_content" => self.tool_scan_content(&arguments),
284            "check_rule" => self.tool_check_rule(&arguments),
285            "list_rules" => self.tool_list_rules(&arguments),
286            "get_fix_suggestion" => self.tool_get_fix_suggestion(&arguments),
287            _ => Err(JsonRpcError {
288                code: -32602,
289                message: format!("Unknown tool: {}", name),
290                data: None,
291            }),
292        }
293    }
294
295    fn tool_scan(&self, args: &Value) -> std::result::Result<Value, JsonRpcError> {
296        let path = args
297            .get("path")
298            .and_then(|v| v.as_str())
299            .ok_or_else(|| JsonRpcError {
300                code: -32602,
301                message: "Missing 'path' argument".to_string(),
302                data: None,
303            })?;
304
305        let path = PathBuf::from(path);
306        let scanner = SkillScanner::new();
307
308        match scanner.scan_path(&path) {
309            Ok(findings) => {
310                let summary = Summary::from_findings(&findings);
311                let risk_score = RiskScore::from_findings(&findings);
312                let result = ScanResult {
313                    version: env!("CARGO_PKG_VERSION").to_string(),
314                    scanned_at: chrono::Utc::now().to_rfc3339(),
315                    target: path.display().to_string(),
316                    summary,
317                    findings,
318                    risk_score: Some(risk_score),
319                    elapsed_ms: 0,
320                };
321                Ok(json!({
322                    "content": [{
323                        "type": "text",
324                        "text": Self::json_pretty(&result)?
325                    }]
326                }))
327            }
328            Err(e) => Err(JsonRpcError {
329                code: -32000,
330                message: format!("Scan failed: {}", e),
331                data: None,
332            }),
333        }
334    }
335
336    fn tool_scan_content(&self, args: &Value) -> std::result::Result<Value, JsonRpcError> {
337        let content = args
338            .get("content")
339            .and_then(|v| v.as_str())
340            .ok_or_else(|| JsonRpcError {
341                code: -32602,
342                message: "Missing 'content' argument".to_string(),
343                data: None,
344            })?;
345
346        let filename = args
347            .get("filename")
348            .and_then(|v| v.as_str())
349            .unwrap_or("content.md");
350
351        let config = ScannerConfig::new();
352        let findings = config.check_content(content, filename);
353
354        let summary = Summary::from_findings(&findings);
355        let risk_score = RiskScore::from_findings(&findings);
356
357        Ok(json!({
358            "content": [{
359                "type": "text",
360                "text": Self::json_pretty(&json!({
361                    "findings": findings,
362                    "summary": summary,
363                    "risk_score": risk_score
364                }))?
365            }]
366        }))
367    }
368
369    fn tool_check_rule(&self, args: &Value) -> std::result::Result<Value, JsonRpcError> {
370        let rule_id = args
371            .get("rule_id")
372            .and_then(|v| v.as_str())
373            .ok_or_else(|| JsonRpcError {
374                code: -32602,
375                message: "Missing 'rule_id' argument".to_string(),
376                data: None,
377            })?;
378
379        let content = args
380            .get("content")
381            .and_then(|v| v.as_str())
382            .ok_or_else(|| JsonRpcError {
383                code: -32602,
384                message: "Missing 'content' argument".to_string(),
385                data: None,
386            })?;
387
388        // Check if rule exists
389        let Some(rule) = self.rule_engine.get_rule(rule_id) else {
390            return Ok(json!({
391                "content": [{
392                    "type": "text",
393                    "text": format!("Rule '{}' not found", rule_id)
394                }]
395            }));
396        };
397
398        // Check if any pattern matches
399        let mut matches = false;
400        for pattern in &rule.patterns {
401            if pattern.is_match(content) {
402                matches = true;
403                break;
404            }
405        }
406
407        Ok(json!({
408            "content": [{
409                "type": "text",
410                "text": Self::json_pretty(&json!({
411                    "rule_id": rule_id,
412                    "rule_name": rule.name,
413                    "severity": format!("{:?}", rule.severity),
414                    "matches": matches,
415                    "message": if matches {
416                        format!("Content matches rule: {}", rule.message)
417                    } else {
418                        "No match found".to_string()
419                    }
420                }))?
421            }]
422        }))
423    }
424
425    fn tool_list_rules(&self, args: &Value) -> std::result::Result<Value, JsonRpcError> {
426        let category_filter = args
427            .get("category")
428            .and_then(|v| v.as_str())
429            .map(|s| s.to_lowercase());
430
431        let rules = self.rule_engine.get_all_rules();
432        let filtered: Vec<_> = rules
433            .iter()
434            .filter(|r| {
435                if let Some(ref cat) = category_filter {
436                    format!("{:?}", r.category).to_lowercase().contains(cat)
437                } else {
438                    true
439                }
440            })
441            .map(|r| {
442                json!({
443                    "id": r.id,
444                    "name": r.name,
445                    "severity": format!("{:?}", r.severity),
446                    "category": format!("{:?}", r.category),
447                    "confidence": format!("{:?}", r.confidence)
448                })
449            })
450            .collect();
451
452        Ok(json!({
453            "content": [{
454                "type": "text",
455                "text": Self::json_pretty(&json!({
456                    "total": filtered.len(),
457                    "rules": filtered
458                }))?
459            }]
460        }))
461    }
462
463    fn tool_get_fix_suggestion(&self, args: &Value) -> std::result::Result<Value, JsonRpcError> {
464        let finding_id = args
465            .get("finding_id")
466            .and_then(|v| v.as_str())
467            .ok_or_else(|| JsonRpcError {
468                code: -32602,
469                message: "Missing 'finding_id' argument".to_string(),
470                data: None,
471            })?;
472
473        let code = args
474            .get("code")
475            .and_then(|v| v.as_str())
476            .ok_or_else(|| JsonRpcError {
477                code: -32602,
478                message: "Missing 'code' argument".to_string(),
479                data: None,
480            })?;
481
482        // Create a mock finding for the fixer
483        let Some(rule) = self.rule_engine.get_rule(finding_id) else {
484            return Ok(json!({
485                "content": [{
486                    "type": "text",
487                    "text": format!("No fix suggestion available for rule '{}'", finding_id)
488                }]
489            }));
490        };
491        let finding = Finding {
492            id: finding_id.to_string(),
493            severity: rule.severity,
494            category: rule.category,
495            confidence: rule.confidence,
496            name: rule.name.to_string(),
497            location: crate::rules::Location {
498                file: "virtual".to_string(),
499                line: 1,
500                column: None,
501            },
502            code: code.to_string(),
503            message: rule.message.to_string(),
504            recommendation: rule.recommendation.to_string(),
505            fix_hint: rule.fix_hint.map(|s| s.to_string()),
506            cwe_ids: rule.cwe_ids.iter().map(|s| s.to_string()).collect(),
507            rule_severity: None,
508            client: None,
509            context: None,
510        };
511
512        let fixer = AutoFixer::new(true);
513        let fixes = fixer.generate_fixes(&[finding]);
514
515        if fixes.is_empty() {
516            Ok(json!({
517                "content": [{
518                    "type": "text",
519                    "text": format!("No automatic fix available for {}. Manual review recommended.\n\nRecommendation: {}", finding_id, rule.recommendation)
520                }]
521            }))
522        } else {
523            let fix = &fixes[0];
524            Ok(json!({
525                "content": [{
526                    "type": "text",
527                    "text": Self::json_pretty(&json!({
528                        "has_fix": true,
529                        "description": fix.description,
530                        "original": fix.original,
531                        "replacement": fix.replacement
532                    }))?
533                }]
534            }))
535        }
536    }
537}
538
539impl Default for McpServer {
540    fn default() -> Self {
541        Self::new()
542    }
543}
544
545#[cfg(test)]
546mod tests {
547    use super::*;
548    use tempfile::TempDir;
549
550    #[test]
551    fn test_mcp_server_new() {
552        let server = McpServer::new();
553        assert!(!server.rule_engine.get_all_rules().is_empty());
554    }
555
556    #[test]
557    fn test_mcp_server_default() {
558        let server = McpServer::default();
559        assert!(!server.rule_engine.get_all_rules().is_empty());
560    }
561
562    #[test]
563    fn test_handle_initialize() {
564        let server = McpServer::new();
565        let result = server.handle_initialize(&None).unwrap();
566
567        assert!(result.get("protocolVersion").is_some());
568        assert!(result.get("serverInfo").is_some());
569    }
570
571    #[test]
572    fn test_handle_initialize_with_params() {
573        let server = McpServer::new();
574        let params = Some(json!({"clientInfo": {"name": "test"}}));
575        let result = server.handle_initialize(&params).unwrap();
576
577        assert!(result.get("protocolVersion").is_some());
578    }
579
580    #[test]
581    fn test_handle_list_tools() {
582        let server = McpServer::new();
583        let result = server.handle_list_tools().unwrap();
584
585        let tools = result.get("tools").unwrap().as_array().unwrap();
586        assert_eq!(tools.len(), 5);
587
588        let tool_names: Vec<&str> = tools
589            .iter()
590            .map(|t| t.get("name").unwrap().as_str().unwrap())
591            .collect();
592        assert!(tool_names.contains(&"scan"));
593        assert!(tool_names.contains(&"scan_content"));
594        assert!(tool_names.contains(&"check_rule"));
595        assert!(tool_names.contains(&"list_rules"));
596        assert!(tool_names.contains(&"get_fix_suggestion"));
597    }
598
599    #[test]
600    fn test_tool_scan_content() {
601        let server = McpServer::new();
602        let args = json!({
603            "content": "allowed-tools: *",
604            "filename": "test.md"
605        });
606
607        let result = server.tool_scan_content(&args).unwrap();
608        let content = result.get("content").unwrap().as_array().unwrap();
609        assert!(!content.is_empty());
610    }
611
612    #[test]
613    fn test_tool_scan_content_no_filename() {
614        let server = McpServer::new();
615        let args = json!({
616            "content": "some safe content"
617        });
618
619        let result = server.tool_scan_content(&args).unwrap();
620        let content = result.get("content").unwrap().as_array().unwrap();
621        assert!(!content.is_empty());
622    }
623
624    #[test]
625    fn test_tool_scan_content_missing_content() {
626        let server = McpServer::new();
627        let args = json!({});
628
629        let result = server.tool_scan_content(&args);
630        assert!(result.is_err());
631    }
632
633    #[test]
634    fn test_tool_list_rules() {
635        let server = McpServer::new();
636        let args = json!({});
637
638        let result = server.tool_list_rules(&args).unwrap();
639        let content = result.get("content").unwrap().as_array().unwrap();
640        assert!(!content.is_empty());
641    }
642
643    #[test]
644    fn test_tool_list_rules_with_category() {
645        let server = McpServer::new();
646        let args = json!({"category": "exfiltration"});
647
648        let result = server.tool_list_rules(&args).unwrap();
649        let content = result.get("content").unwrap().as_array().unwrap();
650        assert!(!content.is_empty());
651    }
652
653    #[test]
654    fn test_tool_check_rule() {
655        let server = McpServer::new();
656        let args = json!({
657            "rule_id": "OP-001",
658            "content": "allowed-tools: *"
659        });
660
661        let result = server.tool_check_rule(&args).unwrap();
662        let content = result.get("content").unwrap().as_array().unwrap();
663        assert!(!content.is_empty());
664    }
665
666    #[test]
667    fn test_tool_check_rule_no_match() {
668        let server = McpServer::new();
669        let args = json!({
670            "rule_id": "OP-001",
671            "content": "allowed-tools: Read, Write"
672        });
673
674        let result = server.tool_check_rule(&args).unwrap();
675        let content = result.get("content").unwrap().as_array().unwrap();
676        let text = content[0].get("text").unwrap().as_str().unwrap();
677        assert!(text.contains("No match found") || text.contains("matches"));
678    }
679
680    #[test]
681    fn test_tool_check_rule_not_found() {
682        let server = McpServer::new();
683        let args = json!({
684            "rule_id": "NONEXISTENT-001",
685            "content": "some content"
686        });
687
688        let result = server.tool_check_rule(&args).unwrap();
689        let content = result.get("content").unwrap().as_array().unwrap();
690        let text = content[0].get("text").unwrap().as_str().unwrap();
691        assert!(text.contains("not found"));
692    }
693
694    #[test]
695    fn test_tool_check_rule_missing_rule_id() {
696        let server = McpServer::new();
697        let args = json!({
698            "content": "some content"
699        });
700
701        let result = server.tool_check_rule(&args);
702        assert!(result.is_err());
703    }
704
705    #[test]
706    fn test_tool_check_rule_missing_content() {
707        let server = McpServer::new();
708        let args = json!({
709            "rule_id": "OP-001"
710        });
711
712        let result = server.tool_check_rule(&args);
713        assert!(result.is_err());
714    }
715
716    #[test]
717    fn test_tool_scan_valid_path() {
718        let temp_dir = TempDir::new().unwrap();
719        let test_file = temp_dir.path().join("SKILL.md");
720        std::fs::write(&test_file, "---\nallowed-tools: *\n---\n").unwrap();
721
722        let server = McpServer::new();
723        let args = json!({"path": test_file.display().to_string()});
724
725        let result = server.tool_scan(&args).unwrap();
726        let content = result.get("content").unwrap().as_array().unwrap();
727        assert!(!content.is_empty());
728    }
729
730    #[test]
731    fn test_tool_scan_invalid_path() {
732        let server = McpServer::new();
733        let args = json!({"path": "/nonexistent/path/that/does/not/exist"});
734
735        let result = server.tool_scan(&args);
736        assert!(result.is_err());
737    }
738
739    #[test]
740    fn test_tool_scan_missing_path() {
741        let server = McpServer::new();
742        let args = json!({});
743
744        let result = server.tool_scan(&args);
745        assert!(result.is_err());
746    }
747
748    #[test]
749    fn test_tool_get_fix_suggestion_valid() {
750        let server = McpServer::new();
751        let args = json!({
752            "finding_id": "OP-001",
753            "code": "allowed-tools: *"
754        });
755
756        let result = server.tool_get_fix_suggestion(&args).unwrap();
757        let content = result.get("content").unwrap().as_array().unwrap();
758        assert!(!content.is_empty());
759    }
760
761    #[test]
762    fn test_tool_get_fix_suggestion_no_fix_available() {
763        let server = McpServer::new();
764        let args = json!({
765            "finding_id": "EX-001",
766            "code": "echo hello"
767        });
768
769        let result = server.tool_get_fix_suggestion(&args).unwrap();
770        let content = result.get("content").unwrap().as_array().unwrap();
771        let text = content[0].get("text").unwrap().as_str().unwrap();
772        assert!(text.contains("No automatic fix") || text.contains("has_fix"));
773    }
774
775    #[test]
776    fn test_tool_get_fix_suggestion_rule_not_found() {
777        let server = McpServer::new();
778        let args = json!({
779            "finding_id": "NONEXISTENT-001",
780            "code": "some code"
781        });
782
783        let result = server.tool_get_fix_suggestion(&args).unwrap();
784        let content = result.get("content").unwrap().as_array().unwrap();
785        let text = content[0].get("text").unwrap().as_str().unwrap();
786        assert!(text.contains("No fix suggestion available"));
787    }
788
789    #[test]
790    fn test_tool_get_fix_suggestion_missing_finding_id() {
791        let server = McpServer::new();
792        let args = json!({
793            "code": "some code"
794        });
795
796        let result = server.tool_get_fix_suggestion(&args);
797        assert!(result.is_err());
798    }
799
800    #[test]
801    fn test_tool_get_fix_suggestion_missing_code() {
802        let server = McpServer::new();
803        let args = json!({
804            "finding_id": "OP-001"
805        });
806
807        let result = server.tool_get_fix_suggestion(&args);
808        assert!(result.is_err());
809    }
810
811    #[test]
812    fn test_handle_request_initialize() {
813        let server = McpServer::new();
814        let request = JsonRpcRequest {
815            jsonrpc: "2.0".to_string(),
816            id: Some(json!(1)),
817            method: "initialize".to_string(),
818            params: None,
819        };
820
821        let response = server.handle_request(request);
822        assert!(response.result.is_some());
823        assert!(response.error.is_none());
824    }
825
826    #[test]
827    fn test_handle_request_tools_list() {
828        let server = McpServer::new();
829        let request = JsonRpcRequest {
830            jsonrpc: "2.0".to_string(),
831            id: Some(json!(2)),
832            method: "tools/list".to_string(),
833            params: None,
834        };
835
836        let response = server.handle_request(request);
837        assert!(response.result.is_some());
838        assert!(response.error.is_none());
839    }
840
841    #[test]
842    fn test_handle_request_shutdown() {
843        let server = McpServer::new();
844        let request = JsonRpcRequest {
845            jsonrpc: "2.0".to_string(),
846            id: Some(json!(3)),
847            method: "shutdown".to_string(),
848            params: None,
849        };
850
851        let response = server.handle_request(request);
852        assert!(response.result.is_some());
853        assert!(response.error.is_none());
854    }
855
856    #[test]
857    fn test_handle_request_unknown_method() {
858        let server = McpServer::new();
859        let request = JsonRpcRequest {
860            jsonrpc: "2.0".to_string(),
861            id: Some(json!(4)),
862            method: "unknown/method".to_string(),
863            params: None,
864        };
865
866        let response = server.handle_request(request);
867        assert!(response.result.is_none());
868        assert!(response.error.is_some());
869        assert_eq!(response.error.as_ref().unwrap().code, -32601);
870    }
871
872    #[test]
873    fn test_handle_tool_call_missing_params() {
874        let server = McpServer::new();
875        let result = server.handle_tool_call(&None);
876        assert!(result.is_err());
877        assert_eq!(result.unwrap_err().code, -32602);
878    }
879
880    #[test]
881    fn test_handle_tool_call_missing_name() {
882        let server = McpServer::new();
883        let params = Some(json!({"arguments": {}}));
884        let result = server.handle_tool_call(&params);
885        assert!(result.is_err());
886    }
887
888    #[test]
889    fn test_handle_tool_call_unknown_tool() {
890        let server = McpServer::new();
891        let params = Some(json!({
892            "name": "unknown_tool",
893            "arguments": {}
894        }));
895        let result = server.handle_tool_call(&params);
896        assert!(result.is_err());
897        assert!(result.unwrap_err().message.contains("Unknown tool"));
898    }
899
900    #[test]
901    fn test_handle_tool_call_scan_content() {
902        let server = McpServer::new();
903        let params = Some(json!({
904            "name": "scan_content",
905            "arguments": {
906                "content": "safe content"
907            }
908        }));
909
910        let result = server.handle_tool_call(&params);
911        assert!(result.is_ok());
912    }
913
914    #[test]
915    fn test_handle_tool_call_list_rules() {
916        let server = McpServer::new();
917        let params = Some(json!({
918            "name": "list_rules",
919            "arguments": {}
920        }));
921
922        let result = server.handle_tool_call(&params);
923        assert!(result.is_ok());
924    }
925
926    #[test]
927    fn test_handle_tool_call_check_rule() {
928        let server = McpServer::new();
929        let params = Some(json!({
930            "name": "check_rule",
931            "arguments": {
932                "rule_id": "OP-001",
933                "content": "allowed-tools: *"
934            }
935        }));
936
937        let result = server.handle_tool_call(&params);
938        assert!(result.is_ok());
939    }
940
941    #[test]
942    fn test_handle_tool_call_get_fix_suggestion() {
943        let server = McpServer::new();
944        let params = Some(json!({
945            "name": "get_fix_suggestion",
946            "arguments": {
947                "finding_id": "OP-001",
948                "code": "allowed-tools: *"
949            }
950        }));
951
952        let result = server.handle_tool_call(&params);
953        assert!(result.is_ok());
954    }
955
956    #[test]
957    fn test_json_rpc_request_debug() {
958        let request = JsonRpcRequest {
959            jsonrpc: "2.0".to_string(),
960            id: Some(json!(1)),
961            method: "test".to_string(),
962            params: None,
963        };
964
965        let debug_str = format!("{:?}", request);
966        assert!(debug_str.contains("JsonRpcRequest"));
967    }
968
969    #[test]
970    fn test_json_rpc_response_serialization() {
971        let response = JsonRpcResponse {
972            jsonrpc: "2.0".to_string(),
973            id: Some(json!(1)),
974            result: Some(json!({"status": "ok"})),
975            error: None,
976        };
977
978        let json_str = serde_json::to_string(&response).unwrap();
979        assert!(json_str.contains("\"jsonrpc\":\"2.0\""));
980        assert!(!json_str.contains("error"));
981    }
982
983    #[test]
984    fn test_json_rpc_error_serialization() {
985        let response = JsonRpcResponse {
986            jsonrpc: "2.0".to_string(),
987            id: Some(json!(1)),
988            result: None,
989            error: Some(JsonRpcError {
990                code: -32600,
991                message: "Invalid request".to_string(),
992                data: None,
993            }),
994        };
995
996        let json_str = serde_json::to_string(&response).unwrap();
997        assert!(json_str.contains("error"));
998        assert!(json_str.contains("-32600"));
999        assert!(!json_str.contains("result"));
1000    }
1001
1002    #[test]
1003    fn test_json_rpc_error_with_data() {
1004        let error = JsonRpcError {
1005            code: -32000,
1006            message: "Server error".to_string(),
1007            data: Some(json!({"details": "additional info"})),
1008        };
1009
1010        let json_str = serde_json::to_string(&error).unwrap();
1011        assert!(json_str.contains("details"));
1012    }
1013
1014    #[test]
1015    fn test_tool_struct_serialization() {
1016        let tool = Tool {
1017            name: "test_tool".to_string(),
1018            description: "A test tool".to_string(),
1019            input_schema: json!({"type": "object"}),
1020        };
1021
1022        let json_str = serde_json::to_string(&tool).unwrap();
1023        assert!(json_str.contains("test_tool"));
1024        assert!(json_str.contains("inputSchema"));
1025    }
1026
1027    #[test]
1028    fn test_handle_request_tools_call_with_scan() {
1029        let temp_dir = TempDir::new().unwrap();
1030        let test_file = temp_dir.path().join("SKILL.md");
1031        std::fs::write(&test_file, "safe content").unwrap();
1032
1033        let server = McpServer::new();
1034        let request = JsonRpcRequest {
1035            jsonrpc: "2.0".to_string(),
1036            id: Some(json!(5)),
1037            method: "tools/call".to_string(),
1038            params: Some(json!({
1039                "name": "scan",
1040                "arguments": {
1041                    "path": test_file.display().to_string()
1042                }
1043            })),
1044        };
1045
1046        let response = server.handle_request(request);
1047        assert!(response.result.is_some());
1048    }
1049}
cc_audit/mcp_server.rs

cc_audit/
mcp_server.rs
