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        };
505
506        let fixer = AutoFixer::new(true);
507        let fixes = fixer.generate_fixes(&[finding]);
508
509        if fixes.is_empty() {
510            Ok(json!({
511                "content": [{
512                    "type": "text",
513                    "text": format!("No automatic fix available for {}. Manual review recommended.\n\nRecommendation: {}", finding_id, rule.recommendation)
514                }]
515            }))
516        } else {
517            let fix = &fixes[0];
518            Ok(json!({
519                "content": [{
520                    "type": "text",
521                    "text": serde_json::to_string_pretty(&json!({
522                        "has_fix": true,
523                        "description": fix.description,
524                        "original": fix.original,
525                        "replacement": fix.replacement
526                    })).unwrap()
527                }]
528            }))
529        }
530    }
531}
532
533impl Default for McpServer {
534    fn default() -> Self {
535        Self::new()
536    }
537}
538
539#[cfg(test)]
540mod tests {
541    use super::*;
542    use tempfile::TempDir;
543
544    #[test]
545    fn test_mcp_server_new() {
546        let server = McpServer::new();
547        assert!(!server.rule_engine.get_all_rules().is_empty());
548    }
549
550    #[test]
551    fn test_mcp_server_default() {
552        let server = McpServer::default();
553        assert!(!server.rule_engine.get_all_rules().is_empty());
554    }
555
556    #[test]
557    fn test_handle_initialize() {
558        let server = McpServer::new();
559        let result = server.handle_initialize(&None).unwrap();
560
561        assert!(result.get("protocolVersion").is_some());
562        assert!(result.get("serverInfo").is_some());
563    }
564
565    #[test]
566    fn test_handle_initialize_with_params() {
567        let server = McpServer::new();
568        let params = Some(json!({"clientInfo": {"name": "test"}}));
569        let result = server.handle_initialize(&params).unwrap();
570
571        assert!(result.get("protocolVersion").is_some());
572    }
573
574    #[test]
575    fn test_handle_list_tools() {
576        let server = McpServer::new();
577        let result = server.handle_list_tools().unwrap();
578
579        let tools = result.get("tools").unwrap().as_array().unwrap();
580        assert_eq!(tools.len(), 5);
581
582        let tool_names: Vec<&str> = tools
583            .iter()
584            .map(|t| t.get("name").unwrap().as_str().unwrap())
585            .collect();
586        assert!(tool_names.contains(&"scan"));
587        assert!(tool_names.contains(&"scan_content"));
588        assert!(tool_names.contains(&"check_rule"));
589        assert!(tool_names.contains(&"list_rules"));
590        assert!(tool_names.contains(&"get_fix_suggestion"));
591    }
592
593    #[test]
594    fn test_tool_scan_content() {
595        let server = McpServer::new();
596        let args = json!({
597            "content": "allowed-tools: *",
598            "filename": "test.md"
599        });
600
601        let result = server.tool_scan_content(&args).unwrap();
602        let content = result.get("content").unwrap().as_array().unwrap();
603        assert!(!content.is_empty());
604    }
605
606    #[test]
607    fn test_tool_scan_content_no_filename() {
608        let server = McpServer::new();
609        let args = json!({
610            "content": "some safe content"
611        });
612
613        let result = server.tool_scan_content(&args).unwrap();
614        let content = result.get("content").unwrap().as_array().unwrap();
615        assert!(!content.is_empty());
616    }
617
618    #[test]
619    fn test_tool_scan_content_missing_content() {
620        let server = McpServer::new();
621        let args = json!({});
622
623        let result = server.tool_scan_content(&args);
624        assert!(result.is_err());
625    }
626
627    #[test]
628    fn test_tool_list_rules() {
629        let server = McpServer::new();
630        let args = json!({});
631
632        let result = server.tool_list_rules(&args).unwrap();
633        let content = result.get("content").unwrap().as_array().unwrap();
634        assert!(!content.is_empty());
635    }
636
637    #[test]
638    fn test_tool_list_rules_with_category() {
639        let server = McpServer::new();
640        let args = json!({"category": "exfiltration"});
641
642        let result = server.tool_list_rules(&args).unwrap();
643        let content = result.get("content").unwrap().as_array().unwrap();
644        assert!(!content.is_empty());
645    }
646
647    #[test]
648    fn test_tool_check_rule() {
649        let server = McpServer::new();
650        let args = json!({
651            "rule_id": "OP-001",
652            "content": "allowed-tools: *"
653        });
654
655        let result = server.tool_check_rule(&args).unwrap();
656        let content = result.get("content").unwrap().as_array().unwrap();
657        assert!(!content.is_empty());
658    }
659
660    #[test]
661    fn test_tool_check_rule_no_match() {
662        let server = McpServer::new();
663        let args = json!({
664            "rule_id": "OP-001",
665            "content": "allowed-tools: Read, Write"
666        });
667
668        let result = server.tool_check_rule(&args).unwrap();
669        let content = result.get("content").unwrap().as_array().unwrap();
670        let text = content[0].get("text").unwrap().as_str().unwrap();
671        assert!(text.contains("No match found") || text.contains("matches"));
672    }
673
674    #[test]
675    fn test_tool_check_rule_not_found() {
676        let server = McpServer::new();
677        let args = json!({
678            "rule_id": "NONEXISTENT-001",
679            "content": "some content"
680        });
681
682        let result = server.tool_check_rule(&args).unwrap();
683        let content = result.get("content").unwrap().as_array().unwrap();
684        let text = content[0].get("text").unwrap().as_str().unwrap();
685        assert!(text.contains("not found"));
686    }
687
688    #[test]
689    fn test_tool_check_rule_missing_rule_id() {
690        let server = McpServer::new();
691        let args = json!({
692            "content": "some content"
693        });
694
695        let result = server.tool_check_rule(&args);
696        assert!(result.is_err());
697    }
698
699    #[test]
700    fn test_tool_check_rule_missing_content() {
701        let server = McpServer::new();
702        let args = json!({
703            "rule_id": "OP-001"
704        });
705
706        let result = server.tool_check_rule(&args);
707        assert!(result.is_err());
708    }
709
710    #[test]
711    fn test_tool_scan_valid_path() {
712        let temp_dir = TempDir::new().unwrap();
713        let test_file = temp_dir.path().join("SKILL.md");
714        std::fs::write(&test_file, "---\nallowed-tools: *\n---\n").unwrap();
715
716        let server = McpServer::new();
717        let args = json!({"path": test_file.display().to_string()});
718
719        let result = server.tool_scan(&args).unwrap();
720        let content = result.get("content").unwrap().as_array().unwrap();
721        assert!(!content.is_empty());
722    }
723
724    #[test]
725    fn test_tool_scan_invalid_path() {
726        let server = McpServer::new();
727        let args = json!({"path": "/nonexistent/path/that/does/not/exist"});
728
729        let result = server.tool_scan(&args);
730        assert!(result.is_err());
731    }
732
733    #[test]
734    fn test_tool_scan_missing_path() {
735        let server = McpServer::new();
736        let args = json!({});
737
738        let result = server.tool_scan(&args);
739        assert!(result.is_err());
740    }
741
742    #[test]
743    fn test_tool_get_fix_suggestion_valid() {
744        let server = McpServer::new();
745        let args = json!({
746            "finding_id": "OP-001",
747            "code": "allowed-tools: *"
748        });
749
750        let result = server.tool_get_fix_suggestion(&args).unwrap();
751        let content = result.get("content").unwrap().as_array().unwrap();
752        assert!(!content.is_empty());
753    }
754
755    #[test]
756    fn test_tool_get_fix_suggestion_no_fix_available() {
757        let server = McpServer::new();
758        let args = json!({
759            "finding_id": "EX-001",
760            "code": "echo hello"
761        });
762
763        let result = server.tool_get_fix_suggestion(&args).unwrap();
764        let content = result.get("content").unwrap().as_array().unwrap();
765        let text = content[0].get("text").unwrap().as_str().unwrap();
766        assert!(text.contains("No automatic fix") || text.contains("has_fix"));
767    }
768
769    #[test]
770    fn test_tool_get_fix_suggestion_rule_not_found() {
771        let server = McpServer::new();
772        let args = json!({
773            "finding_id": "NONEXISTENT-001",
774            "code": "some code"
775        });
776
777        let result = server.tool_get_fix_suggestion(&args).unwrap();
778        let content = result.get("content").unwrap().as_array().unwrap();
779        let text = content[0].get("text").unwrap().as_str().unwrap();
780        assert!(text.contains("No fix suggestion available"));
781    }
782
783    #[test]
784    fn test_tool_get_fix_suggestion_missing_finding_id() {
785        let server = McpServer::new();
786        let args = json!({
787            "code": "some code"
788        });
789
790        let result = server.tool_get_fix_suggestion(&args);
791        assert!(result.is_err());
792    }
793
794    #[test]
795    fn test_tool_get_fix_suggestion_missing_code() {
796        let server = McpServer::new();
797        let args = json!({
798            "finding_id": "OP-001"
799        });
800
801        let result = server.tool_get_fix_suggestion(&args);
802        assert!(result.is_err());
803    }
804
805    #[test]
806    fn test_handle_request_initialize() {
807        let server = McpServer::new();
808        let request = JsonRpcRequest {
809            jsonrpc: "2.0".to_string(),
810            id: Some(json!(1)),
811            method: "initialize".to_string(),
812            params: None,
813        };
814
815        let response = server.handle_request(request);
816        assert!(response.result.is_some());
817        assert!(response.error.is_none());
818    }
819
820    #[test]
821    fn test_handle_request_tools_list() {
822        let server = McpServer::new();
823        let request = JsonRpcRequest {
824            jsonrpc: "2.0".to_string(),
825            id: Some(json!(2)),
826            method: "tools/list".to_string(),
827            params: None,
828        };
829
830        let response = server.handle_request(request);
831        assert!(response.result.is_some());
832        assert!(response.error.is_none());
833    }
834
835    #[test]
836    fn test_handle_request_shutdown() {
837        let server = McpServer::new();
838        let request = JsonRpcRequest {
839            jsonrpc: "2.0".to_string(),
840            id: Some(json!(3)),
841            method: "shutdown".to_string(),
842            params: None,
843        };
844
845        let response = server.handle_request(request);
846        assert!(response.result.is_some());
847        assert!(response.error.is_none());
848    }
849
850    #[test]
851    fn test_handle_request_unknown_method() {
852        let server = McpServer::new();
853        let request = JsonRpcRequest {
854            jsonrpc: "2.0".to_string(),
855            id: Some(json!(4)),
856            method: "unknown/method".to_string(),
857            params: None,
858        };
859
860        let response = server.handle_request(request);
861        assert!(response.result.is_none());
862        assert!(response.error.is_some());
863        assert_eq!(response.error.as_ref().unwrap().code, -32601);
864    }
865
866    #[test]
867    fn test_handle_tool_call_missing_params() {
868        let server = McpServer::new();
869        let result = server.handle_tool_call(&None);
870        assert!(result.is_err());
871        assert_eq!(result.unwrap_err().code, -32602);
872    }
873
874    #[test]
875    fn test_handle_tool_call_missing_name() {
876        let server = McpServer::new();
877        let params = Some(json!({"arguments": {}}));
878        let result = server.handle_tool_call(&params);
879        assert!(result.is_err());
880    }
881
882    #[test]
883    fn test_handle_tool_call_unknown_tool() {
884        let server = McpServer::new();
885        let params = Some(json!({
886            "name": "unknown_tool",
887            "arguments": {}
888        }));
889        let result = server.handle_tool_call(&params);
890        assert!(result.is_err());
891        assert!(result.unwrap_err().message.contains("Unknown tool"));
892    }
893
894    #[test]
895    fn test_handle_tool_call_scan_content() {
896        let server = McpServer::new();
897        let params = Some(json!({
898            "name": "scan_content",
899            "arguments": {
900                "content": "safe content"
901            }
902        }));
903
904        let result = server.handle_tool_call(&params);
905        assert!(result.is_ok());
906    }
907
908    #[test]
909    fn test_handle_tool_call_list_rules() {
910        let server = McpServer::new();
911        let params = Some(json!({
912            "name": "list_rules",
913            "arguments": {}
914        }));
915
916        let result = server.handle_tool_call(&params);
917        assert!(result.is_ok());
918    }
919
920    #[test]
921    fn test_handle_tool_call_check_rule() {
922        let server = McpServer::new();
923        let params = Some(json!({
924            "name": "check_rule",
925            "arguments": {
926                "rule_id": "OP-001",
927                "content": "allowed-tools: *"
928            }
929        }));
930
931        let result = server.handle_tool_call(&params);
932        assert!(result.is_ok());
933    }
934
935    #[test]
936    fn test_handle_tool_call_get_fix_suggestion() {
937        let server = McpServer::new();
938        let params = Some(json!({
939            "name": "get_fix_suggestion",
940            "arguments": {
941                "finding_id": "OP-001",
942                "code": "allowed-tools: *"
943            }
944        }));
945
946        let result = server.handle_tool_call(&params);
947        assert!(result.is_ok());
948    }
949
950    #[test]
951    fn test_json_rpc_request_debug() {
952        let request = JsonRpcRequest {
953            jsonrpc: "2.0".to_string(),
954            id: Some(json!(1)),
955            method: "test".to_string(),
956            params: None,
957        };
958
959        let debug_str = format!("{:?}", request);
960        assert!(debug_str.contains("JsonRpcRequest"));
961    }
962
963    #[test]
964    fn test_json_rpc_response_serialization() {
965        let response = JsonRpcResponse {
966            jsonrpc: "2.0".to_string(),
967            id: Some(json!(1)),
968            result: Some(json!({"status": "ok"})),
969            error: None,
970        };
971
972        let json_str = serde_json::to_string(&response).unwrap();
973        assert!(json_str.contains("\"jsonrpc\":\"2.0\""));
974        assert!(!json_str.contains("error"));
975    }
976
977    #[test]
978    fn test_json_rpc_error_serialization() {
979        let response = JsonRpcResponse {
980            jsonrpc: "2.0".to_string(),
981            id: Some(json!(1)),
982            result: None,
983            error: Some(JsonRpcError {
984                code: -32600,
985                message: "Invalid request".to_string(),
986                data: None,
987            }),
988        };
989
990        let json_str = serde_json::to_string(&response).unwrap();
991        assert!(json_str.contains("error"));
992        assert!(json_str.contains("-32600"));
993        assert!(!json_str.contains("result"));
994    }
995
996    #[test]
997    fn test_json_rpc_error_with_data() {
998        let error = JsonRpcError {
999            code: -32000,
1000            message: "Server error".to_string(),
1001            data: Some(json!({"details": "additional info"})),
1002        };
1003
1004        let json_str = serde_json::to_string(&error).unwrap();
1005        assert!(json_str.contains("details"));
1006    }
1007
1008    #[test]
1009    fn test_tool_struct_serialization() {
1010        let tool = Tool {
1011            name: "test_tool".to_string(),
1012            description: "A test tool".to_string(),
1013            input_schema: json!({"type": "object"}),
1014        };
1015
1016        let json_str = serde_json::to_string(&tool).unwrap();
1017        assert!(json_str.contains("test_tool"));
1018        assert!(json_str.contains("inputSchema"));
1019    }
1020
1021    #[test]
1022    fn test_handle_request_tools_call_with_scan() {
1023        let temp_dir = TempDir::new().unwrap();
1024        let test_file = temp_dir.path().join("SKILL.md");
1025        std::fs::write(&test_file, "safe content").unwrap();
1026
1027        let server = McpServer::new();
1028        let request = JsonRpcRequest {
1029            jsonrpc: "2.0".to_string(),
1030            id: Some(json!(5)),
1031            method: "tools/call".to_string(),
1032            params: Some(json!({
1033                "name": "scan",
1034                "arguments": {
1035                    "path": test_file.display().to_string()
1036                }
1037            })),
1038        };
1039
1040        let response = server.handle_request(request);
1041        assert!(response.result.is_some());
1042    }
1043}
cc_audit/mcp_server.rs

cc_audit/
mcp_server.rs
