diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index ae5fc0da..f59018af 100644 --- a/src-tauri/src/lib.rs +++ b/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( ..Default::default() }; - let (in_msg_tx, in_msg_rx) = tauri::async_runtime::channel::(16); + let (in_msg_tx, in_msg_rx) = tauri::async_runtime::channel::(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( }; 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( 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( 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( .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() diff --git a/src-tauri/yaak-grpc/src/any.rs b/src-tauri/yaak-grpc/src/any.rs new file mode 100644 index 00000000..ed41fbfd --- /dev/null +++ b/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) { + 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) { + 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"]); + } +} diff --git a/src-tauri/yaak-grpc/src/lib.rs b/src-tauri/yaak-grpc/src/lib.rs index 1de396a4..33fdcdda 100644 --- a/src-tauri/yaak-grpc/src/lib.rs +++ b/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; diff --git a/src-tauri/yaak-grpc/src/manager.rs b/src-tauri/yaak-grpc/src/manager.rs index e915966b..4aed8050 100644 --- a/src-tauri/yaak-grpc/src/manager.rs +++ b/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>, conn: Client, BoxBody>, pub uri: Uri, + use_reflection: bool, } #[derive(Default, Debug)] @@ -52,18 +57,19 @@ impl From for StreamError { } impl GrpcConnection { - pub fn service(&self, service: &str) -> Result { - 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 { - let service = self.service(service)?; + pub async fn method(&self, service: &str, method: &str) -> Result { + 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 { + 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, ) -> Result, 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, + stream: ReceiverStream, metadata: &BTreeMap, ) -> Result>, 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, + stream: ReceiverStream, metadata: &BTreeMap, ) -> Result, 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, ) -> Result>, 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, metadata: &BTreeMap, validate_certificates: bool, - ) -> Result<(), String> { - let pool = if proto_files.is_empty() { + ) -> Result { + 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, validate_certificates: bool, ) -> Result { - 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) -> Option<&DescriptorPool> { diff --git a/src-tauri/yaak-grpc/src/reflection.rs b/src-tauri/yaak-grpc/src/reflection.rs index ccd45166..84afff72 100644 --- a/src-tauri/yaak-grpc/src/reflection.rs +++ b/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>, + uri: &Uri, + json: &str, + metadata: &BTreeMap, +) -> 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>, pool: &mut DescriptorPool, client: &mut AutoReflectionClient,