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                    elapsed_ms: 0,
311                };
312                Ok(json!({
313                    "content": [{
314                        "type": "text",
315                        "text": serde_json::to_string_pretty(&result).unwrap()
316                    }]
317                }))
318            }
319            Err(e) => Err(JsonRpcError {
320                code: -32000,
321                message: format!("Scan failed: {}", e),
322                data: None,
323            }),
324        }
325    }
326
327    fn tool_scan_content(&self, args: &Value) -> std::result::Result<Value, JsonRpcError> {
328        let content = args
329            .get("content")
330            .and_then(|v| v.as_str())
331            .ok_or_else(|| JsonRpcError {
332                code: -32602,
333                message: "Missing 'content' argument".to_string(),
334                data: None,
335            })?;
336
337        let filename = args
338            .get("filename")
339            .and_then(|v| v.as_str())
340            .unwrap_or("content.md");
341
342        let config = ScannerConfig::new();
343        let findings = config.check_content(content, filename);
344
345        let summary = Summary::from_findings(&findings);
346        let risk_score = RiskScore::from_findings(&findings);
347
348        Ok(json!({
349            "content": [{
350                "type": "text",
351                "text": serde_json::to_string_pretty(&json!({
352                    "findings": findings,
353                    "summary": summary,
354                    "risk_score": risk_score
355                })).unwrap()
356            }]
357        }))
358    }
359
360    fn tool_check_rule(&self, args: &Value) -> std::result::Result<Value, JsonRpcError> {
361        let rule_id = args
362            .get("rule_id")
363            .and_then(|v| v.as_str())
364            .ok_or_else(|| JsonRpcError {
365                code: -32602,
366                message: "Missing 'rule_id' argument".to_string(),
367                data: None,
368            })?;
369
370        let content = args
371            .get("content")
372            .and_then(|v| v.as_str())
373            .ok_or_else(|| JsonRpcError {
374                code: -32602,
375                message: "Missing 'content' argument".to_string(),
376                data: None,
377            })?;
378
379        // Check if rule exists
380        let Some(rule) = self.rule_engine.get_rule(rule_id) else {
381            return Ok(json!({
382                "content": [{
383                    "type": "text",
384                    "text": format!("Rule '{}' not found", rule_id)
385                }]
386            }));
387        };
388
389        // Check if any pattern matches
390        let mut matches = false;
391        for pattern in &rule.patterns {
392            if pattern.is_match(content) {
393                matches = true;
394                break;
395            }
396        }
397
398        Ok(json!({
399            "content": [{
400                "type": "text",
401                "text": serde_json::to_string_pretty(&json!({
402                    "rule_id": rule_id,
403                    "rule_name": rule.name,
404                    "severity": format!("{:?}", rule.severity),
405                    "matches": matches,
406                    "message": if matches {
407                        format!("Content matches rule: {}", rule.message)
408                    } else {
409                        "No match found".to_string()
410                    }
411                })).unwrap()
412            }]
413        }))
414    }
415
416    fn tool_list_rules(&self, args: &Value) -> std::result::Result<Value, JsonRpcError> {
417        let category_filter = args
418            .get("category")
419            .and_then(|v| v.as_str())
420            .map(|s| s.to_lowercase());
421
422        let rules = self.rule_engine.get_all_rules();
423        let filtered: Vec<_> = rules
424            .iter()
425            .filter(|r| {
426                if let Some(ref cat) = category_filter {
427                    format!("{:?}", r.category).to_lowercase().contains(cat)
428                } else {
429                    true
430                }
431            })
432            .map(|r| {
433                json!({
434                    "id": r.id,
435                    "name": r.name,
436                    "severity": format!("{:?}", r.severity),
437                    "category": format!("{:?}", r.category),
438                    "confidence": format!("{:?}", r.confidence)
439                })
440            })
441            .collect();
442
443        Ok(json!({
444            "content": [{
445                "type": "text",
446                "text": serde_json::to_string_pretty(&json!({
447                    "total": filtered.len(),
448                    "rules": filtered
449                })).unwrap()
450            }]
451        }))
452    }
453
454    fn tool_get_fix_suggestion(&self, args: &Value) -> std::result::Result<Value, JsonRpcError> {
455        let finding_id = args
456            .get("finding_id")
457            .and_then(|v| v.as_str())
458            .ok_or_else(|| JsonRpcError {
459                code: -32602,
460                message: "Missing 'finding_id' argument".to_string(),
461                data: None,
462            })?;
463
464        let code = args
465            .get("code")
466            .and_then(|v| v.as_str())
467            .ok_or_else(|| JsonRpcError {
468                code: -32602,
469                message: "Missing 'code' argument".to_string(),
470                data: None,
471            })?;
472
473        // Create a mock finding for the fixer
474        let Some(rule) = self.rule_engine.get_rule(finding_id) else {
475            return Ok(json!({
476                "content": [{
477                    "type": "text",
478                    "text": format!("No fix suggestion available for rule '{}'", finding_id)
479                }]
480            }));
481        };
482        let finding = Finding {
483            id: finding_id.to_string(),
484            severity: rule.severity,
485            category: rule.category,
486            confidence: rule.confidence,
487            name: rule.name.to_string(),
488            location: crate::rules::Location {
489                file: "virtual".to_string(),
490                line: 1,
491                column: None,
492            },
493            code: code.to_string(),
494            message: rule.message.to_string(),
495            recommendation: rule.recommendation.to_string(),
496            fix_hint: rule.fix_hint.map(|s| s.to_string()),
497            cwe_ids: rule.cwe_ids.iter().map(|s| s.to_string()).collect(),
498            rule_severity: None,
499            client: None,
500            context: None,
501        };
502
503        let fixer = AutoFixer::new(true);
504        let fixes = fixer.generate_fixes(&[finding]);
505
506        if fixes.is_empty() {
507            Ok(json!({
508                "content": [{
509                    "type": "text",
510                    "text": format!("No automatic fix available for {}. Manual review recommended.\n\nRecommendation: {}", finding_id, rule.recommendation)
511                }]
512            }))
513        } else {
514            let fix = &fixes[0];
515            Ok(json!({
516                "content": [{
517                    "type": "text",
518                    "text": serde_json::to_string_pretty(&json!({
519                        "has_fix": true,
520                        "description": fix.description,
521                        "original": fix.original,
522                        "replacement": fix.replacement
523                    })).unwrap()
524                }]
525            }))
526        }
527    }
528}
529
530impl Default for McpServer {
531    fn default() -> Self {
532        Self::new()
533    }
534}
535
536#[cfg(test)]
537mod tests {
538    use super::*;
539    use tempfile::TempDir;
540
541    #[test]
542    fn test_mcp_server_new() {
543        let server = McpServer::new();
544        assert!(!server.rule_engine.get_all_rules().is_empty());
545    }
546
547    #[test]
548    fn test_mcp_server_default() {
549        let server = McpServer::default();
550        assert!(!server.rule_engine.get_all_rules().is_empty());
551    }
552
553    #[test]
554    fn test_handle_initialize() {
555        let server = McpServer::new();
556        let result = server.handle_initialize(&None).unwrap();
557
558        assert!(result.get("protocolVersion").is_some());
559        assert!(result.get("serverInfo").is_some());
560    }
561
562    #[test]
563    fn test_handle_initialize_with_params() {
564        let server = McpServer::new();
565        let params = Some(json!({"clientInfo": {"name": "test"}}));
566        let result = server.handle_initialize(&params).unwrap();
567
568        assert!(result.get("protocolVersion").is_some());
569    }
570
571    #[test]
572    fn test_handle_list_tools() {
573        let server = McpServer::new();
574        let result = server.handle_list_tools().unwrap();
575
576        let tools = result.get("tools").unwrap().as_array().unwrap();
577        assert_eq!(tools.len(), 5);
578
579        let tool_names: Vec<&str> = tools
580            .iter()
581            .map(|t| t.get("name").unwrap().as_str().unwrap())
582            .collect();
583        assert!(tool_names.contains(&"scan"));
584        assert!(tool_names.contains(&"scan_content"));
585        assert!(tool_names.contains(&"check_rule"));
586        assert!(tool_names.contains(&"list_rules"));
587        assert!(tool_names.contains(&"get_fix_suggestion"));
588    }
589
590    #[test]
591    fn test_tool_scan_content() {
592        let server = McpServer::new();
593        let args = json!({
594            "content": "allowed-tools: *",
595            "filename": "test.md"
596        });
597
598        let result = server.tool_scan_content(&args).unwrap();
599        let content = result.get("content").unwrap().as_array().unwrap();
600        assert!(!content.is_empty());
601    }
602
603    #[test]
604    fn test_tool_scan_content_no_filename() {
605        let server = McpServer::new();
606        let args = json!({
607            "content": "some safe content"
608        });
609
610        let result = server.tool_scan_content(&args).unwrap();
611        let content = result.get("content").unwrap().as_array().unwrap();
612        assert!(!content.is_empty());
613    }
614
615    #[test]
616    fn test_tool_scan_content_missing_content() {
617        let server = McpServer::new();
618        let args = json!({});
619
620        let result = server.tool_scan_content(&args);
621        assert!(result.is_err());
622    }
623
624    #[test]
625    fn test_tool_list_rules() {
626        let server = McpServer::new();
627        let args = json!({});
628
629        let result = server.tool_list_rules(&args).unwrap();
630        let content = result.get("content").unwrap().as_array().unwrap();
631        assert!(!content.is_empty());
632    }
633
634    #[test]
635    fn test_tool_list_rules_with_category() {
636        let server = McpServer::new();
637        let args = json!({"category": "exfiltration"});
638
639        let result = server.tool_list_rules(&args).unwrap();
640        let content = result.get("content").unwrap().as_array().unwrap();
641        assert!(!content.is_empty());
642    }
643
644    #[test]
645    fn test_tool_check_rule() {
646        let server = McpServer::new();
647        let args = json!({
648            "rule_id": "OP-001",
649            "content": "allowed-tools: *"
650        });
651
652        let result = server.tool_check_rule(&args).unwrap();
653        let content = result.get("content").unwrap().as_array().unwrap();
654        assert!(!content.is_empty());
655    }
656
657    #[test]
658    fn test_tool_check_rule_no_match() {
659        let server = McpServer::new();
660        let args = json!({
661            "rule_id": "OP-001",
662            "content": "allowed-tools: Read, Write"
663        });
664
665        let result = server.tool_check_rule(&args).unwrap();
666        let content = result.get("content").unwrap().as_array().unwrap();
667        let text = content[0].get("text").unwrap().as_str().unwrap();
668        assert!(text.contains("No match found") || text.contains("matches"));
669    }
670
671    #[test]
672    fn test_tool_check_rule_not_found() {
673        let server = McpServer::new();
674        let args = json!({
675            "rule_id": "NONEXISTENT-001",
676            "content": "some content"
677        });
678
679        let result = server.tool_check_rule(&args).unwrap();
680        let content = result.get("content").unwrap().as_array().unwrap();
681        let text = content[0].get("text").unwrap().as_str().unwrap();
682        assert!(text.contains("not found"));
683    }
684
685    #[test]
686    fn test_tool_check_rule_missing_rule_id() {
687        let server = McpServer::new();
688        let args = json!({
689            "content": "some content"
690        });
691
692        let result = server.tool_check_rule(&args);
693        assert!(result.is_err());
694    }
695
696    #[test]
697    fn test_tool_check_rule_missing_content() {
698        let server = McpServer::new();
699        let args = json!({
700            "rule_id": "OP-001"
701        });
702
703        let result = server.tool_check_rule(&args);
704        assert!(result.is_err());
705    }
706
707    #[test]
708    fn test_tool_scan_valid_path() {
709        let temp_dir = TempDir::new().unwrap();
710        let test_file = temp_dir.path().join("SKILL.md");
711        std::fs::write(&test_file, "---\nallowed-tools: *\n---\n").unwrap();
712
713        let server = McpServer::new();
714        let args = json!({"path": test_file.display().to_string()});
715
716        let result = server.tool_scan(&args).unwrap();
717        let content = result.get("content").unwrap().as_array().unwrap();
718        assert!(!content.is_empty());
719    }
720
721    #[test]
722    fn test_tool_scan_invalid_path() {
723        let server = McpServer::new();
724        let args = json!({"path": "/nonexistent/path/that/does/not/exist"});
725
726        let result = server.tool_scan(&args);
727        assert!(result.is_err());
728    }
729
730    #[test]
731    fn test_tool_scan_missing_path() {
732        let server = McpServer::new();
733        let args = json!({});
734
735        let result = server.tool_scan(&args);
736        assert!(result.is_err());
737    }
738
739    #[test]
740    fn test_tool_get_fix_suggestion_valid() {
741        let server = McpServer::new();
742        let args = json!({
743            "finding_id": "OP-001",
744            "code": "allowed-tools: *"
745        });
746
747        let result = server.tool_get_fix_suggestion(&args).unwrap();
748        let content = result.get("content").unwrap().as_array().unwrap();
749        assert!(!content.is_empty());
750    }
751
752    #[test]
753    fn test_tool_get_fix_suggestion_no_fix_available() {
754        let server = McpServer::new();
755        let args = json!({
756            "finding_id": "EX-001",
757            "code": "echo hello"
758        });
759
760        let result = server.tool_get_fix_suggestion(&args).unwrap();
761        let content = result.get("content").unwrap().as_array().unwrap();
762        let text = content[0].get("text").unwrap().as_str().unwrap();
763        assert!(text.contains("No automatic fix") || text.contains("has_fix"));
764    }
765
766    #[test]
767    fn test_tool_get_fix_suggestion_rule_not_found() {
768        let server = McpServer::new();
769        let args = json!({
770            "finding_id": "NONEXISTENT-001",
771            "code": "some code"
772        });
773
774        let result = server.tool_get_fix_suggestion(&args).unwrap();
775        let content = result.get("content").unwrap().as_array().unwrap();
776        let text = content[0].get("text").unwrap().as_str().unwrap();
777        assert!(text.contains("No fix suggestion available"));
778    }
779
780    #[test]
781    fn test_tool_get_fix_suggestion_missing_finding_id() {
782        let server = McpServer::new();
783        let args = json!({
784            "code": "some code"
785        });
786
787        let result = server.tool_get_fix_suggestion(&args);
788        assert!(result.is_err());
789    }
790
791    #[test]
792    fn test_tool_get_fix_suggestion_missing_code() {
793        let server = McpServer::new();
794        let args = json!({
795            "finding_id": "OP-001"
796        });
797
798        let result = server.tool_get_fix_suggestion(&args);
799        assert!(result.is_err());
800    }
801
802    #[test]
803    fn test_handle_request_initialize() {
804        let server = McpServer::new();
805        let request = JsonRpcRequest {
806            jsonrpc: "2.0".to_string(),
807            id: Some(json!(1)),
808            method: "initialize".to_string(),
809            params: None,
810        };
811
812        let response = server.handle_request(request);
813        assert!(response.result.is_some());
814        assert!(response.error.is_none());
815    }
816
817    #[test]
818    fn test_handle_request_tools_list() {
819        let server = McpServer::new();
820        let request = JsonRpcRequest {
821            jsonrpc: "2.0".to_string(),
822            id: Some(json!(2)),
823            method: "tools/list".to_string(),
824            params: None,
825        };
826
827        let response = server.handle_request(request);
828        assert!(response.result.is_some());
829        assert!(response.error.is_none());
830    }
831
832    #[test]
833    fn test_handle_request_shutdown() {
834        let server = McpServer::new();
835        let request = JsonRpcRequest {
836            jsonrpc: "2.0".to_string(),
837            id: Some(json!(3)),
838            method: "shutdown".to_string(),
839            params: None,
840        };
841
842        let response = server.handle_request(request);
843        assert!(response.result.is_some());
844        assert!(response.error.is_none());
845    }
846
847    #[test]
848    fn test_handle_request_unknown_method() {
849        let server = McpServer::new();
850        let request = JsonRpcRequest {
851            jsonrpc: "2.0".to_string(),
852            id: Some(json!(4)),
853            method: "unknown/method".to_string(),
854            params: None,
855        };
856
857        let response = server.handle_request(request);
858        assert!(response.result.is_none());
859        assert!(response.error.is_some());
860        assert_eq!(response.error.as_ref().unwrap().code, -32601);
861    }
862
863    #[test]
864    fn test_handle_tool_call_missing_params() {
865        let server = McpServer::new();
866        let result = server.handle_tool_call(&None);
867        assert!(result.is_err());
868        assert_eq!(result.unwrap_err().code, -32602);
869    }
870
871    #[test]
872    fn test_handle_tool_call_missing_name() {
873        let server = McpServer::new();
874        let params = Some(json!({"arguments": {}}));
875        let result = server.handle_tool_call(&params);
876        assert!(result.is_err());
877    }
878
879    #[test]
880    fn test_handle_tool_call_unknown_tool() {
881        let server = McpServer::new();
882        let params = Some(json!({
883            "name": "unknown_tool",
884            "arguments": {}
885        }));
886        let result = server.handle_tool_call(&params);
887        assert!(result.is_err());
888        assert!(result.unwrap_err().message.contains("Unknown tool"));
889    }
890
891    #[test]
892    fn test_handle_tool_call_scan_content() {
893        let server = McpServer::new();
894        let params = Some(json!({
895            "name": "scan_content",
896            "arguments": {
897                "content": "safe content"
898            }
899        }));
900
901        let result = server.handle_tool_call(&params);
902        assert!(result.is_ok());
903    }
904
905    #[test]
906    fn test_handle_tool_call_list_rules() {
907        let server = McpServer::new();
908        let params = Some(json!({
909            "name": "list_rules",
910            "arguments": {}
911        }));
912
913        let result = server.handle_tool_call(&params);
914        assert!(result.is_ok());
915    }
916
917    #[test]
918    fn test_handle_tool_call_check_rule() {
919        let server = McpServer::new();
920        let params = Some(json!({
921            "name": "check_rule",
922            "arguments": {
923                "rule_id": "OP-001",
924                "content": "allowed-tools: *"
925            }
926        }));
927
928        let result = server.handle_tool_call(&params);
929        assert!(result.is_ok());
930    }
931
932    #[test]
933    fn test_handle_tool_call_get_fix_suggestion() {
934        let server = McpServer::new();
935        let params = Some(json!({
936            "name": "get_fix_suggestion",
937            "arguments": {
938                "finding_id": "OP-001",
939                "code": "allowed-tools: *"
940            }
941        }));
942
943        let result = server.handle_tool_call(&params);
944        assert!(result.is_ok());
945    }
946
947    #[test]
948    fn test_json_rpc_request_debug() {
949        let request = JsonRpcRequest {
950            jsonrpc: "2.0".to_string(),
951            id: Some(json!(1)),
952            method: "test".to_string(),
953            params: None,
954        };
955
956        let debug_str = format!("{:?}", request);
957        assert!(debug_str.contains("JsonRpcRequest"));
958    }
959
960    #[test]
961    fn test_json_rpc_response_serialization() {
962        let response = JsonRpcResponse {
963            jsonrpc: "2.0".to_string(),
964            id: Some(json!(1)),
965            result: Some(json!({"status": "ok"})),
966            error: None,
967        };
968
969        let json_str = serde_json::to_string(&response).unwrap();
970        assert!(json_str.contains("\"jsonrpc\":\"2.0\""));
971        assert!(!json_str.contains("error"));
972    }
973
974    #[test]
975    fn test_json_rpc_error_serialization() {
976        let response = JsonRpcResponse {
977            jsonrpc: "2.0".to_string(),
978            id: Some(json!(1)),
979            result: None,
980            error: Some(JsonRpcError {
981                code: -32600,
982                message: "Invalid request".to_string(),
983                data: None,
984            }),
985        };
986
987        let json_str = serde_json::to_string(&response).unwrap();
988        assert!(json_str.contains("error"));
989        assert!(json_str.contains("-32600"));
990        assert!(!json_str.contains("result"));
991    }
992
993    #[test]
994    fn test_json_rpc_error_with_data() {
995        let error = JsonRpcError {
996            code: -32000,
997            message: "Server error".to_string(),
998            data: Some(json!({"details": "additional info"})),
999        };
1000
1001        let json_str = serde_json::to_string(&error).unwrap();
1002        assert!(json_str.contains("details"));
1003    }
1004
1005    #[test]
1006    fn test_tool_struct_serialization() {
1007        let tool = Tool {
1008            name: "test_tool".to_string(),
1009            description: "A test tool".to_string(),
1010            input_schema: json!({"type": "object"}),
1011        };
1012
1013        let json_str = serde_json::to_string(&tool).unwrap();
1014        assert!(json_str.contains("test_tool"));
1015        assert!(json_str.contains("inputSchema"));
1016    }
1017
1018    #[test]
1019    fn test_handle_request_tools_call_with_scan() {
1020        let temp_dir = TempDir::new().unwrap();
1021        let test_file = temp_dir.path().join("SKILL.md");
1022        std::fs::write(&test_file, "safe content").unwrap();
1023
1024        let server = McpServer::new();
1025        let request = JsonRpcRequest {
1026            jsonrpc: "2.0".to_string(),
1027            id: Some(json!(5)),
1028            method: "tools/call".to_string(),
1029            params: Some(json!({
1030                "name": "scan",
1031                "arguments": {
1032                    "path": test_file.display().to_string()
1033                }
1034            })),
1035        };
1036
1037        let response = server.handle_request(request);
1038        assert!(response.result.is_some());
1039    }
1040}
cc_audit/mcp_server.rs

cc_audit/
mcp_server.rs
