mockforge_grpc/dynamic/
proto_parser.rs

1//! Proto file parsing and service discovery
2//!
3//! This module handles parsing of .proto files and extracting service definitions
4//! to generate dynamic gRPC service implementations.
5
6use prost_reflect::DescriptorPool;
7use std::collections::HashMap;
8use std::fs;
9use std::path::{Path, PathBuf};
10use std::process::Command;
11use tempfile::TempDir;
12use tracing::{debug, error, info, warn};
13
14/// A parsed proto service definition
15#[derive(Debug, Clone)]
16pub struct ProtoService {
17    /// The service name (e.g., "mockforge.greeter.Greeter")
18    pub name: String,
19    /// The package name (e.g., "mockforge.greeter")
20    pub package: String,
21    /// The short service name (e.g., "Greeter")
22    pub short_name: String,
23    /// List of methods in this service
24    pub methods: Vec<ProtoMethod>,
25}
26
27/// A parsed proto method definition
28#[derive(Debug, Clone)]
29pub struct ProtoMethod {
30    /// The method name (e.g., "SayHello")
31    pub name: String,
32    /// The input message type
33    pub input_type: String,
34    /// The output message type
35    pub output_type: String,
36    /// Whether this is a client streaming method
37    pub client_streaming: bool,
38    /// Whether this is a server streaming method
39    pub server_streaming: bool,
40}
41
42/// A proto file parser that can extract service definitions
43pub struct ProtoParser {
44    /// The descriptor pool containing parsed proto files
45    pool: DescriptorPool,
46    /// Map of service names to their definitions
47    services: HashMap<String, ProtoService>,
48    /// Include paths for proto compilation
49    include_paths: Vec<PathBuf>,
50    /// Temporary directory for compilation artifacts
51    temp_dir: Option<TempDir>,
52}
53
54impl ProtoParser {
55    /// Create a new proto parser
56    pub fn new() -> Self {
57        Self {
58            pool: DescriptorPool::new(),
59            services: HashMap::new(),
60            include_paths: vec![],
61            temp_dir: None,
62        }
63    }
64
65    /// Create a new proto parser with include paths
66    pub fn with_include_paths(include_paths: Vec<PathBuf>) -> Self {
67        Self {
68            pool: DescriptorPool::new(),
69            services: HashMap::new(),
70            include_paths,
71            temp_dir: None,
72        }
73    }
74
75    /// Parse proto files from a directory
76    pub async fn parse_directory(
77        &mut self,
78        proto_dir: &str,
79    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
80        info!("Parsing proto files from directory: {}", proto_dir);
81
82        let proto_path = Path::new(proto_dir);
83        if !proto_path.exists() {
84            return Err(format!("Proto directory does not exist: {}", proto_dir).into());
85        }
86
87        // Discover all proto files
88        let proto_files = self.discover_proto_files(proto_path)?;
89        if proto_files.is_empty() {
90            warn!("No proto files found in directory: {}", proto_dir);
91            return Ok(());
92        }
93
94        info!("Found {} proto files: {:?}", proto_files.len(), proto_files);
95
96        // Parse each proto file
97        for proto_file in proto_files {
98            if let Err(e) = self.parse_proto_file(&proto_file).await {
99                error!("Failed to parse proto file {}: {}", proto_file, e);
100                // Continue with other files
101            }
102        }
103
104        // Extract services from the descriptor pool only if there are any services in the pool
105        if self.pool.services().count() > 0 {
106            self.extract_services()?;
107        } else {
108            debug!("No services found in descriptor pool, keeping mock services");
109        }
110
111        info!("Successfully parsed {} services", self.services.len());
112        Ok(())
113    }
114
115    /// Discover proto files in a directory recursively
116    #[allow(clippy::only_used_in_recursion)]
117    fn discover_proto_files(
118        &self,
119        dir: &Path,
120    ) -> Result<Vec<String>, Box<dyn std::error::Error + Send + Sync>> {
121        let mut proto_files = Vec::new();
122
123        if let Ok(entries) = fs::read_dir(dir) {
124            for entry in entries.flatten() {
125                let path = entry.path();
126
127                if path.is_dir() {
128                    // Recursively search subdirectories
129                    proto_files.extend(self.discover_proto_files(&path)?);
130                } else if path.extension().and_then(|s| s.to_str()) == Some("proto") {
131                    // Found a .proto file
132                    proto_files.push(path.to_string_lossy().to_string());
133                }
134            }
135        }
136
137        Ok(proto_files)
138    }
139
140    /// Parse a single proto file using protoc compilation
141    async fn parse_proto_file(
142        &mut self,
143        proto_file: &str,
144    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
145        debug!("Parsing proto file: {}", proto_file);
146
147        // Create temporary directory for compilation artifacts if not exists
148        if self.temp_dir.is_none() {
149            self.temp_dir = Some(TempDir::new()?);
150        }
151
152        let temp_dir = self.temp_dir.as_ref().unwrap();
153        let descriptor_path = temp_dir.path().join("descriptors.bin");
154
155        // Try real protoc compilation first
156        match self.compile_with_protoc(proto_file, &descriptor_path).await {
157            Ok(()) => {
158                // Load the compiled descriptor set into the pool
159                let descriptor_bytes = fs::read(&descriptor_path)?;
160                match self.pool.decode_file_descriptor_set(&*descriptor_bytes) {
161                    Ok(()) => {
162                        info!("Successfully compiled and loaded proto file: {}", proto_file);
163                        // Extract services from the descriptor pool if successful
164                        if self.pool.services().count() > 0 {
165                            self.extract_services()?;
166                        }
167                        return Ok(());
168                    }
169                    Err(e) => {
170                        warn!("Failed to decode descriptor set, falling back to mock: {}", e);
171                    }
172                }
173            }
174            Err(e) => {
175                warn!("Failed to compile with protoc, falling back to mock: {}", e);
176            }
177        }
178
179        // Fallback to mock service for testing
180        if proto_file.contains("gretter.proto") || proto_file.contains("greeter.proto") {
181            debug!("Adding mock greeter service for {}", proto_file);
182            self.add_mock_greeter_service();
183        }
184
185        Ok(())
186    }
187
188    /// Compile proto file using protoc
189    async fn compile_with_protoc(
190        &self,
191        proto_file: &str,
192        output_path: &Path,
193    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
194        debug!("Compiling proto file with protoc: {}", proto_file);
195
196        // Build protoc command
197        let mut cmd = Command::new("protoc");
198
199        // Add include paths
200        for include_path in &self.include_paths {
201            cmd.arg("-I").arg(include_path);
202        }
203
204        // Add proto file's directory as include path
205        if let Some(parent_dir) = Path::new(proto_file).parent() {
206            cmd.arg("-I").arg(parent_dir);
207        }
208
209        // Add well-known types include path (common protoc installation paths)
210        let well_known_paths = [
211            "/usr/local/include",
212            "/usr/include",
213            "/opt/homebrew/include",
214        ];
215
216        for path in &well_known_paths {
217            if Path::new(path).exists() {
218                cmd.arg("-I").arg(path);
219            }
220        }
221
222        // Set output path and format
223        cmd.arg("--descriptor_set_out")
224            .arg(output_path)
225            .arg("--include_imports")
226            .arg("--include_source_info")
227            .arg(proto_file);
228
229        debug!("Running protoc command: {:?}", cmd);
230
231        // Execute protoc
232        let output = cmd.output()?;
233
234        if !output.status.success() {
235            let stderr = String::from_utf8_lossy(&output.stderr);
236            return Err(format!("protoc failed: {}", stderr).into());
237        }
238
239        info!("Successfully compiled proto file with protoc: {}", proto_file);
240        Ok(())
241    }
242
243    /// Add a mock greeter service (for demonstration)
244    fn add_mock_greeter_service(&mut self) {
245        let service = ProtoService {
246            name: "mockforge.greeter.Greeter".to_string(),
247            package: "mockforge.greeter".to_string(),
248            short_name: "Greeter".to_string(),
249            methods: vec![
250                ProtoMethod {
251                    name: "SayHello".to_string(),
252                    input_type: "mockforge.greeter.HelloRequest".to_string(),
253                    output_type: "mockforge.greeter.HelloReply".to_string(),
254                    client_streaming: false,
255                    server_streaming: false,
256                },
257                ProtoMethod {
258                    name: "SayHelloStream".to_string(),
259                    input_type: "mockforge.greeter.HelloRequest".to_string(),
260                    output_type: "mockforge.greeter.HelloReply".to_string(),
261                    client_streaming: false,
262                    server_streaming: true,
263                },
264                ProtoMethod {
265                    name: "SayHelloClientStream".to_string(),
266                    input_type: "mockforge.greeter.HelloRequest".to_string(),
267                    output_type: "mockforge.greeter.HelloReply".to_string(),
268                    client_streaming: true,
269                    server_streaming: false,
270                },
271                ProtoMethod {
272                    name: "Chat".to_string(),
273                    input_type: "mockforge.greeter.HelloRequest".to_string(),
274                    output_type: "mockforge.greeter.HelloReply".to_string(),
275                    client_streaming: true,
276                    server_streaming: true,
277                },
278            ],
279        };
280
281        self.services.insert(service.name.clone(), service);
282    }
283
284    /// Extract services from the descriptor pool
285    fn extract_services(&mut self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
286        debug!("Extracting services from descriptor pool");
287
288        // Clear existing services (except mock ones)
289        let mock_services: HashMap<String, ProtoService> = self
290            .services
291            .drain()
292            .filter(|(name, _)| name.contains("mockforge.greeter"))
293            .collect();
294
295        self.services = mock_services;
296
297        // Extract services from the descriptor pool
298        for service_descriptor in self.pool.services() {
299            let service_name = service_descriptor.full_name().to_string();
300            let package_name = service_descriptor.parent_file().package_name().to_string();
301            let short_name = service_descriptor.name().to_string();
302
303            debug!("Found service: {} in package: {}", service_name, package_name);
304
305            // Extract methods for this service
306            let mut methods = Vec::new();
307            for method_descriptor in service_descriptor.methods() {
308                let method = ProtoMethod {
309                    name: method_descriptor.name().to_string(),
310                    input_type: method_descriptor.input().full_name().to_string(),
311                    output_type: method_descriptor.output().full_name().to_string(),
312                    client_streaming: method_descriptor.is_client_streaming(),
313                    server_streaming: method_descriptor.is_server_streaming(),
314                };
315
316                debug!(
317                    "  Found method: {} ({} -> {})",
318                    method.name, method.input_type, method.output_type
319                );
320
321                methods.push(method);
322            }
323
324            let service = ProtoService {
325                name: service_name.clone(),
326                package: package_name,
327                short_name,
328                methods,
329            };
330
331            self.services.insert(service_name, service);
332        }
333
334        info!("Extracted {} services from descriptor pool", self.services.len());
335        Ok(())
336    }
337
338    /// Get all discovered services
339    pub fn services(&self) -> &HashMap<String, ProtoService> {
340        &self.services
341    }
342
343    /// Get a specific service by name
344    pub fn get_service(&self, name: &str) -> Option<&ProtoService> {
345        self.services.get(name)
346    }
347
348    /// Get the descriptor pool
349    pub fn pool(&self) -> &DescriptorPool {
350        &self.pool
351    }
352
353    /// Consume the parser and return the descriptor pool
354    pub fn into_pool(self) -> DescriptorPool {
355        self.pool
356    }
357}
358
359impl Default for ProtoParser {
360    fn default() -> Self {
361        Self::new()
362    }
363}
364
365#[cfg(test)]
366mod tests {
367    use super::*;
368
369    #[tokio::test]
370    async fn test_parse_proto_file() {
371        // Test with the existing greeter.proto file
372        let proto_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap() + "/proto";
373        let proto_path = format!("{}/gretter.proto", proto_dir);
374
375        // Parse the proto file
376        let mut parser = ProtoParser::new();
377        parser.parse_proto_file(&proto_path).await.unwrap();
378
379        // Verify the service was parsed correctly
380        let services = parser.services();
381        assert_eq!(services.len(), 1);
382
383        let service_name = "mockforge.greeter.Greeter";
384        assert!(services.contains_key(service_name));
385
386        let service = &services[service_name];
387        assert_eq!(service.name, service_name);
388        assert_eq!(service.methods.len(), 4); // SayHello, SayHelloStream, SayHelloClientStream, Chat
389
390        // Check SayHello method (unary)
391        let say_hello = service.methods.iter().find(|m| m.name == "SayHello").unwrap();
392        assert_eq!(say_hello.input_type, "mockforge.greeter.HelloRequest");
393        assert_eq!(say_hello.output_type, "mockforge.greeter.HelloReply");
394        assert!(!say_hello.client_streaming);
395        assert!(!say_hello.server_streaming);
396
397        // Check SayHelloStream method (server streaming)
398        let say_hello_stream = service.methods.iter().find(|m| m.name == "SayHelloStream").unwrap();
399        assert!(!say_hello_stream.client_streaming);
400        assert!(say_hello_stream.server_streaming);
401    }
402
403    #[tokio::test]
404    async fn test_parse_directory() {
405        // Test with the existing proto directory
406        let proto_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap() + "/proto";
407
408        // Parse the directory
409        let mut parser = ProtoParser::new();
410        parser.parse_directory(&proto_dir).await.unwrap();
411
412        // Verify services were discovered
413        let services = parser.services();
414        assert_eq!(services.len(), 1);
415
416        let service_name = "mockforge.greeter.Greeter";
417        assert!(services.contains_key(service_name));
418
419        let service = &services[service_name];
420        assert_eq!(service.methods.len(), 4);
421
422        // Check all methods exist
423        let method_names: Vec<&str> = service.methods.iter().map(|m| m.name.as_str()).collect();
424        assert!(method_names.contains(&"SayHello"));
425        assert!(method_names.contains(&"SayHelloStream"));
426        assert!(method_names.contains(&"SayHelloClientStream"));
427        assert!(method_names.contains(&"Chat"));
428    }
429}
mockforge_grpc/dynamic/proto_parser.rs

mockforge_grpc/dynamic/
proto_parser.rs
