gRPC schema from files!

src-tauri/grpc/Cargo.toml

@@ -16,5 +16,7 @@ prost-reflect = { version = "0.12.0", features = ["serde", "derive"] }
 log = "0.4.20"
 once_cell = { version = "1.19.0", features = [] }
 anyhow = "1.0.79"
-hyper = { version = "0.14"}
+hyper = { version = "0.14" }
 hyper-rustls = { version = "0.24.0", features = ["http2"] }
+protoc-bin-vendored = "3.0.0"
+uuid = { version = "1.7.0", features = ["v4"] }

src-tauri/grpc/src/manager.rs

@@ -1,4 +1,5 @@
 use std::collections::HashMap;
+use std::path::PathBuf;
 
 use hyper::client::HttpConnector;
 use hyper::Client;
@@ -13,11 +14,9 @@ use tonic::transport::Uri;
 use tonic::{IntoRequest, IntoStreamingRequest, Streaming};
 
 use crate::codec::DynamicCodec;
-use crate::proto::{fill_pool, method_desc_to_path};
+use crate::proto::{fill_pool, fill_pool_from_files, get_transport, method_desc_to_path};
 use crate::{json_schema, MethodDefinition, ServiceDefinition};
 
-type Result<T> = std::result::Result<T, String>;
-
 #[derive(Clone)]
 pub struct GrpcConnection {
     pool: DescriptorPool,
@@ -26,7 +25,12 @@ pub struct GrpcConnection {
 }
 
 impl GrpcConnection {
-    pub async fn unary(&self, service: &str, method: &str, message: &str) -> Result<String> {
+    pub async fn unary(
+        &self,
+        service: &str,
+        method: &str,
+        message: &str,
+    ) -> Result<String, String> {
         let service = self.pool.get_service_by_name(service).unwrap();
         let method = &service.methods().find(|m| m.name() == method).unwrap();
         let input_message = method.input();
@@ -55,7 +59,7 @@ impl GrpcConnection {
         service: &str,
         method: &str,
         stream: ReceiverStream<String>,
-    ) -> Result<Streaming<DynamicMessage>> {
+    ) -> Result<Streaming<DynamicMessage>, String> {
         let service = self.pool.get_service_by_name(service).unwrap();
         let method = &service.methods().find(|m| m.name() == method).unwrap();
 
@@ -87,7 +91,7 @@ impl GrpcConnection {
         service: &str,
         method: &str,
         stream: ReceiverStream<String>,
-    ) -> Result<DynamicMessage> {
+    ) -> Result<DynamicMessage, String> {
         let service = self.pool.get_service_by_name(service).unwrap();
         let method = &service.methods().find(|m| m.name() == method).unwrap();
         let mut client = tonic::client::Grpc::with_origin(self.conn.clone(), self.uri.clone());
@@ -121,7 +125,7 @@ impl GrpcConnection {
         service: &str,
         method: &str,
         message: &str,
-    ) -> Result<Streaming<DynamicMessage>> {
+    ) -> Result<Streaming<DynamicMessage>, String> {
         let service = self.pool.get_service_by_name(service).unwrap();
         let method = &service.methods().find(|m| m.name() == method).unwrap();
         let input_message = method.input();
@@ -159,56 +163,56 @@ impl Default for GrpcHandle {
 }
 
 impl GrpcHandle {
-    pub async fn clean_reflect(&mut self, uri: &Uri) -> Result<Vec<ServiceDefinition>> {
-        self.pools.remove(&uri.to_string());
-        self.reflect(uri).await
+    pub async fn services_from_files(
+        &self,
+        paths: Vec<PathBuf>,
+    ) -> Result<Vec<ServiceDefinition>, String> {
+        let pool = fill_pool_from_files(paths).await?;
+        Ok(self.services_from_pool(&pool))
+    }
+    pub async fn services_from_reflection(
+        &mut self,
+        uri: &Uri,
+    ) -> Result<Vec<ServiceDefinition>, String> {
+        let pool = fill_pool(uri).await?;
+        Ok(self.services_from_pool(&pool))
     }
 
-    pub async fn reflect(&mut self, uri: &Uri) -> Result<Vec<ServiceDefinition>> {
-        let pool = match self.pools.get(&uri.to_string()) {
-            Some(p) => p.clone(),
-            None => {
-                let (pool, _) = fill_pool(uri).await?;
-                self.pools.insert(uri.to_string(), pool.clone());
-                pool
-            }
-        };
-
-        let result =
-            pool.services()
-                .map(|s| {
-                    let mut def = ServiceDefinition {
-                        name: s.full_name().to_string(),
-                        methods: vec![],
-                    };
-                    for method in s.methods() {
-                        let input_message = method.input();
-                        def.methods.push(MethodDefinition {
-                            name: method.name().to_string(),
-                            server_streaming: method.is_server_streaming(),
-                            client_streaming: method.is_client_streaming(),
-                            schema: serde_json::to_string_pretty(
-                                &json_schema::message_to_json_schema(&pool, input_message),
-                            )
-                            .unwrap(),
-                        })
-                    }
-                    def
-                })
-                .collect::<Vec<_>>();
-
-        Ok(result)
+    fn services_from_pool(&self, pool: &DescriptorPool) -> Vec<ServiceDefinition> {
+        pool.services()
+            .map(|s| {
+                let mut def = ServiceDefinition {
+                    name: s.full_name().to_string(),
+                    methods: vec![],
+                };
+                for method in s.methods() {
+                    let input_message = method.input();
+                    def.methods.push(MethodDefinition {
+                        name: method.name().to_string(),
+                        server_streaming: method.is_server_streaming(),
+                        client_streaming: method.is_client_streaming(),
+                        schema: serde_json::to_string_pretty(&json_schema::message_to_json_schema(
+                            &pool,
+                            input_message,
+                        ))
+                        .unwrap(),
+                    })
+                }
+                def
+            })
+            .collect::<Vec<_>>()
     }
 
     pub async fn server_streaming(
         &mut self,
         id: &str,
         uri: Uri,
+        proto_files: Vec<PathBuf>,
         service: &str,
         method: &str,
         message: &str,
-    ) -> Result<Streaming<DynamicMessage>> {
-        self.connect(id, uri)
+    ) -> Result<Streaming<DynamicMessage>, String> {
+        self.connect(id, uri, proto_files)
             .await?
             .server_streaming(service, method, message)
             .await
@@ -218,11 +222,12 @@ impl GrpcHandle {
         &mut self,
         id: &str,
         uri: Uri,
+        proto_files: Vec<PathBuf>,
         service: &str,
         method: &str,
         stream: ReceiverStream<String>,
-    ) -> Result<DynamicMessage> {
-        self.connect(id, uri)
+    ) -> Result<DynamicMessage, String> {
+        self.connect(id, uri, proto_files)
             .await?
             .client_streaming(service, method, stream)
             .await
@@ -232,18 +237,36 @@ impl GrpcHandle {
         &mut self,
         id: &str,
         uri: Uri,
+        proto_files: Vec<PathBuf>,
         service: &str,
         method: &str,
         stream: ReceiverStream<String>,
-    ) -> Result<Streaming<DynamicMessage>> {
-        self.connect(id, uri)
+    ) -> Result<Streaming<DynamicMessage>, String> {
+        self.connect(id, uri, proto_files)
             .await?
             .streaming(service, method, stream)
             .await
     }
 
-    pub async fn connect(&mut self, id: &str, uri: Uri) -> Result<GrpcConnection> {
-        let (pool, conn) = fill_pool(&uri).await?;
+    pub async fn connect(
+        &mut self,
+        id: &str,
+        uri: Uri,
+        proto_files: Vec<PathBuf>,
+    ) -> Result<GrpcConnection, String> {
+        let pool = match self.pools.get(id) {
+            Some(p) => p.clone(),
+            None => match proto_files.len() {
+                0 => fill_pool(&uri).await?,
+                _ => {
+                    let pool = fill_pool_from_files(proto_files).await?;
+                    self.pools.insert(id.to_string(), pool.clone());
+                    pool
+                }
+            },
+        };
+
+        let conn = get_transport();
         let connection = GrpcConnection { pool, conn, uri };
         self.connections.insert(id.to_string(), connection.clone());
         Ok(connection)

src-tauri/grpc/src/proto.rs

@@ -1,15 +1,17 @@
+use std::env::temp_dir;
 use std::ops::Deref;
 use std::path::PathBuf;
+use std::process::Command;
 use std::str::FromStr;
 
 use anyhow::anyhow;
 use hyper::client::HttpConnector;
 use hyper::Client;
 use hyper_rustls::{HttpsConnector, HttpsConnectorBuilder};
-use log::warn;
+use log::{debug, warn};
 use prost::Message;
 use prost_reflect::{DescriptorPool, MethodDescriptor};
-use prost_types::FileDescriptorProto;
+use prost_types::{FileDescriptorProto, FileDescriptorSet};
 use tokio::fs;
 use tokio_stream::StreamExt;
 use tonic::body::BoxBody;
@@ -23,36 +25,71 @@ use tonic_reflection::pb::ServerReflectionRequest;
 
 pub async fn fill_pool_from_files(paths: Vec<PathBuf>) -> Result<DescriptorPool, String> {
     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 bin = protoc_bin_vendored::protoc_bin_path().unwrap();
+
+    let mut cmd = Command::new(bin.clone());
+    cmd.arg("--include_imports")
+        .arg("--include_source_info")
+        .arg("-o")
+        .arg(&desc_path);
+
     for p in paths {
-        let bytes = fs::read(p).await.unwrap();
-        let fdp = FileDescriptorProto::decode(bytes.deref()).unwrap();
-        pool.add_file_descriptor_proto(fdp)
-            .map_err(|e| e.to_string())?;
+        if p.as_path().exists() {
+            cmd.arg(p.as_path().to_string_lossy().as_ref());
+        } else {
+            continue;
+        }
+
+        let parent = p.as_path().parent();
+        if let Some(parent_path) = parent {
+            cmd.arg("-I").arg(parent_path);
+        } else {
+            debug!("ignoring {:?} since it does not exist.", parent)
+        }
     }
+
+    debug!("Running: {:?}", cmd);
+
+    let output = cmd.output().map_err(|e| e.to_string())?;
+    if !output.status.success() {
+        return Err(format!(
+            "protoc failed: {}",
+            String::from_utf8_lossy(&output.stderr)
+        ));
+    }
+
+    let bytes = fs::read(desc_path.as_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.as_path())
+        .await
+        .map_err(|e| e.to_string())?;
+
     Ok(pool)
 }
-pub async fn fill_pool(
-    uri: &Uri,
-) -> Result<
-    (
-        DescriptorPool,
-        Client<HttpsConnector<HttpConnector>, BoxBody>,
-    ),
-    String,
-> {
-    let mut pool = DescriptorPool::new();
+
+pub fn get_transport() -> Client<HttpsConnector<HttpConnector>, BoxBody> {
     let connector = HttpsConnectorBuilder::new().with_native_roots();
     let connector = connector.https_or_http().enable_http2().wrap_connector({
         let mut http_connector = HttpConnector::new();
         http_connector.enforce_http(false);
         http_connector
     });
-    let transport = Client::builder()
+    Client::builder()
         .pool_max_idle_per_host(0)
         .http2_only(true)
-        .build(connector);
+        .build(connector)
+}
 
-    let mut client = ServerReflectionClient::with_origin(transport.clone(), uri.clone());
+pub async fn fill_pool(uri: &Uri) -> Result<DescriptorPool, String> {
+    let mut pool = DescriptorPool::new();
+    let mut client = ServerReflectionClient::with_origin(get_transport(), uri.clone());
 
     for service in list_services(&mut client).await? {
         if service == "grpc.reflection.v1alpha.ServerReflection" {
@@ -61,7 +98,7 @@ pub async fn fill_pool(
         file_descriptor_set_from_service_name(&service, &mut pool, &mut client).await;
     }
 
-    Ok((pool, transport))
+    Ok(pool)
 }
 
 async fn list_services(