mdbook_aquascope/
preprocessor.rs

1//! Parser for Aquascope code blocks within Markdown.
2
3use std::{
4  collections::HashMap,
5  fs,
6  path::PathBuf,
7  process::{Command, Stdio},
8  sync::RwLock,
9  time::Duration,
10};
11
12use anyhow::{bail, Result};
13use aquascope_workspace_utils::{miri_sysroot, run_and_get_output, rustc};
14use mdbook_preprocessor_utils::HtmlElementBuilder;
15use rayon::prelude::*;
16use tempfile::tempdir;
17use wait_timeout::ChildExt;
18
19use crate::{block::AquascopeBlock, cache::Cache};
20
21pub struct AquascopePreprocessor {
22  miri_sysroot: PathBuf,
23  target_libdir: PathBuf,
24  cache: RwLock<Cache<AquascopeBlock, String>>,
25}
26
27impl AquascopePreprocessor {
28  pub fn new() -> Result<Self> {
29    let miri_sysroot = miri_sysroot()?;
30    let rustc = rustc()?;
31    let output = run_and_get_output(
32      Command::new(rustc).args(["--print", "target-libdir"]),
33    )?;
34    let target_libdir = PathBuf::from(output);
35
36    let cache = RwLock::new(Cache::load()?);
37    Ok(AquascopePreprocessor {
38      miri_sysroot,
39      target_libdir,
40      cache,
41    })
42  }
43
44  /// Runs cargo-aquascope on code from a given Aquascope block.
45  fn run_aquascope(&self, block: &AquascopeBlock) -> Result<String> {
46    // TODO: this code shares a lot of structure w/ aquascope_serve.
47    // Can we unify them?
48    let tempdir = tempdir()?;
49    let root = tempdir.path();
50    let status = Command::new("cargo")
51      .args(["new", "--bin", "example"])
52      .current_dir(root)
53      .stdout(Stdio::null())
54      .stderr(Stdio::null())
55      .status()?;
56    if !status.success() {
57      bail!("Cargo failed");
58    }
59
60    fs::write(root.join("example/src/main.rs"), &block.code)?;
61
62    let mut responses = HashMap::new();
63    for operation in &block.operations {
64      let mut cmd = Command::new("cargo");
65      cmd
66        .arg("aquascope")
67        .env("SYSROOT", &self.miri_sysroot)
68        .env("MIRI_SYSROOT", &self.miri_sysroot)
69        .env("DYLD_LIBRARY_PATH", &self.target_libdir)
70        .env("LD_LIBRARY_PATH", &self.target_libdir)
71        .env("RUST_BACKTRACE", "1")
72        .current_dir(root.join("example"));
73
74      let should_fail = block.config.iter().any(|(k, _)| k == "shouldFail");
75      if should_fail {
76        cmd.arg("--should-fail");
77      }
78
79      cmd.arg(operation);
80
81      let show_flows = block.config.iter().any(|(k, _)| k == "showFlows");
82      if show_flows {
83        cmd.arg("--show-flows");
84      }
85
86      let mut child =
87        cmd.stdout(Stdio::piped()).stderr(Stdio::piped()).spawn()?;
88      if child.wait_timeout(Duration::from_secs(10))?.is_none() {
89        child.kill()?;
90        bail!("Aquascope timed out on program:\n{}", block.code)
91      };
92
93      let output = child.wait_with_output()?;
94
95      if !output.status.success() {
96        let error = String::from_utf8(output.stderr)?;
97        bail!(
98          "Aquascope failed for program:\n{}\nwith error:\n{error}",
99          block.code
100        )
101      }
102
103      let response = String::from_utf8(output.stdout)?;
104      let response_json: serde_json::Value = serde_json::from_str(&response)?;
105      let is_err = match (response_json.as_object(), response_json.as_array()) {
106        (Some(obj), _) => obj.get("Err").is_some(),
107        (_, Some(arr)) => arr.iter().any(|obj| obj.get("Err").is_some()),
108        _ => false,
109      };
110      if is_err {
111        let stderr = String::from_utf8(output.stderr)?;
112        bail!(
113          "Aquascope failed for program:\n{}\nwith error:\n{stderr}",
114          block.code,
115        )
116      }
117
118      if let Some("BuildError") =
119        response_json.get("type").and_then(|ty| ty.as_str())
120      {
121        bail!("Aquascope failed for program:\n{}", block.code)
122      }
123
124      responses.insert(operation, response_json);
125    }
126
127    Ok(serde_json::to_string(&responses)?)
128  }
129
130  /// Get the HTML output for an Aquascope block
131  fn process_code(&self, block: AquascopeBlock) -> Result<String> {
132    let cached_response = {
133      let cache = self.cache.read().unwrap();
134      cache.get(&block).cloned()
135    };
136    let response_str = match cached_response {
137      Some(response) => response,
138      None => {
139        let response = self.run_aquascope(&block)?;
140        self
141          .cache
142          .write()
143          .unwrap()
144          .set(block.clone(), response.clone());
145        response
146      }
147    };
148    let response: serde_json::Value =
149      serde_json::from_str(response_str.trim_end())?;
150
151    let mut html = HtmlElementBuilder::new();
152    html
153      .attr("class", "aquascope-embed")
154      .data("code", &block.code)?
155      .data("annotations", &block.annotations)?
156      .data("operations", &block.operations)?
157      .data("responses", response)?;
158
159    let config = block
160      .config
161      .iter()
162      .map(|(k, v)| (k, v))
163      .collect::<HashMap<_, _>>();
164    html.data("config", &config)?;
165
166    // TODO: make this configurable
167    html.data("no-interact", true)?;
168
169    // TODO: add a code path to enable this from config
170    // add_data("show-bug-reporter", true)?;
171
172    Ok(html.finish())
173  }
174
175  pub fn replacements(
176    &self,
177    content: &str,
178  ) -> Result<Vec<(std::ops::Range<usize>, String)>> {
179    let to_process = AquascopeBlock::parse_all(content);
180    to_process
181      .into_par_iter()
182      .map(|(range, block)| {
183        let html = self.process_code(block)?;
184        Ok((range, html))
185      })
186      .chain(crate::permissions::parse_perms(content).par_bridge())
187      .collect()
188  }
189
190  pub fn save_cache(&mut self) {
191    self.cache.write().unwrap().save().unwrap();
192  }
193}
mdbook_aquascope/preprocessor.rs

mdbook_aquascope/
preprocessor.rs
