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
main()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
codegen( root_dir: &Path, iface_files: &[&str], include_dirs: &[&str], out_dir: &Path, file_descriptor_set_path: &Path, build_client: bool, build_server: bool, )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
write_fds(fds: &FileDescriptorSet, path: &Path)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