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

cc_audit/
mcp_server.rs
