#![deny(intra_doc_link_resolution_failure)]
#![deny(missing_docs)]
use protoc::Protoc;
use protoc_rust::Customize;
use std::fs;
use std::io;
use std::io::Read;
use std::io::Write;
use std::path::{Path, PathBuf};
pub type Error = io::Error;
pub type Result<T> = io::Result<T>;
#[derive(Debug, Default)]
pub struct Codegen {
protoc: Option<Protoc>,
out_dir: PathBuf,
includes: Vec<PathBuf>,
inputs: Vec<PathBuf>,
rust_protobuf: bool,
rust_protobuf_customize: protoc_rust::Customize,
}
impl Codegen {
pub fn new() -> Codegen {
Default::default()
}
pub fn out_dir(&mut self, out_dir: impl AsRef<Path>) -> &mut Self {
self.out_dir = out_dir.as_ref().to_owned();
self
}
pub fn include(&mut self, include: impl AsRef<Path>) -> &mut Self {
self.includes.push(include.as_ref().to_owned());
self
}
pub fn includes(&mut self, includes: impl IntoIterator<Item = impl AsRef<Path>>) -> &mut Self {
for include in includes {
self.include(include);
}
self
}
pub fn input(&mut self, input: impl AsRef<Path>) -> &mut Self {
self.inputs.push(input.as_ref().to_owned());
self
}
pub fn inputs(&mut self, inputs: impl IntoIterator<Item = impl AsRef<Path>>) -> &mut Self {
for input in inputs {
self.input(input);
}
self
}
pub fn rust_protobuf(&mut self, rust_protobuf: bool) -> &mut Self {
self.rust_protobuf = rust_protobuf;
self
}
pub fn rust_protobuf_customize(&mut self, rust_protobuf_customize: Customize) -> &mut Self {
self.rust_protobuf_customize = rust_protobuf_customize;
self
}
pub fn run(&self) -> Result<()> {
let protoc = self
.protoc
.clone()
.unwrap_or_else(|| protoc::Protoc::from_env_path());
let version = protoc.version().expect("protoc version");
if !version.is_3() {
panic!("protobuf must have version 3");
}
if self.rust_protobuf {
protoc_rust::Codegen::new()
.out_dir(&self.out_dir)
.includes(&self.includes)
.inputs(&self.inputs)
.customize(self.rust_protobuf_customize.clone())
.run()?;
}
let temp_dir = tempdir::TempDir::new("protoc-rust")?;
let temp_file = temp_dir.path().join("descriptor.pbbin");
let temp_file = temp_file.to_str().expect("utf-8 file name");
let includes: Vec<&str> = self
.includes
.iter()
.map(|p| p.as_os_str().to_str().unwrap())
.collect();
let inputs: Vec<&str> = self
.inputs
.iter()
.map(|p| p.as_os_str().to_str().unwrap())
.collect();
protoc.write_descriptor_set(protoc::DescriptorSetOutArgs {
out: temp_file,
includes: &includes,
input: &inputs,
include_imports: true,
})?;
let mut fds = Vec::new();
let mut file = fs::File::open(temp_file)?;
file.read_to_end(&mut fds)?;
drop(file);
drop(temp_dir);
let fds: protobuf::descriptor::FileDescriptorSet = protobuf::parse_from_bytes(&fds)
.map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
let mut includes = self.includes.clone();
if includes.is_empty() {
includes = vec![PathBuf::from(".")];
}
let mut files_to_generate = Vec::new();
'outer: for file in &self.inputs {
for include in &includes {
if let Some(truncated) =
remove_path_prefix(file.to_str().unwrap(), include.to_str().unwrap())
{
files_to_generate.push(truncated.to_owned());
continue 'outer;
}
}
return Err(Error::new(
io::ErrorKind::Other,
format!(
"file {:?} is not found in includes {:?}",
file, self.includes
),
));
}
let gen_result = grpc_compiler::codegen::gen(fds.get_file(), &files_to_generate);
for r in gen_result {
let r: protobuf::compiler_plugin::GenResult = r;
let file = format!("{}/{}", self.out_dir.display(), r.name);
let mut file = fs::File::create(&file)?;
file.write_all(&r.content)?;
file.flush()?;
}
Ok(())
}
}
fn remove_dot_slash(path: &str) -> &str {
if path == "." {
""
} else if path.starts_with("./") || path.starts_with(".\\") {
&path[2..]
} else {
path
}
}
fn remove_path_prefix<'a>(mut path: &'a str, mut prefix: &str) -> Option<&'a str> {
path = remove_dot_slash(path);
prefix = remove_dot_slash(prefix);
if prefix == "" {
return Some(path);
}
if prefix.ends_with("/") || prefix.ends_with("\\") {
prefix = &prefix[..prefix.len() - 1];
}
if !path.starts_with(prefix) {
return None;
}
if path.len() <= prefix.len() {
return None;
}
if path.as_bytes()[prefix.len()] == b'/' || path.as_bytes()[prefix.len()] == b'\\' {
return Some(&path[prefix.len() + 1..]);
} else {
return None;
}
}
#[cfg(test)]
mod test {
#[test]
fn remove_path_prefix() {
assert_eq!(
Some("abc.proto"),
super::remove_path_prefix("xxx/abc.proto", "xxx")
);
assert_eq!(
Some("abc.proto"),
super::remove_path_prefix("xxx/abc.proto", "xxx/")
);
assert_eq!(
Some("abc.proto"),
super::remove_path_prefix("../xxx/abc.proto", "../xxx/")
);
assert_eq!(
Some("abc.proto"),
super::remove_path_prefix("abc.proto", ".")
);
assert_eq!(
Some("abc.proto"),
super::remove_path_prefix("abc.proto", "./")
);
assert_eq!(None, super::remove_path_prefix("xxx/abc.proto", "yyy"));
assert_eq!(None, super::remove_path_prefix("xxx/abc.proto", "yyy/"));
}
}