Support Any type for gRPC reflection

src-tauri/src/lib.rs

@@ -29,8 +29,8 @@ use tokio::sync::Mutex;
 use tokio::task::block_in_place;
 use tokio::time;
 use yaak_common::window::WorkspaceWindowTrait;
-use yaak_grpc::manager::{DynamicMessage, GrpcHandle};
-use yaak_grpc::{Code, ServiceDefinition, deserialize_message, serialize_message};
+use yaak_grpc::manager::GrpcHandle;
+use yaak_grpc::{Code, ServiceDefinition, serialize_message};
 use yaak_models::models::{
     AnyModel, CookieJar, Environment, GrpcConnection, GrpcConnectionState, GrpcEvent,
     GrpcEventType, GrpcRequest, HttpRequest, HttpResponse, HttpResponseState, Plugin, Workspace,
@@ -252,7 +252,7 @@ async fn cmd_grpc_go<R: Runtime>(
         ..Default::default()
     };
 
-    let (in_msg_tx, in_msg_rx) = tauri::async_runtime::channel::<DynamicMessage>(16);
+    let (in_msg_tx, in_msg_rx) = tauri::async_runtime::channel::<String>(16);
     let maybe_in_msg_tx = std::sync::Mutex::new(Some(in_msg_tx.clone()));
     let (cancelled_tx, mut cancelled_rx) = tokio::sync::watch::channel(false);
 
@@ -298,7 +298,7 @@ async fn cmd_grpc_go<R: Runtime>(
     };
 
     let method_desc =
-        connection.method(&service, &method).map_err(|e| GenericError(e.to_string()))?;
+        connection.method(&service, &method).await.map_err(|e| GenericError(e.to_string()))?;
 
     #[derive(serde::Deserialize)]
     enum IncomingMsg {
@@ -313,7 +313,6 @@ async fn cmd_grpc_go<R: Runtime>(
         let environment_chain = environment_chain.clone();
         let window = window.clone();
         let base_msg = base_msg.clone();
-        let method_desc = method_desc.clone();
 
         move |ev: tauri::Event| {
             if *cancelled_rx.borrow() {
@@ -335,7 +334,6 @@ async fn cmd_grpc_go<R: Runtime>(
                     let window = window.clone();
                     let app_handle = app_handle.clone();
                     let base_msg = base_msg.clone();
-                    let method_desc = method_desc.clone();
                     let environment_chain = environment_chain.clone();
                     let msg = block_in_place(|| {
                         tauri::async_runtime::block_on(async {
@@ -355,27 +353,7 @@ async fn cmd_grpc_go<R: Runtime>(
                             .expect("Failed to render template")
                         })
                     });
-                    let d_msg: DynamicMessage = match deserialize_message(msg.as_str(), method_desc)
-                    {
-                        Ok(d_msg) => d_msg,
-                        Err(e) => {
-                            tauri::async_runtime::spawn(async move {
-                                app_handle
-                                    .db()
-                                    .upsert_grpc_event(
-                                        &GrpcEvent {
-                                            event_type: GrpcEventType::Error,
-                                            content: e.to_string(),
-                                            ..base_msg.clone()
-                                        },
-                                        &UpdateSource::from_window(&window),
-                                    )
-                                    .unwrap();
-                            });
-                            return;
-                        }
-                    };
-                    in_msg_tx.try_send(d_msg).unwrap();
+                    in_msg_tx.try_send(msg.clone()).unwrap();
                     tauri::async_runtime::spawn(async move {
                         app_handle
                             .db()

src-tauri/yaak-grpc/src/any.rs

@@ -0,0 +1,60 @@
+use log::error;
+
+pub(crate) fn collect_any_types(json: &str, out: &mut Vec<String>) {
+    let value = match serde_json::from_str(json).map_err(|e| e.to_string()) {
+        Ok(v) => v,
+        Err(e) => {
+            error!("Failed to parse gRPC message JSON: {e:?}");
+            return;
+        }
+    };
+    collect_any_types_value(&value, out);
+}
+
+fn collect_any_types_value(json: &serde_json::Value, out: &mut Vec<String>) {
+    match json {
+        serde_json::Value::Object(map) => {
+            if let Some(t) = map.get("@type").and_then(|v| v.as_str()) {
+                if let Some(full_name) = t.rsplit_once('/').map(|(_, n)| n) {
+                    out.push(full_name.to_string());
+                }
+            }
+
+            for v in map.values() {
+                collect_any_types_value(v, out);
+            }
+        }
+        serde_json::Value::Array(arr) => {
+            for v in arr {
+                collect_any_types_value(v, out);
+            }
+        }
+        _ => {}
+    }
+}
+
+// Write tests for this
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn test_collect_any_types() {
+        let json = r#"{
+            "mounts": [
+              {
+                "mountSource": {
+                  "@type": "type.googleapis.com/mount_source.MountSourceRBDVolume",
+                  "volumeID": "volumes/rbd"
+                }
+              }
+            ],
+            "foo": {
+              "@type": "type.googleapis.com/foo.bar",
+              "foo": "fooo"
+            }
+        }"#;
+
+        let mut out = Vec::new();
+        super::collect_any_types(json, &mut out);
+        assert_eq!(out, vec!["foo.bar", "mount_source.MountSourceRBDVolume"]);
+    }
+}

src-tauri/yaak-grpc/src/lib.rs

@@ -8,6 +8,7 @@ mod json_schema;
 pub mod manager;
 mod reflection;
 mod transport;
+mod any;
 
 pub use tonic::metadata::*;
 pub use tonic::Code;

src-tauri/yaak-grpc/src/manager.rs

@@ -1,30 +1,35 @@
-use std::collections::BTreeMap;
-use std::path::PathBuf;
-use std::str::FromStr;
-
+use crate::codec::DynamicCodec;
+use crate::reflection::{
+    fill_pool_from_files, fill_pool_from_reflection, method_desc_to_path, reflect_types_for_message,
+};
+use crate::transport::get_transport;
+use crate::{MethodDefinition, ServiceDefinition, json_schema};
 use hyper_rustls::HttpsConnector;
-use hyper_util::client::legacy::connect::HttpConnector;
 use hyper_util::client::legacy::Client;
+use hyper_util::client::legacy::connect::HttpConnector;
+use log::warn;
 pub use prost_reflect::DynamicMessage;
 use prost_reflect::{DescriptorPool, MethodDescriptor, ServiceDescriptor};
 use serde_json::Deserializer;
+use std::collections::BTreeMap;
+use std::path::PathBuf;
+use std::str::FromStr;
+use std::sync::Arc;
 use tauri::AppHandle;
+use tokio::sync::RwLock;
+use tokio_stream::StreamExt;
 use tokio_stream::wrappers::ReceiverStream;
 use tonic::body::BoxBody;
 use tonic::metadata::{MetadataKey, MetadataValue};
 use tonic::transport::Uri;
 use tonic::{IntoRequest, IntoStreamingRequest, Request, Response, Status, Streaming};
 
-use crate::codec::DynamicCodec;
-use crate::reflection::{fill_pool_from_files, fill_pool_from_reflection, method_desc_to_path};
-use crate::transport::get_transport;
-use crate::{json_schema, MethodDefinition, ServiceDefinition};
-
 #[derive(Clone)]
 pub struct GrpcConnection {
-    pool: DescriptorPool,
+    pool: Arc<RwLock<DescriptorPool>>,
     conn: Client<HttpsConnector<HttpConnector>, BoxBody>,
     pub uri: Uri,
+    use_reflection: bool,
 }
 
 #[derive(Default, Debug)]
@@ -52,18 +57,19 @@ impl From<Status> for StreamError {
 }
 
 impl GrpcConnection {
-    pub fn service(&self, service: &str) -> Result<ServiceDescriptor, String> {
-        let service = self.pool.get_service_by_name(service).ok_or("Failed to find service")?;
-        Ok(service)
-    }
-
-    pub fn method(&self, service: &str, method: &str) -> Result<MethodDescriptor, String> {
-        let service = self.service(service)?;
+    pub async fn method(&self, service: &str, method: &str) -> Result<MethodDescriptor, String> {
+        let service = self.service(service).await?;
         let method =
             service.methods().find(|m| m.name() == method).ok_or("Failed to find method")?;
         Ok(method)
     }
 
+    async fn service(&self, service: &str) -> Result<ServiceDescriptor, String> {
+        let pool = self.pool.read().await;
+        let service = pool.get_service_by_name(service).ok_or("Failed to find service")?;
+        Ok(service)
+    }
+
     pub async fn unary(
         &self,
         service: &str,
@@ -71,7 +77,10 @@ impl GrpcConnection {
         message: &str,
         metadata: &BTreeMap<String, String>,
     ) -> Result<Response<DynamicMessage>, StreamError> {
-        let method = &self.method(&service, &method)?;
+        if self.use_reflection {
+            reflect_types_for_message(self.pool.clone(), &self.uri, message, metadata).await?;
+        }
+        let method = &self.method(&service, &method).await?;
         let input_message = method.input();
 
         let mut deserializer = Deserializer::from_str(message);
@@ -95,18 +104,47 @@ impl GrpcConnection {
         &self,
         service: &str,
         method: &str,
-        stream: ReceiverStream<DynamicMessage>,
+        stream: ReceiverStream<String>,
         metadata: &BTreeMap<String, String>,
     ) -> Result<Response<Streaming<DynamicMessage>>, StreamError> {
-        let method = &self.method(&service, &method)?;
+        let method = &self.method(&service, &method).await?;
+        let mapped_stream = {
+            let input_message = method.input();
+            let pool = self.pool.clone();
+            let uri = self.uri.clone();
+            let md = metadata.clone();
+            let use_reflection = self.use_reflection.clone();
+            stream.filter_map(move |json| {
+                let pool = pool.clone();
+                let uri = uri.clone();
+                let input_message = input_message.clone();
+                let md = md.clone();
+                let use_reflection = use_reflection.clone();
+                tauri::async_runtime::block_on(async move {
+                    if use_reflection {
+                        if let Err(e) = reflect_types_for_message(pool, &uri, &json, &md).await {
+                            warn!("Failed to resolve Any types: {e}");
+                        }
+                    }
+                    let mut de = Deserializer::from_str(&json);
+                    match DynamicMessage::deserialize(input_message, &mut de) {
+                        Ok(m) => Some(m),
+                        Err(e) => {
+                            warn!("Failed to deserialize message: {e}");
+                            None
+                        }
+                    }
+                })
+            })
+        };
+
         let mut client = tonic::client::Grpc::with_origin(self.conn.clone(), self.uri.clone());
-
-        let mut req = stream.into_streaming_request();
-
-        decorate_req(metadata, &mut req).map_err(|e| e.to_string())?;
-
         let path = method_desc_to_path(method);
         let codec = DynamicCodec::new(method.clone());
+
+        let mut req = mapped_stream.into_streaming_request();
+        decorate_req(metadata, &mut req).map_err(|e| e.to_string())?;
+
         client.ready().await.map_err(|e| e.to_string())?;
         Ok(client.streaming(req, path, codec).await?)
     }
@@ -115,16 +153,47 @@ impl GrpcConnection {
         &self,
         service: &str,
         method: &str,
-        stream: ReceiverStream<DynamicMessage>,
+        stream: ReceiverStream<String>,
         metadata: &BTreeMap<String, String>,
     ) -> Result<Response<DynamicMessage>, StreamError> {
-        let method = &self.method(&service, &method)?;
-        let mut client = tonic::client::Grpc::with_origin(self.conn.clone(), self.uri.clone());
-        let mut req = stream.into_streaming_request();
-        decorate_req(metadata, &mut req).map_err(|e| e.to_string())?;
+        let method = &self.method(&service, &method).await?;
+        let mapped_stream = {
+            let input_message = method.input();
+            let pool = self.pool.clone();
+            let uri = self.uri.clone();
+            let md = metadata.clone();
+            let use_reflection = self.use_reflection.clone();
+            stream.filter_map(move |json| {
+                let pool = pool.clone();
+                let uri = uri.clone();
+                let input_message = input_message.clone();
+                let md = md.clone();
+                let use_reflection = use_reflection.clone();
+                tauri::async_runtime::block_on(async move {
+                    if use_reflection {
+                        if let Err(e) = reflect_types_for_message(pool, &uri, &json, &md).await {
+                            warn!("Failed to resolve Any types: {e}");
+                        }
+                    }
+                    let mut de = Deserializer::from_str(&json);
+                    match DynamicMessage::deserialize(input_message, &mut de) {
+                        Ok(m) => Some(m),
+                        Err(e) => {
+                            warn!("Failed to deserialize message: {e}");
+                            None
+                        }
+                    }
+                })
+            })
+        };
 
+        let mut client = tonic::client::Grpc::with_origin(self.conn.clone(), self.uri.clone());
         let path = method_desc_to_path(method);
         let codec = DynamicCodec::new(method.clone());
+
+        let mut req = mapped_stream.into_streaming_request();
+        decorate_req(metadata, &mut req).map_err(|e| e.to_string())?;
+
         client.ready().await.unwrap();
         client.client_streaming(req, path, codec).await.map_err(|e| StreamError {
             message: e.message().to_string(),
@@ -139,7 +208,7 @@ impl GrpcConnection {
         message: &str,
         metadata: &BTreeMap<String, String>,
     ) -> Result<Response<Streaming<DynamicMessage>>, StreamError> {
-        let method = &self.method(&service, &method)?;
+        let method = &self.method(&service, &method).await?;
         let input_message = method.input();
 
         let mut deserializer = Deserializer::from_str(message);
@@ -182,8 +251,9 @@ impl GrpcHandle {
         proto_files: &Vec<PathBuf>,
         metadata: &BTreeMap<String, String>,
         validate_certificates: bool,
-    ) -> Result<(), String> {
-        let pool = if proto_files.is_empty() {
+    ) -> Result<bool, String> {
+        let server_reflection = proto_files.is_empty();
+        let pool = if server_reflection {
             let full_uri = uri_from_str(uri)?;
             fill_pool_from_reflection(&full_uri, metadata, validate_certificates).await
         } else {
@@ -191,7 +261,7 @@ impl GrpcHandle {
         }?;
 
         self.pools.insert(make_pool_key(id, uri, proto_files), pool.clone());
-        Ok(())
+        Ok(server_reflection)
     }
 
     pub async fn services(
@@ -242,17 +312,17 @@ impl GrpcHandle {
         metadata: &BTreeMap<String, String>,
         validate_certificates: bool,
     ) -> Result<GrpcConnection, String> {
-        self.reflect(id, uri, proto_files, metadata, validate_certificates).await?;
-        let pool = self.get_pool(id, uri, proto_files).ok_or("Failed to get pool")?;
-
+        let use_reflection =
+            self.reflect(id, uri, proto_files, metadata, validate_certificates).await?;
+        let pool = self.get_pool(id, uri, proto_files).ok_or("Failed to get pool")?.clone();
         let uri = uri_from_str(uri)?;
         let conn = get_transport(validate_certificates);
-        let connection = GrpcConnection {
-            pool: pool.clone(),
+        Ok(GrpcConnection {
+            pool: Arc::new(RwLock::new(pool)),
+            use_reflection,
             conn,
             uri,
-        };
-        Ok(connection)
+        })
     }
 
     fn get_pool(&self, id: &str, uri: &str, proto_files: &Vec<PathBuf>) -> Option<&DescriptorPool> {

src-tauri/yaak-grpc/src/reflection.rs

@@ -1,9 +1,4 @@
-use std::collections::{BTreeMap, HashSet};
-use std::env::temp_dir;
-use std::ops::Deref;
-use std::path::{Path, PathBuf};
-use std::str::FromStr;
-
+use crate::any::collect_any_types;
 use crate::client::AutoReflectionClient;
 use anyhow::anyhow;
 use async_recursion::async_recursion;
@@ -11,10 +6,17 @@ use log::{debug, info, warn};
 use prost::Message;
 use prost_reflect::{DescriptorPool, MethodDescriptor};
 use prost_types::{FileDescriptorProto, FileDescriptorSet};
+use std::collections::{BTreeMap, HashSet};
+use std::env::temp_dir;
+use std::ops::Deref;
+use std::path::{Path, PathBuf};
+use std::str::FromStr;
+use std::sync::Arc;
 use tauri::path::BaseDirectory;
 use tauri::{AppHandle, Manager};
 use tauri_plugin_shell::ShellExt;
 use tokio::fs;
+use tokio::sync::RwLock;
 use tonic::codegen::http::uri::PathAndQuery;
 use tonic::transport::Uri;
 use tonic_reflection::pb::v1::server_reflection_request::MessageRequest;
@@ -130,9 +132,9 @@ pub async fn fill_pool_from_reflection(
             continue;
         }
         if service == "grpc.reflection.v1.ServerReflection" {
-            // TODO: update reflection client to use v1
             continue;
         }
+        debug!("Fetching descriptors for {}", service);
         file_descriptor_set_from_service_name(&service, &mut pool, &mut client, metadata).await;
     }
 
@@ -188,8 +190,54 @@ async fn file_descriptor_set_from_service_name(
     .await;
 }
 
+pub(crate) async fn reflect_types_for_message(
+    pool: Arc<RwLock<DescriptorPool>>,
+    uri: &Uri,
+    json: &str,
+    metadata: &BTreeMap<String, String>,
+) -> Result<(), String> {
+    // 1. Collect all Any types in the JSON
+    let mut extra_types = Vec::new();
+    collect_any_types(json, &mut extra_types);
+
+    if extra_types.is_empty() {
+        return Ok(()); // nothing to do
+    }
+
+    let mut client = AutoReflectionClient::new(uri, false);
+    for extra_type in extra_types {
+        {
+            let guard = pool.read().await;
+            if guard.get_message_by_name(&extra_type).is_some() {
+                continue;
+            }
+        }
+        info!("Adding file descriptor for {:?} from reflection", extra_type);
+        let req = MessageRequest::FileContainingSymbol(extra_type.clone().into());
+        let resp = match client.send_reflection_request(req, metadata).await {
+            Ok(r) => r,
+            Err(e) => {
+                return Err(format!(
+                    "Error sending reflection request for @type \"{extra_type}\": {e}",
+                ));
+            }
+        };
+        let files = match resp {
+            MessageResponse::FileDescriptorResponse(resp) => resp.file_descriptor_proto,
+            _ => panic!("Expected a FileDescriptorResponse variant"),
+        };
+
+        {
+            let mut guard = pool.write().await;
+            add_file_descriptors_to_pool(files, &mut *guard, &mut client, metadata).await;
+        }
+    }
+
+    Ok(())
+}
+
 #[async_recursion]
-async fn add_file_descriptors_to_pool(
+pub(crate) async fn add_file_descriptors_to_pool(
     fds: Vec<Vec<u8>>,
     pool: &mut DescriptorPool,
     client: &mut AutoReflectionClient,