folk_plugin_grpc/

plugin.rs

use std::collections::BTreeSet;
use std::path::{Path, PathBuf};

use anyhow::{Context, Result};
use async_trait::async_trait;
use axum::Router;
use folk_api::{PluginContext, RpcMethodDef, ServerPlugin};
use tonic::transport::Server;
use tracing::{debug, info, warn};

use crate::config::GrpcConfig;
use crate::service::{GrpcState, grpc_handler};

pub struct GrpcPlugin {
    config: GrpcConfig,
}

impl GrpcPlugin {
    pub fn new(config: GrpcConfig) -> Self {
        Self { config }
    }
}

#[async_trait]
impl ServerPlugin for GrpcPlugin {
    fn name(&self) -> &'static str {
        "grpc"
    }

    async fn run(&self, ctx: PluginContext) -> Result<()> {
        let state = GrpcState {
            executor: ctx.executor.clone(),
        };

        let router: axum::Router = Router::new().fallback(grpc_handler).with_state(state);
        let mut sd = ctx.shutdown.clone();

        info!(listen = %self.config.listen, "gRPC server listening");

        let routes = tonic::service::Routes::from(router);

        if !self.config.proto.is_empty() {
            match build_reflection_descriptors(&self.config.proto) {
                Ok(encoded_fds) => {
                    info!(proto_files = self.config.proto.len(), "gRPC reflection enabled");
                    let reflection = tonic_reflection::server::Builder::configure()
                        .register_encoded_file_descriptor_set(&encoded_fds)
                        .build_v1()
                        .context("build reflection service")?;
                    Server::builder()
                        .add_routes(routes)
                        .add_service(reflection)
                        .serve_with_shutdown(self.config.listen, async move {
                            sd.changed().await.ok();
                        })
                        .await?;
                    return Ok(());
                },
                Err(e) => {
                    warn!(error = %e, "failed to build reflection; starting without it");
                },
            }
        }

        Server::builder()
            .add_routes(routes)
            .serve_with_shutdown(self.config.listen, async move {
                sd.changed().await.ok();
            })
            .await?;

        Ok(())
    }

    fn rpc_methods(&self) -> Vec<RpcMethodDef> {
        vec![RpcMethodDef::new(
            "grpc.services",
            "list registered gRPC service names",
        )]
    }
}

/// Compile proto files and return encoded FileDescriptorSet.
///
/// Import paths are resolved automatically:
/// - Parses `import` statements from each proto file
/// - Searches for imported files within the project root (cwd)
/// - Never searches outside the working directory
fn build_reflection_descriptors(proto_files: &[String]) -> Result<Vec<u8>> {
    let cwd = std::env::current_dir().context("get cwd")?;
    let mut include_paths = BTreeSet::new();

    // Resolve imports starting from each proto file
    let mut visited = BTreeSet::new();
    for proto in proto_files {
        let abs = cwd.join(proto);
        resolve_imports(&abs, &cwd, &mut include_paths, &mut visited);
    }

    // Always include the directory of each proto file itself
    for proto in proto_files {
        let abs = cwd.join(proto);
        if let Some(parent) = abs.parent() {
            include_paths.insert(parent.to_path_buf());
        }
    }

    debug!(
        paths = include_paths.len(),
        "resolved proto include paths"
    );

    let paths: Vec<_> = include_paths.iter().collect();
    let mut compiler = protox::Compiler::new(paths.iter().map(|p| p.as_path()))
        .context("create protox compiler")?;
    compiler.include_imports(true);

    for proto in proto_files {
        let path = Path::new(proto);
        let file_name = path
            .file_name()
            .and_then(|n| n.to_str())
            .context("invalid proto file name")?;
        compiler
            .open_file(file_name)
            .with_context(|| format!("compile {proto}"))?;
    }

    let fds = compiler.file_descriptor_set();
    Ok(prost::Message::encode_to_vec(&fds))
}

/// Parse a proto file for `import` statements and find each imported file
/// within the project root. Adds the containing directory to include_paths.
/// Recurses into found imports.
fn resolve_imports(
    proto_path: &Path,
    project_root: &Path,
    include_paths: &mut BTreeSet<PathBuf>,
    visited: &mut BTreeSet<PathBuf>,
) {
    if !proto_path.is_file() || !visited.insert(proto_path.to_path_buf()) {
        return;
    }

    let content = match std::fs::read_to_string(proto_path) {
        Ok(c) => c,
        Err(_) => return,
    };

    for line in content.lines() {
        let line = line.trim();
        // Match: import "path/to/file.proto";
        // Match: import public "path/to/file.proto";
        if !line.starts_with("import ") {
            continue;
        }
        let Some(start) = line.find('"') else { continue };
        let Some(end) = line[start + 1..].find('"') else { continue };
        let import_path = &line[start + 1..start + 1 + end];

        // Skip well-known google types — protox has them built-in
        if import_path.starts_with("google/protobuf/") {
            continue;
        }

        // Find the best match for this import within the project
        if let Some(found) = find_best_match(project_root, import_path) {
            // Compute the include base: found = base / import_path
            if let Some(base) = found
                .to_str()
                .and_then(|f| f.strip_suffix(import_path))
                .map(PathBuf::from)
            {
                if !base.as_os_str().is_empty() {
                    include_paths.insert(base);
                }
            } else if let Some(parent) = found.parent() {
                include_paths.insert(parent.to_path_buf());
            }

            resolve_imports(&found, project_root, include_paths, visited);
        }
    }
}

/// Find the best match for a proto import within the project tree.
/// If multiple files match, picks the largest one (avoids stubs/overlays).
fn find_best_match(root: &Path, relative: &str) -> Option<PathBuf> {
    let mut candidates = Vec::new();

    // Direct from root
    let direct = root.join(relative);
    if direct.is_file() {
        candidates.push(direct);
    }

    find_all_recursive(root, relative, &mut candidates);

    if candidates.len() <= 1 {
        return candidates.into_iter().next();
    }

    // Multiple matches — pick the largest file (real source, not a stub)
    candidates
        .into_iter()
        .max_by_key(|p| std::fs::metadata(p).map(|m| m.len()).unwrap_or(0))
}

fn find_all_recursive(dir: &Path, relative: &str, results: &mut Vec<PathBuf>) {
    let candidate = dir.join(relative);
    if candidate.is_file() && !results.contains(&candidate) {
        results.push(candidate);
    }

    let entries = match std::fs::read_dir(dir) {
        Ok(e) => e,
        Err(_) => return,
    };

    for entry in entries.flatten() {
        let path = entry.path();
        if path.is_dir() {
            // Only skip hidden directories
            let name = entry.file_name();
            if name.to_str().is_some_and(|n| n.starts_with('.')) {
                continue;
            }
            find_all_recursive(&path, relative, results);
        }
    }
}