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: Vec<PathBuf> = iface_files 82 .iter() 83 .map(|&path| root_dir.join(path)) 84 .collect(); 85 86 let include_dirs: Vec<PathBuf> = include_dirs 87 .iter() 88 .map(|&path| root_dir.join(path)) 89 .collect(); 90 let out_dir = root_dir.join(out_dir); 91 let file_descriptor_set_path = root_dir.join(file_descriptor_set_path); 92 93 let fds = protox::compile(&iface_files, &include_dirs).unwrap(); 94 95 write_fds(&fds, &file_descriptor_set_path); 96 97 tonic_build::configure() 98 .build_client(build_client) 99 .build_server(build_server) 100 .out_dir(&tempdir) 101 .compile_fds(fds) 102 .unwrap(); 103 104 for path in std::fs::read_dir(tempdir.path()).unwrap() { 105 let path = path.unwrap().path(); 106 let to = out_dir.join( 107 path.file_name() 108 .unwrap() 109 .to_str() 110 .unwrap() 111 .strip_suffix(".rs") 112 .unwrap() 113 .replace('.', "_") 114 + ".rs", 115 ); 116 std::fs::copy(&path, &to).unwrap(); 117 } 118 } 119 120 fn write_fds(fds: &FileDescriptorSet, path: &Path) { 121 const GENERATED_COMMENT: &str = "// This file is @generated by codegen."; 122 123 let mut fds = fds.clone(); 124 for fd in fds.file.iter_mut() { 125 fd.source_code_info = None; 126 } 127 128 let fds_raw = fds.encode_to_vec(); 129 let tokens = quote! { 130 /// Byte encoded FILE_DESCRIPTOR_SET. 131 pub const FILE_DESCRIPTOR_SET: &[u8] = &[#(#fds_raw),*]; 132 }; 133 let ast = syn::parse2(tokens).unwrap(); 134 let formatted = prettyplease::unparse(&ast); 135 136 let content = format!("{GENERATED_COMMENT}\n{formatted}"); 137 138 let mut writer = BufWriter::new(File::create(path).unwrap()); 139 writer.write_all(content.as_bytes()).unwrap(); 140 } 141