crank/crates/crank-adapter-grpc/src/client.rs

use std::{collections::BTreeMap, str::FromStr, time::Duration};

use base64::{Engine as _, engine::general_purpose::STANDARD};
use crank_core::GrpcTarget;
use futures_util::StreamExt;
use prost::Message;
use prost_reflect::{DescriptorPool, MethodDescriptor, prost_types::FileDescriptorSet};
use serde_json::json;
use tokio::time::{Instant, timeout_at};
use tonic::{
    Request,
    client::Grpc,
    metadata::{KeyAndValueRef, MetadataKey, MetadataValue},
    transport::Endpoint,
};

use crate::{
    GrpcAdapterError, GrpcRequest, GrpcResponse, GrpcWindowRequest, GrpcWindowResponse,
    codec::JsonCodec,
};

#[derive(Clone, Debug, Default)]
pub struct GrpcAdapter;

impl GrpcAdapter {
    pub fn new() -> Self {
        Self
    }

    pub async fn execute(
        &self,
        target: &GrpcTarget,
        request: &GrpcRequest,
    ) -> Result<GrpcResponse, GrpcAdapterError> {
        let method = resolve_method(target)?;
        if method.is_client_streaming() || method.is_server_streaming() {
            return Err(GrpcAdapterError::UnsupportedMethodKind {
                service: format!("{}.{}", target.package, target.service),
                method: target.method.clone(),
            });
        }

        let mut grpc = connect(target, request.timeout_ms).await?;
        let path = build_path(&method)?;
        let codec = JsonCodec::new(method.input(), method.output());
        let mut tonic_request = Request::new(request.body.clone());

        apply_headers(&mut tonic_request, &request.headers)?;

        let response = grpc.unary(tonic_request, path, codec).await?;
        let headers = normalize_headers(response.metadata());
        let body = response.into_inner();

        Ok(GrpcResponse {
            status_code: 200,
            headers,
            body,
        })
    }

    pub async fn execute_window(
        &self,
        target: &GrpcTarget,
        request: &GrpcWindowRequest,
    ) -> Result<GrpcWindowResponse, GrpcAdapterError> {
        let method = resolve_method(target)?;
        if method.is_client_streaming() || !method.is_server_streaming() {
            return Err(GrpcAdapterError::UnsupportedStreamingMethodKind {
                service: format!("{}.{}", target.package, target.service),
                method: target.method.clone(),
            });
        }

        let mut grpc = connect(target, request.request.timeout_ms).await?;
        let path = build_path(&method)?;
        let codec = JsonCodec::new(method.input(), method.output());
        let mut tonic_request = Request::new(request.request.body.clone());

        apply_headers(&mut tonic_request, &request.request.headers)?;

        let response = grpc.server_streaming(tonic_request, path, codec).await?;
        let headers = normalize_headers(response.metadata());
        let mut stream = response.into_inner();
        let deadline = Instant::now() + Duration::from_millis(request.window_duration_ms);
        let mut items = Vec::new();
        let mut done = true;

        loop {
            if request
                .max_items
                .is_some_and(|limit| items.len() >= limit as usize)
            {
                done = false;
                break;
            }

            let next_message = match timeout_at(deadline, stream.next()).await {
                Ok(next_message) => next_message,
                Err(_) => break,
            };

            let Some(next_message) = next_message else {
                break;
            };

            let message = next_message?;
            items.push(message);
        }

        Ok(GrpcWindowResponse {
            status_code: 200,
            headers,
            body: json!({
                "items": items,
                "done": done,
            }),
        })
    }
}

async fn connect(
    target: &GrpcTarget,
    timeout_ms: u64,
) -> Result<Grpc<tonic::transport::Channel>, GrpcAdapterError> {
    let endpoint = Endpoint::from_shared(target.server_addr.clone())?
        .timeout(Duration::from_millis(timeout_ms));
    let channel = endpoint.connect().await?;
    let mut grpc = Grpc::new(channel);
    grpc.ready()
        .await
        .map_err(|error| GrpcAdapterError::Status {
            code: tonic::Code::Unavailable,
            message: error.to_string(),
        })?;

    Ok(grpc)
}

fn build_path(
    method: &MethodDescriptor,
) -> Result<tonic::codegen::http::uri::PathAndQuery, GrpcAdapterError> {
    let service_name = method.parent_service().full_name().to_owned();
    let method_name = method.name().to_owned();
    tonic::codegen::http::uri::PathAndQuery::from_str(&format!("/{service_name}/{method_name}"))
        .map_err(|_| GrpcAdapterError::InvalidMethodPath {
            service: service_name,
            method: method_name,
        })
}

