use std::collections::BTreeMap; use std::env::temp_dir; use std::ops::Deref; use std::path::PathBuf; use std::str::FromStr; use crate::client::AutoReflectionClient; use anyhow::anyhow; use async_recursion::async_recursion; use log::{debug, warn}; use prost::Message; use prost_reflect::{DescriptorPool, MethodDescriptor}; use prost_types::{FileDescriptorProto, FileDescriptorSet}; use tauri::path::BaseDirectory; use tauri::{AppHandle, Manager}; use tauri_plugin_shell::ShellExt; use tokio::fs; use tonic::codegen::http::uri::PathAndQuery; use tonic::transport::Uri; use tonic_reflection::pb::v1::server_reflection_request::MessageRequest; use tonic_reflection::pb::v1::server_reflection_response::MessageResponse; pub async fn fill_pool_from_files( app_handle: &AppHandle, paths: &Vec, ) -> Result { let mut pool = DescriptorPool::new(); let random_file_name = format!("{}.desc", uuid::Uuid::new_v4()); let desc_path = temp_dir().join(random_file_name); let global_import_dir = app_handle .path() .resolve("vendored/protoc/include", BaseDirectory::Resource) .expect("failed to resolve protoc include directory"); // HACK: Remove UNC prefix for Windows paths let global_import_dir = dunce::simplified(global_import_dir.as_path()).to_string_lossy().to_string(); let desc_path = dunce::simplified(desc_path.as_path()); let mut args = vec![ "--include_imports".to_string(), "--include_source_info".to_string(), "-I".to_string(), global_import_dir, "-o".to_string(), desc_path.to_string_lossy().to_string(), ]; for p in paths { if p.as_path().exists() { args.push(p.to_string_lossy().to_string()); } else { continue; } let parent = p.as_path().parent(); if let Some(parent_path) = parent { args.push("-I".to_string()); args.push(parent_path.to_string_lossy().to_string()); args.push("-I".to_string()); args.push(parent_path.parent().unwrap().to_string_lossy().to_string()); } else { debug!("ignoring {:?} since it does not exist.", parent) } } let out = app_handle .shell() .sidecar("yaakprotoc") .expect("yaakprotoc not found") .args(args) .output() .await .expect("yaakprotoc failed to run"); if !out.status.success() { return Err(format!( "protoc failed with status {}: {}", out.status.code().unwrap(), String::from_utf8_lossy(out.stderr.as_slice()) )); } let bytes = fs::read(desc_path).await.map_err(|e| e.to_string())?; let fdp = FileDescriptorSet::decode(bytes.deref()).map_err(|e| e.to_string())?; pool.add_file_descriptor_set(fdp).map_err(|e| e.to_string())?; fs::remove_file(desc_path).await.map_err(|e| e.to_string())?; Ok(pool) } pub async fn fill_pool_from_reflection( uri: &Uri, metadata: &BTreeMap, validate_certificates: bool, ) -> Result { let mut pool = DescriptorPool::new(); let mut client = AutoReflectionClient::new(uri, validate_certificates); for service in list_services(&mut client, metadata).await? { if service == "grpc.reflection.v1alpha.ServerReflection" { continue; } if service == "grpc.reflection.v1.ServerReflection" { // TODO: update reflection client to use v1 continue; } file_descriptor_set_from_service_name(&service, &mut pool, &mut client, metadata).await; } Ok(pool) } async fn list_services( client: &mut AutoReflectionClient, metadata: &BTreeMap, ) -> Result, String> { let response = client.send_reflection_request(MessageRequest::ListServices("".into()), metadata).await?; let list_services_response = match response { MessageResponse::ListServicesResponse(resp) => resp, _ => panic!("Expected a ListServicesResponse variant"), }; Ok(list_services_response.service.iter().map(|s| s.name.clone()).collect::>()) } async fn file_descriptor_set_from_service_name( service_name: &str, pool: &mut DescriptorPool, client: &mut AutoReflectionClient, metadata: &BTreeMap, ) { let response = match client .send_reflection_request( MessageRequest::FileContainingSymbol(service_name.into()), metadata, ) .await { Ok(resp) => resp, Err(e) => { warn!("Error fetching file descriptor for service {}: {}", service_name, e); return; } }; let file_descriptor_response = match response { MessageResponse::FileDescriptorResponse(resp) => resp, _ => panic!("Expected a FileDescriptorResponse variant"), }; add_file_descriptors_to_pool( file_descriptor_response.file_descriptor_proto, pool, client, metadata, ) .await; } #[async_recursion] async fn add_file_descriptors_to_pool( fds: Vec>, pool: &mut DescriptorPool, client: &mut AutoReflectionClient, metadata: &BTreeMap, ) { let mut topo_sort = topology::SimpleTopoSort::new(); let mut fd_mapping = std::collections::HashMap::with_capacity(fds.len()); for fd in fds { let fdp = FileDescriptorProto::decode(fd.deref()).unwrap(); topo_sort.insert(fdp.name().to_string(), fdp.dependency.clone()); fd_mapping.insert(fdp.name().to_string(), fdp); } for node in topo_sort { match node { Ok(node) => { if let Some(fdp) = fd_mapping.remove(&node) { pool.add_file_descriptor_proto(fdp).expect("add file descriptor proto"); } else { file_descriptor_set_by_filename(node.as_str(), pool, client, metadata).await; } } Err(_) => panic!("proto file got cycle!"), } } } async fn file_descriptor_set_by_filename( filename: &str, pool: &mut DescriptorPool, client: &mut AutoReflectionClient, metadata: &BTreeMap, ) { // We already fetched this file if let Some(_) = pool.get_file_by_name(filename) { return; } let msg = MessageRequest::FileByFilename(filename.into()); let response = client.send_reflection_request(msg, metadata).await; let file_descriptor_response = match response { Ok(MessageResponse::FileDescriptorResponse(resp)) => resp, Ok(_) => { panic!("Expected a FileDescriptorResponse variant") } Err(e) => { warn!("Error fetching file descriptor for {}: {}", filename, e); return; } }; add_file_descriptors_to_pool( file_descriptor_response.file_descriptor_proto, pool, client, metadata, ) .await; } pub fn method_desc_to_path(md: &MethodDescriptor) -> PathAndQuery { let full_name = md.full_name(); let (namespace, method_name) = full_name .rsplit_once('.') .ok_or_else(|| anyhow!("invalid method path")) .expect("invalid method path"); PathAndQuery::from_str(&format!("/{}/{}", namespace, method_name)).expect("invalid method path") } mod topology { use std::collections::{HashMap, HashSet}; pub struct SimpleTopoSort { out_graph: HashMap>, in_graph: HashMap>, } impl SimpleTopoSort where T: Eq + std::hash::Hash + Clone, { pub fn new() -> Self { SimpleTopoSort { out_graph: HashMap::new(), in_graph: HashMap::new(), } } pub fn insert>(&mut self, node: T, deps: I) { self.out_graph.entry(node.clone()).or_insert(HashSet::new()); for dep in deps { self.out_graph.entry(node.clone()).or_insert(HashSet::new()).insert(dep.clone()); self.in_graph.entry(dep.clone()).or_insert(HashSet::new()).insert(node.clone()); } } } impl IntoIterator for SimpleTopoSort where T: Eq + std::hash::Hash + Clone, { type Item = as Iterator>::Item; type IntoIter = SimpleTopoSortIter; fn into_iter(self) -> Self::IntoIter { SimpleTopoSortIter::new(self) } } pub struct SimpleTopoSortIter { data: SimpleTopoSort, zero_indegree: Vec, } impl SimpleTopoSortIter where T: Eq + std::hash::Hash + Clone, { pub fn new(data: SimpleTopoSort) -> Self { let mut zero_indegree = Vec::new(); for (node, _) in data.in_graph.iter() { if !data.out_graph.contains_key(node) { zero_indegree.push(node.clone()); } } for (node, deps) in data.out_graph.iter() { if deps.is_empty() { zero_indegree.push(node.clone()); } } SimpleTopoSortIter { data, zero_indegree, } } } impl Iterator for SimpleTopoSortIter where T: Eq + std::hash::Hash + Clone, { type Item = Result; fn next(&mut self) -> Option { if self.zero_indegree.is_empty() { if self.data.out_graph.is_empty() { return None; } return Some(Err("Cycle detected")); } let node = self.zero_indegree.pop().unwrap(); if let Some(parents) = self.data.in_graph.get(&node) { for parent in parents.iter() { let deps = self.data.out_graph.get_mut(parent).unwrap(); deps.remove(&node); if deps.is_empty() { self.zero_indegree.push(parent.clone()); } } } self.data.out_graph.remove(&node); Some(Ok(node)) } } #[test] fn test_sort() { { let mut topo_sort = SimpleTopoSort::new(); topo_sort.insert("a", []); for node in topo_sort { match node { Ok(n) => assert_eq!(n, "a"), Err(e) => panic!("err {}", e), } } } { let mut topo_sort = SimpleTopoSort::new(); topo_sort.insert("a", ["b"]); topo_sort.insert("b", []); let mut iter = topo_sort.into_iter(); match iter.next() { Some(Ok(n)) => assert_eq!(n, "b"), _ => panic!("err"), } match iter.next() { Some(Ok(n)) => assert_eq!(n, "a"), _ => panic!("err"), } assert_eq!(iter.next(), None); } } }