protoc_grpc/
lib.rs

1extern crate grpcio_compiler;
2extern crate protobuf;
3extern crate protoc;
4extern crate protoc_rust;
5extern crate tempdir;
6
7use std::fs;
8use std::io;
9use std::io::Read;
10use std::io::Write;
11
12pub type Error = io::Error;
13pub type Result<T> = io::Result<T>;
14
15#[derive(Debug, Default)]
16pub struct Args<'a> {
17    /// --lang_out= param
18    pub out_dir: &'a str,
19    /// -I args
20    pub includes: &'a [&'a str],
21    /// List of .proto files to compile
22    pub input: &'a [&'a str],
23    /// Generate rust-protobuf files along with rust-gprc
24    pub rust_protobuf: bool,
25}
26
27pub fn run(args: Args) -> Result<()> {
28    let protoc = protoc::Protoc::from_env_path();
29    let version = protoc.version().expect("protoc version");
30    if !version.is_3() {
31        panic!("protobuf must have version 3");
32    }
33
34    if args.rust_protobuf {
35        protoc_rust::run(protoc_rust::Args {
36            out_dir: args.out_dir,
37            includes: args.includes,
38            input: args.input,
39            ..Default::default()
40        })?;
41    }
42
43    let temp_dir = tempdir::TempDir::new("protoc-rust")?;
44    let temp_file = temp_dir.path().join("descriptor.pbbin");
45    let temp_file = temp_file.to_str().expect("utf-8 file name");
46
47    protoc.write_descriptor_set(protoc::DescriptorSetOutArgs {
48        out: temp_file,
49        includes: args.includes,
50        input: args.input,
51        include_imports: true,
52    })?;
53
54    let mut fds = Vec::new();
55    let mut file = fs::File::open(temp_file)?;
56    file.read_to_end(&mut fds)?;
57
58    drop(file);
59    drop(temp_dir);
60
61    let fds: protobuf::descriptor::FileDescriptorSet =
62        protobuf::parse_from_bytes(&fds).map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
63
64    let mut includes = args.includes;
65    if includes.is_empty() {
66        static DOT_SLICE: &'static [&'static str] = &["."];
67        includes = DOT_SLICE;
68    }
69
70    let mut files_to_generate = Vec::new();
71    'outer: for file in args.input {
72        for include in includes {
73            if let Some(truncated) = remove_path_prefix(file, include) {
74                files_to_generate.push(truncated.to_owned());
75                continue 'outer;
76            }
77        }
78
79        return Err(Error::new(
80            io::ErrorKind::Other,
81            format!(
82                "file {:?} is not found in includes {:?}",
83                file, args.includes
84            ),
85        ));
86    }
87
88    let gen_result = grpcio_compiler::codegen::gen(fds.get_file(), &files_to_generate);
89
90    for r in gen_result {
91        let r: protobuf::compiler_plugin::GenResult = r;
92        let file = format!("{}/{}", args.out_dir, r.name);
93        let mut file = fs::File::create(&file)?;
94        file.write_all(&r.content)?;
95        file.flush()?;
96    }
97
98    Ok(())
99}
100
101fn remove_dot_slash(path: &str) -> &str {
102    if path == "." {
103        ""
104    } else if path.starts_with("./") || path.starts_with(".\\") {
105        &path[2..]
106    } else {
107        path
108    }
109}
110
111fn remove_path_prefix<'a>(mut path: &'a str, mut prefix: &str) -> Option<&'a str> {
112    path = remove_dot_slash(path);
113    prefix = remove_dot_slash(prefix);
114
115    if prefix == "" {
116        return Some(path);
117    }
118
119    if prefix.ends_with("/") || prefix.ends_with("\\") {
120        prefix = &prefix[..prefix.len() - 1];
121    }
122
123    if !path.starts_with(prefix) {
124        return None;
125    }
126
127    if path.len() <= prefix.len() {
128        return None;
129    }
130
131    if path.as_bytes()[prefix.len()] == b'/' || path.as_bytes()[prefix.len()] == b'\\' {
132        return Some(&path[prefix.len() + 1..]);
133    } else {
134        return None;
135    }
136}
137
138#[cfg(test)]
139mod test {
140    #[test]
141    fn remove_path_prefix() {
142        assert_eq!(
143            Some("abc.proto"),
144            super::remove_path_prefix("xxx/abc.proto", "xxx")
145        );
146        assert_eq!(
147            Some("abc.proto"),
148            super::remove_path_prefix("xxx/abc.proto", "xxx/")
149        );
150        assert_eq!(
151            Some("abc.proto"),
152            super::remove_path_prefix("../xxx/abc.proto", "../xxx/")
153        );
154        assert_eq!(
155            Some("abc.proto"),
156            super::remove_path_prefix("abc.proto", ".")
157        );
158        assert_eq!(
159            Some("abc.proto"),
160            super::remove_path_prefix("abc.proto", "./")
161        );
162        assert_eq!(None, super::remove_path_prefix("xxx/abc.proto", "yyy"));
163        assert_eq!(None, super::remove_path_prefix("xxx/abc.proto", "yyy/"));
164    }
165}
protoc_grpc/lib.rs

protoc_grpc/
lib.rs
