xref: /tonic/codegen/src/main.rs (revision 6856ffe3)

use std::{
    fs::File,
    io::{BufWriter, Write as _},
    path::{Path, PathBuf},
};

use protox::prost::Message as _;
use quote::quote;
use tonic_build::FileDescriptorSet;

fn main() {
    // tonic-health
    codegen(
        &PathBuf::from(std::env!("CARGO_MANIFEST_DIR"))
            .parent()
            .unwrap()
            .join("tonic-health"),
        &["proto/health.proto"],
        &["proto"],
        &PathBuf::from("src/generated"),
        &PathBuf::from("src/generated/grpc_health_v1_fds.rs"),
        true,
        true,
    );

    // tonic-reflection
    codegen(
        &PathBuf::from(std::env!("CARGO_MANIFEST_DIR"))
            .parent()
            .unwrap()
            .join("tonic-reflection"),
        &["proto/reflection_v1.proto"],
        &["proto"],
        &PathBuf::from("src/generated"),
        &PathBuf::from("src/generated/reflection_v1_fds.rs"),
        true,
        true,
    );
    codegen(
        &PathBuf::from(std::env!("CARGO_MANIFEST_DIR"))
            .parent()
            .unwrap()
            .join("tonic-reflection"),
        &["proto/reflection_v1alpha.proto"],
        &["proto"],
        &PathBuf::from("src/generated"),
        &PathBuf::from("src/generated/reflection_v1alpha1_fds.rs"),
        true,
        true,
    );

    // tonic-types
    codegen(
        &PathBuf::from(std::env!("CARGO_MANIFEST_DIR"))
            .parent()
            .unwrap()
            .join("tonic-types"),
        &["proto/status.proto", "proto/error_details.proto"],
        &["proto"],
        &PathBuf::from("src/generated"),
        &PathBuf::from("src/generated/types_fds.rs"),
        false,
        false,
    );
}

fn codegen(
    root_dir: &Path,
    iface_files: &[&str],
    include_dirs: &[&str],
    out_dir: &Path,
    file_descriptor_set_path: &Path,
    build_client: bool,
    build_server: bool,
) {
    let tempdir = tempfile::Builder::new()
        .prefix("tonic-codegen-")
        .tempdir()
        .unwrap();

    let iface_files = iface_files.iter().map(|&path| root_dir.join(path));
    let include_dirs = include_dirs.iter().map(|&path| root_dir.join(path));
    let out_dir = root_dir.join(out_dir);
    let file_descriptor_set_path = root_dir.join(file_descriptor_set_path);

    let fds = protox::compile(iface_files, include_dirs).unwrap();

    write_fds(&fds, &file_descriptor_set_path);

    tonic_build::configure()
        .build_client(build_client)
        .build_server(build_server)
        .out_dir(&tempdir)
        .compile_fds(fds)
        .unwrap();

    for path in std::fs::read_dir(tempdir.path()).unwrap() {
        let path = path.unwrap().path();
        let to = out_dir.join(
            path.file_name()
                .unwrap()
                .to_str()
                .unwrap()
                .strip_suffix(".rs")
                .unwrap()
                .replace('.', "_")
                + ".rs",
        );
        std::fs::copy(&path, &to).unwrap();
    }
}

fn write_fds(fds: &FileDescriptorSet, path: &Path) {
    const GENERATED_COMMENT: &str = "// This file is @generated by codegen.";

    let mut fds = fds.clone();
    for fd in fds.file.iter_mut() {
        fd.source_code_info = None;
    }

    let fds_raw = fds.encode_to_vec();
    let tokens = quote! {
        /// Byte encoded FILE_DESCRIPTOR_SET.
        pub const FILE_DESCRIPTOR_SET: &[u8] = &[#(#fds_raw),*];
    };
    let ast = syn::parse2(tokens).unwrap();
    let formatted = prettyplease::unparse(&ast);

    let content = format!("{GENERATED_COMMENT}\n{formatted}");

    let mut writer = BufWriter::new(File::create(path).unwrap());
    writer.write_all(content.as_bytes()).unwrap();
}