cc_audit/
mcp_server.rs

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

cc_audit/
mcp_server.rs
