1use std::any::Any;
2use std::collections::{BTreeMap, HashMap, HashSet};
3use std::convert::Into;
4use std::fs::canonicalize;
5use std::fs::File;
6use std::io::Read;
7use std::path::{Path, PathBuf};
8
9use anyhow::anyhow;
10use anyhow::{Error, Result};
11use askama::filters::format;
12use askama::Template;
13use bytes::Bytes;
14use clap::arg;
15use links::{Backlinks, ProtoSymbol};
16use log::{debug, info, warn};
17use mdbook::book::{Book, Chapter, SectionNumber};
18use mdbook::preprocess::{Preprocessor, PreprocessorContext};
19use mdbook::BookItem;
20use prost::Message;
21use prost_types::field_descriptor_proto::Type;
22use prost_types::source_code_info::Location;
23use prost_types::{
24 DescriptorProto, EnumDescriptorProto, FieldDescriptorProto, FileDescriptorProto,
25 FileDescriptorSet, ServiceDescriptorProto,
26};
27use toml_edit::Value;
28
29mod links;
30mod primitive;
31mod view;
32
33use links::SymbolLink;
34use view::{ProtoFileDescriptorTemplate, ProtoNamespaceTemplate};
35use crate::links::BASE_URL;
36
37pub fn read_file_descriptor_set(path: &Path) -> Result<FileDescriptorSet> {
38 info!("Attempting to read {}", path.display());
39
40 let mut file = File::open(path).map_err(|e| {
41 anyhow!(
42 "Could not read file at path `{}`, does it exist here?",
43 path.display()
44 )
45 })?;
46
47 info!("File descriptor set file found at {}", path.display());
48 let mut buffer = Vec::new();
49
50 file.read_to_end(&mut buffer)?;
51
52 let bytes = Bytes::from(buffer);
53
54 let decoded = FileDescriptorSet::decode(bytes)
55 .map_err(|e| anyhow!("failed to parse file descriptor set as protobuf"))?;
56
57 info!("Successfully decoded file descriptor set");
58 Ok(decoded)
59}
60
61const PREPROCESSOR_NAME: &'static str = "protobuf";
62
63pub struct ProtobufPreprocessor;
64
65impl ProtobufPreprocessor {
66 pub fn new() -> ProtobufPreprocessor {
67 ProtobufPreprocessor
68 }
69}
70
71pub struct ProtobufPreprocessorArgs {
72 nest_under: Option<String>,
73 file_descriptor_path: PathBuf,
74 proto_url_root: Option<String>,
75}
76
77impl ProtobufPreprocessorArgs {
78 pub fn new(ctx: &PreprocessorContext) -> Result<Self> {
79 let config = ctx
80 .config
81 .get_preprocessor(PREPROCESSOR_NAME)
82 .ok_or(anyhow!("Expected config"))?;
83
84 let file_descriptor_path = config
85 .get("proto_descriptor")
86 .ok_or(anyhow!("expected `proto_descriptor` key in config"))?;
87
88 let mut path = ctx.root.clone();
89 path.push(
90 file_descriptor_path
91 .as_str()
92 .ok_or(anyhow!("`proto_descriptor` should be a string"))?,
93 );
94
95 let file_descriptor_path = canonicalize(path.clone()).map_err(|e| {
96 anyhow!(
97 "Failed to find `proto_descriptor` at path {}",
98 path.display()
99 )
100 })?;
101
102 Ok(Self {
103 file_descriptor_path,
104 nest_under: config
105 .get("nest_under")
106 .and_then(|v| v.as_str().map(|s| s.to_string())),
107 proto_url_root: config
108 .get("proto_url_root")
109 .and_then(|v| v.as_str().map(|s| s.to_string())),
110 })
111 }
112}
113
114impl Preprocessor for ProtobufPreprocessor {
115 fn name(&self) -> &str {
116 PREPROCESSOR_NAME
117 }
118
119 fn run(&self, ctx: &PreprocessorContext, mut book: Book) -> Result<Book, Error> {
120 let args = ProtobufPreprocessorArgs::new(ctx)?;
121
122 let file_descriptor_set = read_file_descriptor_set(args.file_descriptor_path.as_path())?;
123
124 info!("found {} proto files", file_descriptor_set.file.len());
125
126 let mut namespaces: BTreeMap<String, ProtoNamespaceTemplate> = BTreeMap::new();
127
128 let mut symbol_usages: HashMap<SymbolLink, Vec<links::Backlink>> = HashMap::new();
129
130 let packages: HashSet<String> = file_descriptor_set
131 .file
132 .iter()
133 .map(|f| f.package().to_string())
134 .collect();
135
136 for file_descriptor in file_descriptor_set.file {
137 let value = namespaces
138 .entry(file_descriptor.package().to_string())
139 .or_default();
140
141 value.add_file(ProtoFileDescriptorTemplate::from_descriptor(
142 file_descriptor,
143 &packages,
144 &mut symbol_usages,
145 ));
146 }
147
148 let nest_under_path: Option<PathBuf> = if let Some(nest_under) = args.nest_under {
149 book.sections.iter().find_map(|s| match s {
150 BookItem::Chapter(c) => {
151 if c.name == nest_under {
152 c.path.clone()
153 } else {
154 None
155 }
156 }
157 _ => None,
158 })
159 } else {
160 None
161 };
162
163 let base_url = if let Some(site_url) = ctx.config.get("output.html.site-url").and_then(|u|u.as_str()) {
164 site_url.to_string()
165 } else {
166 "/".to_string()
167 };
168
169 if let Some(ref path) = nest_under_path {
170 let path = path.with_extension("");
171
172 let base_url = format!("{}{}", base_url, path.display());
173 info!("setting base url to {}", base_url);
174 BASE_URL.set(base_url).expect("Base url should not already be set");
175 } else {
176 BASE_URL.set(base_url).expect("Base url should not already be set");
177 }
178
179 for book_item in &mut book.sections {
180 if let BookItem::Chapter(chapter) = book_item {
181 links::link_proto_symbols(chapter, &mut symbol_usages)?;
182 }
183 }
184
185 links::assign_backlinks(&mut namespaces, symbol_usages);
186
187 if let Some(source_url) = args.proto_url_root {
188 info!("assigning source url to proto symbols: {}", &source_url);
189 links::assign_source_url(&mut namespaces, source_url);
190 } else {
191 warn!("proto_url_root was not set, so `[src]` links will not go to the correct destination");
192 }
193
194 let mut target_chapter = if let Some(nest_under) = &nest_under_path {
196 let found_section = book.sections.iter_mut().find_map(|s| match s {
197 BookItem::Chapter(c) => {
198 if c.path.as_ref() == Some(nest_under) {
199 Some(c)
200 } else {
201 None
202 }
203 }
204 _ => None,
205 });
206
207 if let None = found_section {
208 warn!("`nest_under` config was defined, but no chapter matching path `{}` was found. Note nested chapters are not yet supported.", nest_under.display());
209 }
210
211 found_section
212 } else {
213 None
214 };
215
216
217 let chapters: Result<Vec<Chapter>> = namespaces
218 .iter()
219 .map(|(namespace_key, namespace)| {
220 let content = namespace.render()?;
221 let path = PathBuf::from(format!("proto/{}", &namespace_key.replace(".", "/")));
222 Ok(Chapter::new(
223 namespace_key.as_ref(),
224 content,
225 path,
226 Vec::new(),
227 ))
228 })
229 .collect();
230
231 if let Some(target) = target_chapter {
232 for (idx, mut chapter) in chapters?.into_iter().enumerate() {
233 let mut section_number = target.clone().number.unwrap().0;
234 section_number.push((idx + 1) as u32);
235 chapter.number = Some(SectionNumber(section_number));
236 chapter.parent_names.extend(target.parent_names.clone());
237 chapter.parent_names.push(target.name.clone());
238
239 let section = BookItem::Chapter(chapter);
240
241 target.sub_items.push(section);
242 }
243 } else {
244 book.sections
245 .extend(chapters?.into_iter().map(BookItem::Chapter));
246 }
247
248 Ok(book)
249 }
250
251 fn supports_renderer(&self, renderer: &str) -> bool {
252 renderer != "not-supported"
253 }
254}
255
256#[cfg(test)]
257mod test {
258 use super::*;
259
260 #[test]
261 fn it_should_read_proto_descriptor() {
262 let path = Path::new("../demo/docs/build/proto_file_descriptor_set.pb");
263 let descriptor = read_file_descriptor_set(path);
264
265 assert!(descriptor.is_ok());
266 dbg!(&descriptor);
267 }
268
269 #[test]
270 fn preprocessor_run() {
271 let input_json = r##"[
272 {
273 "root": "./",
274 "config": {
275 "book": {
276 "authors": ["AUTHOR"],
277 "language": "en",
278 "multilingual": false,
279 "src": "src",
280 "title": "TITLE"
281 },
282 "preprocessor": {
283 "protobuf": {
284 "proto_descriptor": "../demo/docs/build/proto_file_descriptor_set.pb",
285 "proto_url_root": "http://example.com/proto/",
286 "nest_under": "Chapter 1"
287 }
288 },
289 "output": {
290 "html": {
291 "site-url": "/sdk/"
292 }
293 }
294 },
295 "renderer": "html",
296 "mdbook_version": "0.4.21"
297 },
298 {
299 "sections": [
300 {
301 "Chapter": {
302 "name": "Chapter 1",
303 "content": "# Chapter 1\n [Message](proto!(Message)) [](proto!(MessageEmpty))",
304 "number": [1],
305 "sub_items": [],
306 "path": "chapter_1.md",
307 "source_path": "chapter_1.md",
308 "parent_names": []
309 }
310 }
311 ],
312 "__non_exhaustive": null
313 }
314 ]"##;
315 let input_json = input_json.as_bytes();
316
317 let (ctx, book) = mdbook::preprocess::CmdPreprocessor::parse_input(input_json).unwrap();
318 let expected_book = book.clone();
319 let result = ProtobufPreprocessor::new().run(&ctx, book);
320 assert!(result.is_ok());
321 }
322
323}