protoc_plugin_by_closure/
lib.rs

1// Copyright 2021 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//      http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15#![doc = include_str!("../readme.md")]
16
17use ipc_channel::ipc::{IpcBytesReceiver, IpcBytesSender, IpcOneShotServer};
18use std::env;
19use std::path::PathBuf;
20use std::process::{Command, ExitStatus};
21use std::time::Duration;
22#[cfg(feature = "on-memory")]
23use tempfile::TempDir;
24use thiserror::Error;
25use wait_timeout::ChildExt;
26
27const PLUGIN_PATH: &'static str = env!("CARGO_BIN_FILE_PROTOC_PLUGIN_BIN");
28
29/// Error type for this crate.
30#[derive(Error, Debug)]
31pub enum ErrorKind {
32    #[error("IpcIpcError: {0}")]
33    IpcIpcError(#[from] ::ipc_channel::ipc::IpcError),
34    #[error("IpcError: {0}")]
35    IpcError(#[from] ::ipc_channel::Error),
36    #[error("IoError: {0}")]
37    IoError(#[from] ::std::io::Error),
38    #[error("CallbackError: {0}")]
39    CallbackError(String),
40    #[error("ProtocTimeoutError")]
41    ProtocTimeoutError,
42    #[error("ProtocProcessError: {0}")]
43    ProtocProcessError(ExitStatus),
44    #[error("FileNameError")]
45    FileNameError,
46}
47
48/// Result type for this crate.
49pub type Result<T> = ::std::result::Result<T, ErrorKind>;
50
51/// A convenient wrapper for running protoc command with your own plugin code as a closure.
52///
53/// See the [crate level documentation](crate) for the basic explanation.
54///
55/// # Example
56/// ```no_run
57/// # fn run_protoc() {
58///     use protoc_plugin_by_closure::Protoc;
59///     use std::time::Duration;
60///     Protoc::new()
61///         .proto_file("my_protobuf_file.proto")
62///         .proto_file("my_protobuf_file2.proto")
63///         .proto_path("path/to/my/input_proto_dir/")
64///         .out_dir("path/to/my/output_dir/")
65///         .run(Duration::from_sec(3), |request_bytes| {
66///             // Your plugin logic here, which takes the CodeGeneratorRequest bytes
67///             // and returns the Result of CodeGeneratorResponse bytes.
68/// #           unimplemented!()
69///         })
70///         .unwrap();
71///
72///     // The generated file names depend on your plugin logic and the contents of
73///     // the input proto files, but typically they will be like this:
74///     assert!(std::path::Path("path/to/my/output_dir/my_protobuf_file.rs").exists());
75///     assert!(std::path::Path("path/to/my/output_dir/my_protobuf_file2.rs").exists());
76/// # }
77/// ```
78pub struct Protoc {
79    protoc_path: PathBuf,
80    out_dir: Option<PathBuf>,
81    proto_files: Vec<PathBuf>,
82    proto_paths: Vec<PathBuf>,
83}
84
85impl Protoc {
86    /// Creates a new `Protoc` instance.
87    pub fn new() -> Self {
88        Self {
89            protoc_path: "protoc".into(),
90            out_dir: None,
91            proto_files: Vec::new(),
92            proto_paths: Vec::new(),
93        }
94    }
95    /// Sets the path to the `protoc` command. Default is `"protoc"`.
96    pub fn protoc_path(mut self, path: impl Into<PathBuf>) -> Self {
97        self.protoc_path = path.into();
98        self
99    }
100    /// Sets the output directory for the generated files. Corresponds to `--rust_out` option of `protoc`.
101    pub fn out_dir(mut self, path: impl Into<PathBuf>) -> Self {
102        self.out_dir = Some(path.into());
103        self
104    }
105    /// Sets the path to the input proto file. Corresponds to the unnamed argument of `protoc`.
106    pub fn proto_file(mut self, path: impl Into<PathBuf>) -> Self {
107        self.proto_files.push(path.into());
108        self
109    }
110    /// Sets the paths to the input proto files. Corresponds to the unnamed arguments of `protoc`.
111    pub fn proto_files<I>(mut self, paths: I) -> Self
112    where
113        I: IntoIterator,
114        I::Item: Into<PathBuf>,
115    {
116        self.proto_files.extend(paths.into_iter().map(|p| p.into()));
117        self
118    }
119    /// Sets the path to the input proto file directory. Corresponds to `--proto_path` option of `protoc`.
120    pub fn proto_path(mut self, path: impl Into<PathBuf>) -> Self {
121        self.proto_paths.push(path.into());
122        self
123    }
124
125    /// Runs the `protoc` command with the given closure as a plugin code.
126    ///
127    /// The `body` param can be any `FnOnce` closure which takes the encoded `CodeGeneratorRequest` bytes
128    /// and returns the `Result` of encoded `CodeGeneratorResponse` bytes.
129    ///
130    /// Set the `timeout` to the maximum duration of the `protoc` command execution.
131    pub fn run<F>(self, timeout: Duration, body: F) -> Result<()>
132    where
133        F: FnOnce(&[u8]) -> ::std::result::Result<Vec<u8>, String>,
134    {
135        let (ipc_init_server, ipc_init_name) = IpcOneShotServer::new()?;
136
137        let mut process = Command::new(&self.protoc_path)
138            .args(&[
139                // We name our plugin binary name as "rust-ppbc" here.
140                format!("--plugin=protoc-gen-rust-ppbc={}", PLUGIN_PATH),
141                format!(
142                    "--rust-ppbc_out={}",
143                    self.out_dir
144                        .as_ref()
145                        .map(|p| p.to_str().ok_or(ErrorKind::FileNameError))
146                        .transpose()?
147                        .unwrap_or(".")
148                ),
149                format!("--rust-ppbc_opt={}", ipc_init_name),
150            ])
151            .args(
152                self.proto_paths
153                    .iter()
154                    .map(|x| {
155                        Ok(format!(
156                            "--proto_path={}",
157                            x.to_str().ok_or(ErrorKind::FileNameError)?
158                        ))
159                    })
160                    .collect::<Result<Vec<_>>>()?,
161            )
162            .args(&self.proto_files)
163            .spawn()?;
164
165        {
166            // recieve the ipc channels from the plugin exe.
167            let (req_recv, res_send): (IpcBytesReceiver, IpcBytesSender) =
168                ipc_init_server.accept()?.1;
169
170            let req = req_recv.recv()?;
171            let res = (body)(&req).map_err(|x| ErrorKind::CallbackError(x))?;
172
173            res_send.send(&res)?;
174        }
175
176        let Some(exit_code) = process.wait_timeout(timeout)? else {
177            return Err(ErrorKind::ProtocTimeoutError);
178        };
179        if !exit_code.success() {
180            return Err(ErrorKind::ProtocProcessError(exit_code));
181        }
182
183        Ok(())
184    }
185}
186
187/// A variant of [`Protoc`] which you can run the `protoc` command without touching the actual filesystem.
188///
189/// Instead of using the actual filesystem, you can pass the name-value pairs of
190/// proto files to this struct, and it returns the generated files as name-value pairs.
191///
192/// This is useful when you want to test your plugin code, or when you want to implement your
193/// procedual macro which generates the inlined generated code.
194///
195/// See the [crate level documentation](crate) or [`Protoc`] for the basic explanations.
196///
197/// # Example
198/// ```no_run
199/// # fn run_protoc() {
200///     use protoc_plugin_by_closure::ProtocOnMemory;
201///     use std::time::Duration;
202///     let result_files = Protoc::new()
203///         .add_file("my_protobuf_file.proto", r#"
204/// syntax = "proto3";
205/// package my_package;
206/// message MyMessage {
207///   string name = 1;
208/// }"#)
209///         .add_file("another/path/to/my_protobuf_file2.proto", r#"
210/// syntax = "proto3";
211/// package my_package2;
212/// message MyMessage2 {
213///   string name2 = 2;
214/// }"#)
215///         .run(Duration::from_sec(3), |request_bytes| {
216///             // Your plugin logic here, which takes the CodeGeneratorRequest bytes
217///             // and returns the Result of CodeGeneratorResponse bytes.
218/// #           unimplemented!()
219///         })
220///         .unwrap();
221///
222///     // The generated filenames depend on your plugin logic, but typically they will be like this:
223///     assert!(result_files.iter().any(|(name, _)| name == "my_package.rs"));
224///     assert!(result_files.iter().any(|(name, _)| name == "my_package2.rs"));
225/// # }
226/// ```
227#[cfg(feature = "on-memory")]
228pub struct ProtocOnMemory {
229    protoc: Protoc,
230    in_files: Vec<(String, String)>,
231}
232
233impl ProtocOnMemory {
234    /// Creates a new `ProtocOnMemory` instance.
235    pub fn new() -> Self {
236        Self {
237            protoc: Protoc::new(),
238            in_files: Vec::new(),
239        }
240    }
241    /// Sets the path to the `protoc` command. Default is `"protoc"`.
242    pub fn protoc_path(mut self, path: impl Into<PathBuf>) -> Self {
243        self.protoc = self.protoc.protoc_path(path);
244        self
245    }
246    /// Adds a (virtual) input proto file. Corresponds to the `protoc` command's unnamed argument.
247    pub fn add_file(mut self, name: &str, content: &str) -> Self {
248        self.in_files.push((name.to_string(), content.to_string()));
249        self
250    }
251    /// Adds (virtual) input proto files. Corresponds to the `protoc` command's unnamed arguments.
252    pub fn add_files<I>(mut self, files: I) -> Self
253    where
254        I: IntoIterator<Item = (String, String)>,
255    {
256        self.in_files.extend(files);
257        self
258    }
259
260    /// Runs the `protoc` command with the given closure as a plugin code.
261    ///
262    /// The `body` param can be any `FnOnce` closure which takes the encoded `CodeGeneratorRequest` bytes
263    /// and returns the `Result` of encoded `CodeGeneratorResponse` bytes.
264    ///
265    /// Set the `timeout` to the maximum duration of the `protoc` command execution.
266    pub fn run<F>(self, timeout: Duration, func: F) -> Result<Vec<(String, String)>>
267    where
268        F: FnOnce(&[u8]) -> ::std::result::Result<Vec<u8>, String>,
269    {
270        let proto_dir = TempDir::new()?;
271        let out_dir = TempDir::new()?;
272        let out_dir_path = out_dir.path().to_str().unwrap().to_string();
273
274        // write the proto files to the temp dir.
275        for (name, content) in &self.in_files {
276            let path = proto_dir.path().join(&name);
277            ::std::fs::write(&path, &content)?;
278        }
279
280        // run the protoc
281        let proto_file_paths = self
282            .in_files
283            .iter()
284            .map(|(name, _)| {
285                proto_dir
286                    .path()
287                    .join(name)
288                    .to_str()
289                    .ok_or(ErrorKind::FileNameError)
290                    .map(str::to_string)
291            })
292            .collect::<Result<Vec<_>>>()?;
293
294        self.protoc
295            .out_dir(&out_dir_path)
296            .proto_path(proto_dir.path().to_str().unwrap())
297            .proto_files(proto_file_paths)
298            .run(timeout, func)?;
299
300        // read the generated files
301        let output_files = ::std::fs::read_dir(out_dir.path())?
302            .map(|entry| -> Result<_> {
303                let path = entry?.path();
304                let name = path
305                    .strip_prefix(out_dir.path())
306                    .map_err(|_| ErrorKind::FileNameError)?
307                    .to_str()
308                    .ok_or(ErrorKind::FileNameError)?
309                    .to_string();
310                let content = ::std::fs::read_to_string(&path)?;
311                Ok((name, content))
312            })
313            .collect::<Result<Vec<_>>>()?;
314
315        Ok(output_files)
316    }
317}
protoc_plugin_by_closure/lib.rs

protoc_plugin_by_closure/
lib.rs