fn resolve_method(target: &GrpcTarget) -> Result<MethodDescriptor, GrpcAdapterError> {
    let bytes = STANDARD
        .decode(&target.descriptor_set_b64)
        .map_err(|_| GrpcAdapterError::InvalidDescriptorEncoding)?;
    let descriptor_set = FileDescriptorSet::decode(bytes.as_slice())
        .map_err(|_| GrpcAdapterError::InvalidDescriptorSet)?;
    let pool = DescriptorPool::from_file_descriptor_set(descriptor_set)
        .map_err(|_| GrpcAdapterError::InvalidDescriptorPool)?;
    let full_service_name = if target.package.is_empty() {
        target.service.clone()
    } else {
        format!("{}.{}", target.package, target.service)
    };
    let service = pool
        .get_service_by_name(&full_service_name)
        .ok_or_else(|| GrpcAdapterError::ServiceNotFound {
            service: full_service_name.clone(),
        })?;

    service
        .methods()
        .find(|method| method.name() == target.method)
        .ok_or_else(|| GrpcAdapterError::MethodNotFound {
            service: full_service_name,
            method: target.method.clone(),
        })
}

fn apply_headers(
    request: &mut Request<serde_json::Value>,
    headers: &BTreeMap<String, String>,
) -> Result<(), GrpcAdapterError> {
    for (key, value) in headers {
        let metadata_key =
            MetadataKey::from_str(key).map_err(|source| GrpcAdapterError::InvalidMetadataKey {
                key: key.clone(),
                source,
            })?;
        let metadata_value = MetadataValue::from_str(value).map_err(|source| {
            GrpcAdapterError::InvalidMetadataValue {
                key: key.clone(),
                source,
            }
        })?;

        request.metadata_mut().insert(metadata_key, metadata_value);
    }

    Ok(())
}

fn normalize_headers(metadata: &tonic::metadata::MetadataMap) -> BTreeMap<String, String> {
    metadata
        .iter()
        .filter_map(|entry| match entry {
            KeyAndValueRef::Ascii(key, value) => {
                Some((key.as_str().to_owned(), value.to_str().ok()?.to_owned()))
            }
            KeyAndValueRef::Binary(_, _) => None,
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;

    use crank_core::{DescriptorId, GrpcTarget};
    use serde_json::json;

    use crate::{GrpcAdapter, GrpcAdapterError, GrpcRequest, GrpcWindowRequest, test_support};

    #[tokio::test]
    async fn executes_unary_grpc_request() {
        let server_addr = test_support::spawn_unary_echo_server().await;
        let adapter = GrpcAdapter::new();
        let target = GrpcTarget {
            server_addr,
            package: "echo".to_owned(),
            service: "EchoService".to_owned(),
            method: "UnaryEcho".to_owned(),
            descriptor_ref: DescriptorId::new("desc_echo"),
            descriptor_set_b64: test_support::descriptor_set_b64(),
        };
        let request = GrpcRequest {
            headers: BTreeMap::new(),
            body: json!({ "message": "hello" }),
            timeout_ms: 1_000,
        };

        let response = adapter.execute(&target, &request).await.unwrap();

        assert_eq!(response.body, json!({ "message": "hello" }));
    }

    #[tokio::test]
    async fn collects_server_stream_messages_with_window_bounds() {
        let server_addr = test_support::spawn_unary_echo_server().await;
        let adapter = GrpcAdapter::new();
        let target = GrpcTarget {
            server_addr,
            package: "echo".to_owned(),
            service: "EchoService".to_owned(),
            method: "ServerEcho".to_owned(),
            descriptor_ref: DescriptorId::new("desc_echo"),
            descriptor_set_b64: test_support::descriptor_set_b64(),
        };
        let request = GrpcWindowRequest {
            request: GrpcRequest {
                headers: BTreeMap::new(),
                body: json!({ "message": "hello" }),
                timeout_ms: 1_000,
            },
            window_duration_ms: 1_000,
            max_items: Some(2),
        };

        let response = adapter.execute_window(&target, &request).await.unwrap();

        assert_eq!(
            response.body,
            json!({
                "items": [
                    { "message": "hello-one" },
                    { "message": "hello-two" }
                ],
                "done": false
            })
        );
    }

    #[tokio::test]
    async fn rejects_unary_method_in_window_mode() {
        let server_addr = test_support::spawn_unary_echo_server().await;
        let adapter = GrpcAdapter::new();
        let target = GrpcTarget {
            server_addr,
            package: "echo".to_owned(),
            service: "EchoService".to_owned(),
            method: "UnaryEcho".to_owned(),
            descriptor_ref: DescriptorId::new("desc_echo"),
            descriptor_set_b64: test_support::descriptor_set_b64(),
        };
        let request = GrpcWindowRequest {
            request: GrpcRequest {
                headers: BTreeMap::new(),
                body: json!({ "message": "hello" }),
                timeout_ms: 1_000,
            },
            window_duration_ms: 1_000,
            max_items: Some(10),
        };

        let error = adapter.execute_window(&target, &request).await.unwrap_err();

        assert!(matches!(
            error,
            GrpcAdapterError::UnsupportedStreamingMethodKind { .. }
        ));
    }
}