|
| 1 | +use std::ffi::OsStr; |
| 2 | +use std::path::PathBuf; |
| 3 | + |
1 | 4 | use fs_err as fs; |
2 | | -use miette::{Context, IntoDiagnostic}; |
| 5 | +use miette::{IntoDiagnostic, miette}; |
3 | 6 | use protox::prost::Message; |
4 | 7 |
|
5 | | -const RPC_PROTO: &str = "rpc.proto"; |
6 | | -// Unified internal store API (store.Rpc, store.BlockProducer, store.NtxBuilder). |
7 | | -// We compile the same file three times to preserve existing descriptor names. |
8 | | -const STORE_RPC_PROTO: &str = "internal/store.proto"; |
9 | | -const STORE_NTX_BUILDER_PROTO: &str = "internal/store.proto"; |
10 | | -const STORE_BLOCK_PRODUCER_PROTO: &str = "internal/store.proto"; |
11 | | -const BLOCK_PRODUCER_PROTO: &str = "internal/block_producer.proto"; |
12 | | -const REMOTE_PROVER_PROTO: &str = "remote_prover.proto"; |
13 | | -const VALIDATOR_PROTO: &str = "internal/validator.proto"; |
14 | | - |
15 | | -const RPC_DESCRIPTOR: &str = "rpc_file_descriptor.bin"; |
16 | | -const STORE_RPC_DESCRIPTOR: &str = "store_rpc_file_descriptor.bin"; |
17 | | -const STORE_NTX_BUILDER_DESCRIPTOR: &str = "store_ntx_builder_file_descriptor.bin"; |
18 | | -const STORE_BLOCK_PRODUCER_DESCRIPTOR: &str = "store_block_producer_file_descriptor.bin"; |
19 | | -const BLOCK_PRODUCER_DESCRIPTOR: &str = "block_producer_file_descriptor.bin"; |
20 | | -const REMOTE_PROVER_DESCRIPTOR: &str = "remote_prover_file_descriptor.bin"; |
21 | | -const VALIDATOR_DESCRIPTOR: &str = "validator_file_descriptor.bin"; |
22 | | - |
23 | | -/// Generates Rust protobuf bindings from .proto files. |
| 8 | +/// Compiles each gRPC service definitions into a |
| 9 | +/// [`FileDescriptorSet`](tonic_prost_build::FileDescriptorSet) and exposes it as a function: |
24 | 10 | /// |
25 | | -/// This is done only if `BUILD_PROTO` environment variable is set to `1` to avoid running the |
26 | | -/// script on crates.io where repo-level .proto files are not available. |
| 11 | +/// ```rust |
| 12 | +/// fn <service>_api_descriptor() -> FileDescriptorSet; |
| 13 | +/// ``` |
27 | 14 | fn main() -> miette::Result<()> { |
28 | 15 | build_rs::output::rerun_if_changed("./proto"); |
| 16 | + build_rs::output::rerun_if_changed("Cargo.toml"); |
29 | 17 |
|
30 | 18 | let out_dir = build_rs::input::out_dir(); |
31 | | - let crate_root = build_rs::input::cargo_manifest_dir(); |
32 | | - let proto_src_dir = crate_root.join("proto"); |
33 | | - let includes = &[proto_src_dir]; |
34 | | - |
35 | | - let rpc_file_descriptor = protox::compile([RPC_PROTO], includes)?; |
36 | | - let rpc_path = out_dir.join(RPC_DESCRIPTOR); |
37 | | - fs::write(&rpc_path, rpc_file_descriptor.encode_to_vec()) |
38 | | - .into_diagnostic() |
39 | | - .wrap_err("writing rpc file descriptor")?; |
40 | | - |
41 | | - let remote_prover_file_descriptor = protox::compile([REMOTE_PROVER_PROTO], includes)?; |
42 | | - let remote_prover_path = out_dir.join(REMOTE_PROVER_DESCRIPTOR); |
43 | | - fs::write(&remote_prover_path, remote_prover_file_descriptor.encode_to_vec()) |
44 | | - .into_diagnostic() |
45 | | - .wrap_err("writing remote prover file descriptor")?; |
46 | | - |
47 | | - let store_rpc_file_descriptor = protox::compile([STORE_RPC_PROTO], includes)?; |
48 | | - let store_rpc_path = out_dir.join(STORE_RPC_DESCRIPTOR); |
49 | | - fs::write(&store_rpc_path, store_rpc_file_descriptor.encode_to_vec()) |
50 | | - .into_diagnostic() |
51 | | - .wrap_err("writing store rpc file descriptor")?; |
52 | | - |
53 | | - let store_ntx_builder_file_descriptor = protox::compile([STORE_NTX_BUILDER_PROTO], includes)?; |
54 | | - let store_ntx_builder_path = out_dir.join(STORE_NTX_BUILDER_DESCRIPTOR); |
55 | | - fs::write(&store_ntx_builder_path, store_ntx_builder_file_descriptor.encode_to_vec()) |
56 | | - .into_diagnostic() |
57 | | - .wrap_err("writing store ntx builder file descriptor")?; |
58 | | - |
59 | | - let store_block_producer_file_descriptor = |
60 | | - protox::compile([STORE_BLOCK_PRODUCER_PROTO], includes)?; |
61 | | - let store_block_producer_path = out_dir.join(STORE_BLOCK_PRODUCER_DESCRIPTOR); |
62 | | - fs::write(&store_block_producer_path, store_block_producer_file_descriptor.encode_to_vec()) |
63 | | - .into_diagnostic() |
64 | | - .wrap_err("writing store block producer file descriptor")?; |
65 | | - |
66 | | - let block_producer_file_descriptor = protox::compile([BLOCK_PRODUCER_PROTO], includes)?; |
67 | | - let block_producer_path = out_dir.join(BLOCK_PRODUCER_DESCRIPTOR); |
68 | | - fs::write(&block_producer_path, block_producer_file_descriptor.encode_to_vec()) |
69 | | - .into_diagnostic() |
70 | | - .wrap_err("writing block producer file descriptor")?; |
71 | | - |
72 | | - let validator_file_descriptor = protox::compile([VALIDATOR_PROTO], includes)?; |
73 | | - let validator_path = out_dir.join(VALIDATOR_DESCRIPTOR); |
74 | | - fs::write(&validator_path, validator_file_descriptor.encode_to_vec()) |
75 | | - .into_diagnostic() |
76 | | - .wrap_err("writing validator file descriptor")?; |
| 19 | + let schema_dir = build_rs::input::cargo_manifest_dir().join("proto"); |
| 20 | + |
| 21 | + // Codegen which will hold the file descriptor functions. |
| 22 | + // |
| 23 | + // `protox::prost::Message` is a trait which brings into scope the encoding and decoding of file |
| 24 | + // descriptors. This is required so because we serialize the descriptors in code as a `Vec<u8>` |
| 25 | + // and then decode it again inline. |
| 26 | + let mut code = codegen::Scope::new(); |
| 27 | + code.import("tonic_prost_build", "FileDescriptorSet"); |
| 28 | + code.import("protox::prost", "Message"); |
| 29 | + |
| 30 | + // We split our gRPC services into public and internal. |
| 31 | + // |
| 32 | + // This is easy to do since public services are listed in the root of the schema folder, |
| 33 | + // and internal services are nested in the `internal` folder. |
| 34 | + for public_api in proto_files_in_directory(&schema_dir)? { |
| 35 | + let file_descriptor_fn = generate_file_descriptor(&public_api, &schema_dir)?; |
| 36 | + code.push_fn(file_descriptor_fn); |
| 37 | + } |
| 38 | + |
| 39 | + // Internal gRPC services need an additional feature gate `#[cfg(feature = "internal")]`. |
| 40 | + for internal_api in proto_files_in_directory(&schema_dir.join("internal"))? { |
| 41 | + let mut file_descriptor_fn = generate_file_descriptor(&internal_api, &schema_dir)?; |
| 42 | + file_descriptor_fn.attr("cfg(feature = \"internal\")"); |
| 43 | + code.push_fn(file_descriptor_fn); |
| 44 | + } |
| 45 | + |
| 46 | + fs::write(out_dir.join("file_descriptors.rs"), code.to_string()).into_diagnostic()?; |
77 | 47 |
|
78 | 48 | Ok(()) |
79 | 49 | } |
| 50 | + |
| 51 | +/// The list of `*.proto` files in the given directory. |
| 52 | +/// |
| 53 | +/// Does _not_ recurse into folders; only top level files are returned. |
| 54 | +fn proto_files_in_directory(directory: &PathBuf) -> Result<Vec<PathBuf>, miette::Error> { |
| 55 | + let mut proto_files = Vec::new(); |
| 56 | + for entry in fs::read_dir(directory).into_diagnostic()? { |
| 57 | + let entry = entry.into_diagnostic()?; |
| 58 | + |
| 59 | + // Skip non-files |
| 60 | + if !entry.file_type().into_diagnostic()?.is_file() { |
| 61 | + continue; |
| 62 | + } |
| 63 | + |
| 64 | + // Skip non-protobuf files |
| 65 | + if PathBuf::from(entry.file_name()).extension().is_none_or(|ext| ext != "proto") { |
| 66 | + continue; |
| 67 | + } |
| 68 | + |
| 69 | + proto_files.push(entry.path()); |
| 70 | + } |
| 71 | + Ok(proto_files) |
| 72 | +} |
| 73 | + |
| 74 | +/// Creates a function which emits the file descriptor of the given gRPC service file. |
| 75 | +/// |
| 76 | +/// The function looks as follows: |
| 77 | +/// |
| 78 | +/// ```rust |
| 79 | +/// fn <file_stem>_api_descriptor() -> FileDescriptorSet { |
| 80 | +/// FileDescriptorSet::decode(vec![<encoded>].as_slice()) |
| 81 | +/// .expect("encoded file descriptor should decode") |
| 82 | +/// } |
| 83 | +/// ``` |
| 84 | +/// |
| 85 | +/// where `<encoded>` is bytes of the compiled gRPC service. |
| 86 | +fn generate_file_descriptor( |
| 87 | + grpc_service: &PathBuf, |
| 88 | + includes: &PathBuf, |
| 89 | +) -> Result<codegen::Function, miette::Error> { |
| 90 | + let file_name = grpc_service |
| 91 | + .file_stem() |
| 92 | + .and_then(OsStr::to_str) |
| 93 | + .ok_or_else(|| miette!("invalid file name for {grpc_service:?}"))?; |
| 94 | + |
| 95 | + let file_descriptor = protox::compile([grpc_service], includes)?; |
| 96 | + let file_descriptor = file_descriptor.encode_to_vec(); |
| 97 | + |
| 98 | + let mut f = codegen::Function::new(format!("{file_name}_api_descriptor")); |
| 99 | + f.vis("pub") |
| 100 | + .ret("FileDescriptorSet") |
| 101 | + .line(format!("FileDescriptorSet::decode(vec!{file_descriptor:?}.as_slice())")) |
| 102 | + .line(".expect(\"we just encoded this so it should decode\")"); |
| 103 | + |
| 104 | + Ok(f) |
| 105 | +} |
0 commit comments