1 use std::{ 2 fs::File, 3 io::{BufWriter, Write as _}, 4 path::{Path, PathBuf}, 5 }; 6 7 use protox::prost::Message as _; 8 use quote::quote; 9 use tonic_build::FileDescriptorSet; 10 11 fn main() { 12 // tonic-health 13 codegen( 14 &PathBuf::from(std::env!("CARGO_MANIFEST_DIR")) 15 .parent() 16 .unwrap() 17 .join("tonic-health"), 18 &["proto/health.proto"], 19 &["proto"], 20 &PathBuf::from("src/generated"), 21 &PathBuf::from("src/generated/grpc_health_v1_fds.rs"), 22 true, 23 true, 24 ); 25 26 // tonic-reflection 27 codegen( 28 &PathBuf::from(std::env!("CARGO_MANIFEST_DIR")) 29 .parent() 30 .unwrap() 31 .join("tonic-reflection"), 32 &["proto/reflection_v1.proto"], 33 &["proto"], 34 &PathBuf::from("src/generated"), 35 &PathBuf::from("src/generated/reflection_v1_fds.rs"), 36 true, 37 true, 38 ); 39 codegen( 40 &PathBuf::from(std::env!("CARGO_MANIFEST_DIR")) 41 .parent() 42 .unwrap() 43 .join("tonic-reflection"), 44 &["proto/reflection_v1alpha.proto"], 45 &["proto"], 46 &PathBuf::from("src/generated"), 47 &PathBuf::from("src/generated/reflection_v1alpha1_fds.rs"), 48 true, 49 true, 50 ); 51 52 // tonic-types 53 codegen( 54 &PathBuf::from(std::env!("CARGO_MANIFEST_DIR")) 55 .parent() 56 .unwrap() 57 .join("tonic-types"), 58 &["proto/status.proto", "proto/error_details.proto"], 59 &["proto"], 60 &PathBuf::from("src/generated"), 61 &PathBuf::from("src/generated/types_fds.rs"), 62 false, 63 false, 64 ); 65 } 66 67 fn codegen( 68 root_dir: &Path, 69 iface_files: &[&str], 70 include_dirs: &[&str], 71 out_dir: &Path, 72 file_descriptor_set_path: &Path, 73 build_client: bool, 74 build_server: bool, 75 ) { 76 let tempdir = tempfile::Builder::new() 77 .prefix("tonic-codegen-") 78 .tempdir() 79 .unwrap(); 80 81 let iface_files = iface_files.iter().map(|&path| root_dir.join(path)); 82 let include_dirs = include_dirs.iter().map(|&path| root_dir.join(path)); 83 let out_dir = root_dir.join(out_dir); 84 let file_descriptor_set_path = root_dir.join(file_descriptor_set_path); 85 86 let fds = protox::compile(iface_files, include_dirs).unwrap(); 87 88 write_fds(&fds, &file_descriptor_set_path); 89 90 tonic_build::configure() 91 .build_client(build_client) 92 .build_server(build_server) 93 .out_dir(&tempdir) 94 .compile_fds(fds) 95 .unwrap(); 96 97 for path in std::fs::read_dir(tempdir.path()).unwrap() { 98 let path = path.unwrap().path(); 99 let to = out_dir.join( 100 path.file_name() 101 .unwrap() 102 .to_str() 103 .unwrap() 104 .strip_suffix(".rs") 105 .unwrap() 106 .replace('.', "_") 107 + ".rs", 108 ); 109 std::fs::copy(&path, &to).unwrap(); 110 } 111 } 112 113 fn write_fds(fds: &FileDescriptorSet, path: &Path) { 114 const GENERATED_COMMENT: &str = "// This file is @generated by codegen."; 115 116 let mut fds = fds.clone(); 117 for fd in fds.file.iter_mut() { 118 fd.source_code_info = None; 119 } 120 121 let fds_raw = fds.encode_to_vec(); 122 let tokens = quote! { 123 /// Byte encoded FILE_DESCRIPTOR_SET. 124 pub const FILE_DESCRIPTOR_SET: &[u8] = &[#(#fds_raw),*]; 125 }; 126 let ast = syn::parse2(tokens).unwrap(); 127 let formatted = prettyplease::unparse(&ast); 128 129 let content = format!("{GENERATED_COMMENT}\n{formatted}"); 130 131 let mut writer = BufWriter::new(File::create(path).unwrap()); 132 writer.write_all(content.as_bytes()).unwrap(); 133 } 134