Decouple core Yaak logic from Tauri (#354)

crates/yaak-grpc/src/client.rs

@@ -0,0 +1,182 @@
+use crate::error::Error::GenericError;
+use crate::error::Result;
+use crate::manager::decorate_req;
+use crate::transport::get_transport;
+use async_recursion::async_recursion;
+use hyper_rustls::HttpsConnector;
+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_reflection::pb::v1::server_reflection_request::MessageRequest;
+use tonic_reflection::pb::v1::server_reflection_response::MessageResponse;
+use tonic_reflection::pb::v1::{
+    ErrorResponse, ExtensionNumberResponse, ListServiceResponse, ServerReflectionRequest,
+    ServiceResponse,
+};
+use tonic_reflection::pb::v1::{ExtensionRequest, FileDescriptorResponse};
+use tonic_reflection::pb::{v1, v1alpha};
+use yaak_tls::ClientCertificateConfig;
+
+pub struct AutoReflectionClient<T = Client<HttpsConnector<HttpConnector>, BoxBody>> {
+    use_v1alpha: bool,
+    client_v1: v1::server_reflection_client::ServerReflectionClient<T>,
+    client_v1alpha: v1alpha::server_reflection_client::ServerReflectionClient<T>,
+}
+
+impl AutoReflectionClient {
+    pub fn new(
+        uri: &Uri,
+        validate_certificates: bool,
+        client_cert: Option<ClientCertificateConfig>,
+    ) -> Result<Self> {
+        let client_v1 = v1::server_reflection_client::ServerReflectionClient::with_origin(
+            get_transport(validate_certificates, client_cert.clone())?,
+            uri.clone(),
+        );
+        let client_v1alpha = v1alpha::server_reflection_client::ServerReflectionClient::with_origin(
+            get_transport(validate_certificates, client_cert.clone())?,
+            uri.clone(),
+        );
+        Ok(AutoReflectionClient { use_v1alpha: false, client_v1, client_v1alpha })
+    }
+
+    #[async_recursion]
+    pub async fn send_reflection_request(
+        &mut self,
+        message: MessageRequest,
+        metadata: &BTreeMap<String, String>,
+    ) -> Result<MessageResponse> {
+        let reflection_request = ServerReflectionRequest {
+            host: "".into(), // Doesn't matter
+            message_request: Some(message.clone()),
+        };
+
+        if self.use_v1alpha {
+            let mut request =
+                Request::new(tokio_stream::once(to_v1alpha_request(reflection_request)));
+            decorate_req(metadata, &mut request)?;
+
+            self.client_v1alpha
+                .server_reflection_info(request)
+                .await
+                .map_err(|e| match e.code() {
+                    tonic::Code::Unavailable => {
+                        GenericError("Failed to connect to endpoint".to_string())
+                    }
+                    tonic::Code::Unauthenticated => {
+                        GenericError("Authentication failed".to_string())
+                    }
+                    tonic::Code::DeadlineExceeded => GenericError("Deadline exceeded".to_string()),
+                    _ => GenericError(e.to_string()),
+                })?
+                .into_inner()
+                .next()
+                .await
+                .ok_or(GenericError("Missing reflection message".to_string()))??
+                .message_response
+                .ok_or(GenericError("No reflection response".to_string()))
+                .map(|resp| to_v1_msg_response(resp))
+        } else {
+            let mut request = Request::new(tokio_stream::once(reflection_request));
+            decorate_req(metadata, &mut request)?;
+
+            let resp = self.client_v1.server_reflection_info(request).await;
+            match resp {
+                Ok(r) => Ok(r),
+                Err(e) => match e.code().clone() {
+                    tonic::Code::Unimplemented => {
+                        // 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, metadata).await;
+                    }
+                    _ => Err(e),
+                },
+            }
+            .map_err(|e| match e.code() {
+                tonic::Code::Unavailable => {
+                    GenericError("Failed to connect to endpoint".to_string())
+                }
+                tonic::Code::Unauthenticated => GenericError("Authentication failed".to_string()),
+                tonic::Code::DeadlineExceeded => GenericError("Deadline exceeded".to_string()),
+                _ => GenericError(e.to_string()),
+            })?
+            .into_inner()
+            .next()
+            .await
+            .ok_or(GenericError("Missing reflection message".to_string()))??
+            .message_response
+            .ok_or(GenericError("No reflection response".to_string()))
+        }
+    }
+}
+
+fn to_v1_msg_response(
+    response: v1alpha::server_reflection_response::MessageResponse,
+) -> MessageResponse {
+    match response {
+        v1alpha::server_reflection_response::MessageResponse::FileDescriptorResponse(v) => {
+            MessageResponse::FileDescriptorResponse(FileDescriptorResponse {
+                file_descriptor_proto: v.file_descriptor_proto,
+            })
+        }
+        v1alpha::server_reflection_response::MessageResponse::AllExtensionNumbersResponse(v) => {
+            MessageResponse::AllExtensionNumbersResponse(ExtensionNumberResponse {
+                extension_number: v.extension_number,
+                base_type_name: v.base_type_name,
+            })
+        }
+        v1alpha::server_reflection_response::MessageResponse::ListServicesResponse(v) => {
+            MessageResponse::ListServicesResponse(ListServiceResponse {
+                service: v
+                    .service
+                    .iter()
+                    .map(|s| ServiceResponse { name: s.name.clone() })
+                    .collect(),
+            })
+        }
+        v1alpha::server_reflection_response::MessageResponse::ErrorResponse(v) => {
+            MessageResponse::ErrorResponse(ErrorResponse {
+                error_code: v.error_code,
+                error_message: v.error_message,
+            })
+        }
+    }
+}
+
+fn to_v1alpha_request(request: ServerReflectionRequest) -> v1alpha::ServerReflectionRequest {
+    v1alpha::ServerReflectionRequest {
+        host: request.host,
+        message_request: request.message_request.map(|m| to_v1alpha_msg_request(m)),
+    }
+}
+
+fn to_v1alpha_msg_request(
+    message: MessageRequest,
+) -> v1alpha::server_reflection_request::MessageRequest {
+    match message {
+        MessageRequest::FileByFilename(v) => {
+            v1alpha::server_reflection_request::MessageRequest::FileByFilename(v)
+        }
+        MessageRequest::FileContainingSymbol(v) => {
+            v1alpha::server_reflection_request::MessageRequest::FileContainingSymbol(v)
+        }
+        MessageRequest::FileContainingExtension(ExtensionRequest {
+            extension_number,
+            containing_type,
+        }) => v1alpha::server_reflection_request::MessageRequest::FileContainingExtension(
+            v1alpha::ExtensionRequest { extension_number, containing_type },
+        ),
+        MessageRequest::AllExtensionNumbersOfType(v) => {
+            v1alpha::server_reflection_request::MessageRequest::AllExtensionNumbersOfType(v)
+        }
+        MessageRequest::ListServices(v) => {
+            v1alpha::server_reflection_request::MessageRequest::ListServices(v)
+        }
+    }
+}