From 035fe54df0c01eb7ec2a04a14dc5e7cc18c45e13 Mon Sep 17 00:00:00 2001 From: Gregory Schier Date: Sun, 11 May 2025 07:20:57 -0700 Subject: [PATCH] Send grpc metadata/auth with reflection requests Closes https://feedback.yaak.app/p/send-metadata-during-grpc-reflection --- src-tauri/src/grpc.rs | 53 +++++++++++++++++++++++++- src-tauri/src/lib.rs | 55 ++++++--------------------- src-tauri/yaak-grpc/src/client.rs | 17 ++++++--- src-tauri/yaak-grpc/src/manager.rs | 22 +++++++---- src-tauri/yaak-grpc/src/reflection.rs | 48 +++++++++++++++++------ 5 files changed, 124 insertions(+), 71 deletions(-) diff --git a/src-tauri/src/grpc.rs b/src-tauri/src/grpc.rs index f3ec14f6..94ab06f4 100644 --- a/src-tauri/src/grpc.rs +++ b/src-tauri/src/grpc.rs @@ -1,10 +1,14 @@ use std::collections::BTreeMap; +use crate::error::Result; use KeyAndValueRef::{Ascii, Binary}; - +use tauri::{Manager, Runtime, WebviewWindow}; use yaak_grpc::{KeyAndValueRef, MetadataMap}; +use yaak_models::models::GrpcRequest; +use yaak_plugins::events::{CallHttpAuthenticationRequest, HttpHeader}; +use yaak_plugins::manager::PluginManager; -pub fn metadata_to_map(metadata: MetadataMap) -> BTreeMap { +pub(crate) fn metadata_to_map(metadata: MetadataMap) -> BTreeMap { let mut entries = BTreeMap::new(); for r in metadata.iter() { match r { @@ -14,3 +18,48 @@ pub fn metadata_to_map(metadata: MetadataMap) -> BTreeMap { } entries } + +pub(crate) async fn build_metadata( + window: &WebviewWindow, + request: &GrpcRequest, +) -> Result> { + let plugin_manager = window.state::(); + let mut metadata = BTreeMap::new(); + + // Add the rest of metadata + for h in request.clone().metadata { + if h.name.is_empty() && h.value.is_empty() { + continue; + } + + if !h.enabled { + continue; + } + + metadata.insert(h.name, h.value); + } + + if let Some(auth_name) = request.authentication_type.clone() { + let auth = request.authentication.clone(); + let plugin_req = CallHttpAuthenticationRequest { + context_id: format!("{:x}", md5::compute(request.id.clone())), + values: serde_json::from_value(serde_json::to_value(&auth).unwrap()).unwrap(), + method: "POST".to_string(), + url: request.url.clone(), + headers: metadata + .iter() + .map(|(name, value)| HttpHeader { + name: name.to_string(), + value: value.to_string(), + }) + .collect(), + }; + let plugin_result = + plugin_manager.call_http_authentication(&window, &auth_name, plugin_req).await?; + for header in plugin_result.set_headers { + metadata.insert(header.name, header.value); + } + } + + Ok(metadata) +} diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 8fcdaf28..000d3772 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -1,7 +1,7 @@ extern crate core; use crate::encoding::read_response_body; use crate::error::Error::GenericError; -use crate::grpc::metadata_to_map; +use crate::grpc::{build_metadata, metadata_to_map}; use crate::http_request::send_http_request; use crate::notifications::YaakNotifier; use crate::render::{render_grpc_request, render_template}; @@ -38,9 +38,9 @@ use yaak_models::util::{ BatchUpsertResult, UpdateSource, get_workspace_export_resources, maybe_gen_id, maybe_gen_id_opt, }; use yaak_plugins::events::{ - BootResponse, CallHttpAuthenticationRequest, CallHttpRequestActionRequest, FilterResponse, + BootResponse, CallHttpRequestActionRequest, FilterResponse, GetHttpAuthenticationConfigResponse, GetHttpAuthenticationSummaryResponse, - GetHttpRequestActionsResponse, GetTemplateFunctionsResponse, HttpHeader, InternalEvent, + GetHttpRequestActionsResponse, GetTemplateFunctionsResponse, InternalEvent, InternalEventPayload, JsonPrimitive, PluginWindowContext, RenderPurpose, }; use yaak_plugins::manager::PluginManager; @@ -166,6 +166,7 @@ async fn cmd_grpc_reflect( .await?; let uri = safe_uri(&req.url); + let metadata = build_metadata(&window, &req).await?; Ok(grpc_handle .lock() @@ -174,6 +175,7 @@ async fn cmd_grpc_reflect( &req.id, &uri, &proto_files.iter().map(|p| PathBuf::from_str(p).unwrap()).collect(), + &metadata, ) .await .map_err(|e| GenericError(e.to_string()))?) @@ -186,7 +188,6 @@ async fn cmd_grpc_go( proto_files: Vec, app_handle: AppHandle, window: WebviewWindow, - plugin_manager: State<'_, PluginManager>, grpc_handle: State<'_, Mutex>, ) -> YaakResult { let environment = match environment_id { @@ -208,42 +209,7 @@ async fn cmd_grpc_go( ) .await?; - let mut metadata = BTreeMap::new(); - - // Add the rest of metadata - for h in request.clone().metadata { - if h.name.is_empty() && h.value.is_empty() { - continue; - } - - if !h.enabled { - continue; - } - - metadata.insert(h.name, h.value); - } - - if let Some(auth_name) = request.authentication_type.clone() { - let auth = request.authentication.clone(); - let plugin_req = CallHttpAuthenticationRequest { - context_id: format!("{:x}", md5::compute(request_id.to_string())), - values: serde_json::from_value(serde_json::to_value(&auth).unwrap()).unwrap(), - method: "POST".to_string(), - url: request.url.clone(), - headers: metadata - .iter() - .map(|(name, value)| HttpHeader { - name: name.to_string(), - value: value.to_string(), - }) - .collect(), - }; - let plugin_result = - plugin_manager.call_http_authentication(&window, &auth_name, plugin_req).await?; - for header in plugin_result.set_headers { - metadata.insert(header.name, header.value); - } - } + let metadata = build_metadata(&window, &request).await?; let conn = app_handle.db().upsert_grpc_connection( &GrpcConnection { @@ -291,6 +257,7 @@ async fn cmd_grpc_go( &request.clone().id, uri.as_str(), &proto_files.iter().map(|p| PathBuf::from_str(p).unwrap()).collect(), + &metadata, ) .await; @@ -448,7 +415,7 @@ async fn cmd_grpc_go( match (method_desc.is_client_streaming(), method_desc.is_server_streaming()) { (true, true) => ( Some( - connection.streaming(&service, &method, in_msg_stream, metadata).await, + connection.streaming(&service, &method, in_msg_stream, &metadata).await, ), None, ), @@ -456,16 +423,16 @@ async fn cmd_grpc_go( None, Some( connection - .client_streaming(&service, &method, in_msg_stream, metadata) + .client_streaming(&service, &method, in_msg_stream, &metadata) .await, ), ), (false, true) => ( - Some(connection.server_streaming(&service, &method, &msg, metadata).await), + Some(connection.server_streaming(&service, &method, &msg, &metadata).await), None, ), (false, false) => { - (None, Some(connection.unary(&service, &method, &msg, metadata).await)) + (None, Some(connection.unary(&service, &method, &msg, &metadata).await)) } }; diff --git a/src-tauri/yaak-grpc/src/client.rs b/src-tauri/yaak-grpc/src/client.rs index c7d1afe6..5c185368 100644 --- a/src-tauri/yaak-grpc/src/client.rs +++ b/src-tauri/yaak-grpc/src/client.rs @@ -1,13 +1,15 @@ +use crate::manager::decorate_req; use crate::transport::get_transport; use async_recursion::async_recursion; 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::debug; +use std::collections::BTreeMap; use tokio_stream::StreamExt; +use tonic::Request; use tonic::body::BoxBody; use tonic::transport::Uri; -use tonic::Request; use tonic_reflection::pb::v1::server_reflection_request::MessageRequest; use tonic_reflection::pb::v1::server_reflection_response::MessageResponse; use tonic_reflection::pb::v1::{ @@ -44,6 +46,7 @@ impl AutoReflectionClient { pub async fn send_reflection_request( &mut self, message: MessageRequest, + metadata: &BTreeMap, ) -> Result { let reflection_request = ServerReflectionRequest { host: "".into(), // Doesn't matter @@ -51,7 +54,9 @@ impl AutoReflectionClient { }; if self.use_v1alpha { - let request = Request::new(tokio_stream::once(to_v1alpha_request(reflection_request))); + let mut request = Request::new(tokio_stream::once(to_v1alpha_request(reflection_request))); + decorate_req(metadata, &mut request).map_err(|e| e.to_string())?; + self.client_v1alpha .server_reflection_info(request) .await @@ -70,7 +75,9 @@ impl AutoReflectionClient { .ok_or("No reflection response".to_string()) .map(|resp| to_v1_msg_response(resp)) } else { - let request = Request::new(tokio_stream::once(reflection_request)); + let mut request = Request::new(tokio_stream::once(reflection_request)); + decorate_req(metadata, &mut request).map_err(|e| e.to_string())?; + let resp = self.client_v1.server_reflection_info(request).await; match resp { Ok(r) => Ok(r), @@ -79,7 +86,7 @@ impl AutoReflectionClient { // If v1 fails, change to v1alpha and try again debug!("gRPC schema reflection falling back to v1alpha"); self.use_v1alpha = true; - return self.send_reflection_request(message).await; + return self.send_reflection_request(message, metadata).await; } _ => Err(e), }, diff --git a/src-tauri/yaak-grpc/src/manager.rs b/src-tauri/yaak-grpc/src/manager.rs index a39b59a9..c4e19efc 100644 --- a/src-tauri/yaak-grpc/src/manager.rs +++ b/src-tauri/yaak-grpc/src/manager.rs @@ -69,7 +69,7 @@ impl GrpcConnection { service: &str, method: &str, message: &str, - metadata: BTreeMap, + metadata: &BTreeMap, ) -> Result, StreamError> { let method = &self.method(&service, &method)?; let input_message = method.input(); @@ -96,7 +96,7 @@ impl GrpcConnection { service: &str, method: &str, stream: ReceiverStream, - metadata: BTreeMap, + metadata: &BTreeMap, ) -> Result>, StreamError> { let method = &self.method(&service, &method)?; let mut client = tonic::client::Grpc::with_origin(self.conn.clone(), self.uri.clone()); @@ -116,7 +116,7 @@ impl GrpcConnection { service: &str, method: &str, stream: ReceiverStream, - metadata: BTreeMap, + metadata: &BTreeMap, ) -> Result, StreamError> { let method = &self.method(&service, &method)?; let mut client = tonic::client::Grpc::with_origin(self.conn.clone(), self.uri.clone()); @@ -137,7 +137,7 @@ impl GrpcConnection { service: &str, method: &str, message: &str, - metadata: BTreeMap, + metadata: &BTreeMap, ) -> Result>, StreamError> { let method = &self.method(&service, &method)?; let input_message = method.input(); @@ -180,10 +180,11 @@ impl GrpcHandle { id: &str, uri: &str, proto_files: &Vec, + metadata: &BTreeMap, ) -> Result<(), String> { let pool = if proto_files.is_empty() { let full_uri = uri_from_str(uri)?; - fill_pool_from_reflection(&full_uri).await + fill_pool_from_reflection(&full_uri, metadata).await } else { fill_pool_from_files(&self.app_handle, proto_files).await }?; @@ -197,9 +198,10 @@ impl GrpcHandle { id: &str, uri: &str, proto_files: &Vec, + metadata: &BTreeMap, ) -> Result, String> { // Ensure reflection is up-to-date - self.reflect(id, uri, proto_files).await?; + self.reflect(id, uri, proto_files, metadata).await?; let pool = self.get_pool(id, uri, proto_files).ok_or("Failed to get pool".to_string())?; Ok(self.services_from_pool(&pool)) @@ -235,8 +237,9 @@ impl GrpcHandle { id: &str, uri: &str, proto_files: &Vec, + metadata: &BTreeMap, ) -> Result { - self.reflect(id, uri, proto_files).await?; + self.reflect(id, uri, proto_files, metadata).await?; let pool = self.get_pool(id, uri, proto_files).ok_or("Failed to get pool")?; let uri = uri_from_str(uri)?; @@ -254,7 +257,10 @@ impl GrpcHandle { } } -fn decorate_req(metadata: BTreeMap, req: &mut Request) -> Result<(), String> { +pub(crate) fn decorate_req( + metadata: &BTreeMap, + req: &mut Request, +) -> Result<(), String> { for (k, v) in metadata { req.metadata_mut().insert( MetadataKey::from_str(k.as_str()).map_err(|e| e.to_string())?, diff --git a/src-tauri/yaak-grpc/src/reflection.rs b/src-tauri/yaak-grpc/src/reflection.rs index 27db0a88..1449a6d8 100644 --- a/src-tauri/yaak-grpc/src/reflection.rs +++ b/src-tauri/yaak-grpc/src/reflection.rs @@ -1,3 +1,4 @@ +use std::collections::BTreeMap; use std::env::temp_dir; use std::ops::Deref; use std::path::PathBuf; @@ -89,11 +90,14 @@ pub async fn fill_pool_from_files( Ok(pool) } -pub async fn fill_pool_from_reflection(uri: &Uri) -> Result { +pub async fn fill_pool_from_reflection( + uri: &Uri, + metadata: &BTreeMap, +) -> Result { let mut pool = DescriptorPool::new(); let mut client = AutoReflectionClient::new(uri); - for service in list_services(&mut client).await? { + for service in list_services(&mut client, metadata).await? { if service == "grpc.reflection.v1alpha.ServerReflection" { continue; } @@ -101,14 +105,18 @@ pub async fn fill_pool_from_reflection(uri: &Uri) -> Result Result, String> { - let response = client.send_reflection_request(MessageRequest::ListServices("".into())).await?; +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, @@ -122,9 +130,13 @@ 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())) + .send_reflection_request( + MessageRequest::FileContainingSymbol(service_name.into()), + metadata, + ) .await { Ok(resp) => resp, @@ -139,8 +151,13 @@ async fn file_descriptor_set_from_service_name( _ => panic!("Expected a FileDescriptorResponse variant"), }; - add_file_descriptors_to_pool(file_descriptor_response.file_descriptor_proto, pool, client) - .await; + add_file_descriptors_to_pool( + file_descriptor_response.file_descriptor_proto, + pool, + client, + metadata, + ) + .await; } #[async_recursion] @@ -148,6 +165,7 @@ 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()); @@ -165,7 +183,7 @@ async fn add_file_descriptors_to_pool( 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).await; + file_descriptor_set_by_filename(node.as_str(), pool, client, metadata).await; } } Err(_) => panic!("proto file got cycle!"), @@ -177,6 +195,7 @@ 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) { @@ -184,7 +203,7 @@ async fn file_descriptor_set_by_filename( } let msg = MessageRequest::FileByFilename(filename.into()); - let response = client.send_reflection_request(msg).await; + let response = client.send_reflection_request(msg, metadata).await; let file_descriptor_response = match response { Ok(MessageResponse::FileDescriptorResponse(resp)) => resp, Ok(_) => { @@ -196,8 +215,13 @@ async fn file_descriptor_set_by_filename( } }; - add_file_descriptors_to_pool(file_descriptor_response.file_descriptor_proto, pool, client) - .await; + add_file_descriptors_to_pool( + file_descriptor_response.file_descriptor_proto, + pool, + client, + metadata, + ) + .await; } pub fn method_desc_to_path(md: &MethodDescriptor) -> PathAndQuery {