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

cc_audit/
mcp_server.rs
