1use std::ffi::OsString;
2use std::path::{Path, PathBuf};
3use std::process::{ExitStatus, Stdio};
4use std::time::Duration;
5
6use kellnr_docs::docs_error::DocsError;
7use kellnr_settings::cli::{DOC_BUILD_CRATESIO_INDEX_ARG, DOC_BUILD_SUBCOMMAND};
8use tokio::io::{AsyncRead, AsyncReadExt};
9use tokio::process::{Child, Command};
10use tokio_util::sync::CancellationToken;
11use tracing::warn;
12
13const ALLOWED_ENV_VARS: &[&str] = &[
16 "PATH",
17 "HOME",
18 "USER",
19 "LANG",
20 "LC_ALL",
21 "TMPDIR",
22 "TEMP",
23 "TMP",
24 "RUSTUP_HOME",
25 "RUSTUP_TOOLCHAIN",
26 "HTTP_PROXY",
27 "HTTPS_PROXY",
28 "ALL_PROXY",
29 "NO_PROXY",
30 "http_proxy",
31 "https_proxy",
32 "all_proxy",
33 "no_proxy",
34 "SSL_CERT_FILE",
35 "SSL_CERT_DIR",
36 "SYSTEMROOT",
38 "USERPROFILE",
39 "APPDATA",
40 "LOCALAPPDATA",
41 "PATHEXT",
42 "COMSPEC",
43];
44
45const CARGO_ENV_PREFIX: &str = "CARGO_";
51
52const OUTPUT_TAIL_BYTES: usize = 16 * 1024;
54const OUTPUT_GRACE_PERIOD: Duration = Duration::from_secs(5);
57
58#[derive(Debug, Clone)]
65pub struct DocBuilder {
66 program: PathBuf,
67 wrapper: Vec<String>,
68 cratesio_index: Option<String>,
69 timeout: Option<Duration>,
70 shutdown: CancellationToken,
71}
72
73enum Waited {
75 Exited(std::io::Result<ExitStatus>),
76 TimedOut(Duration),
77 ShutDown,
78}
79
80impl DocBuilder {
81 pub fn new(program: PathBuf, wrapper: Vec<String>, cratesio_index: Option<String>) -> Self {
82 Self {
83 program,
84 wrapper,
85 cratesio_index,
86 timeout: None,
87 shutdown: CancellationToken::new(),
88 }
89 }
90
91 #[must_use]
94 pub fn with_shutdown(mut self, shutdown: CancellationToken) -> Self {
95 self.shutdown = shutdown;
96 self
97 }
98
99 #[must_use]
102 pub fn with_timeout(mut self, timeout: Duration) -> Self {
103 self.timeout = (!timeout.is_zero()).then_some(timeout);
104 self
105 }
106
107 pub async fn build(&self, crate_path: &Path) -> Result<(), DocsError> {
109 let mut child = self.command(crate_path).spawn().map_err(|e| {
110 DocsError::CargoError(format!("failed to start docs build process: {e}"))
111 })?;
112 let output = tokio::spawn(read_tail(
113 child.stderr.take().expect("stderr is piped"),
114 OUTPUT_TAIL_BYTES,
115 ));
116
117 let status = match self.wait(&mut child).await {
118 Waited::Exited(status) => status?,
119 Waited::TimedOut(timeout) => {
120 kill_build(&mut child).await;
121 output.abort();
122 return Err(DocsError::BuildTimeout(timeout.as_secs()));
123 }
124 Waited::ShutDown => {
125 kill_build(&mut child).await;
126 output.abort();
127 return Err(DocsError::BuildStopped);
128 }
129 };
130 let output = match tokio::time::timeout(OUTPUT_GRACE_PERIOD, output).await {
131 Ok(output) => output.unwrap_or_default(),
132 Err(_) => String::from("<build output not available>"),
133 };
134 if status.success() {
135 Ok(())
136 } else {
137 Err(DocsError::CargoError(format!(
138 "docs build failed ({status}): {output}"
139 )))
140 }
141 }
142
143 async fn wait(&self, child: &mut Child) -> Waited {
144 let timed_out = async {
145 match self.timeout {
146 Some(timeout) => tokio::time::sleep(timeout).await,
147 None => std::future::pending().await,
148 }
149 };
150 tokio::select! {
151 status = child.wait() => Waited::Exited(status),
152 () = timed_out => Waited::TimedOut(self.timeout.unwrap_or_default()),
153 () = self.shutdown.cancelled() => Waited::ShutDown,
154 }
155 }
156
157 fn command(&self, crate_path: &Path) -> Command {
158 let mut args: Vec<OsString> = self.wrapper.iter().map(OsString::from).collect();
159 args.push(self.program.clone().into_os_string());
160 args.push(DOC_BUILD_SUBCOMMAND.into());
161 args.push(crate_path.into());
162 if let Some(index) = &self.cratesio_index {
163 args.push(format!("--{DOC_BUILD_CRATESIO_INDEX_ARG}").into());
164 args.push(index.into());
165 }
166
167 let mut command = Command::new(&args[0]);
168 command
169 .args(&args[1..])
170 .env_clear()
171 .envs(allowed_env(std::env::vars_os()))
172 .stdin(Stdio::null())
173 .stdout(Stdio::null())
174 .stderr(Stdio::piped())
175 .kill_on_drop(true);
176 #[cfg(unix)]
179 command.process_group(0);
180 command
181 }
182}
183
184async fn kill_build(child: &mut Child) {
186 #[cfg(unix)]
187 if let Some(pid) = child
188 .id()
189 .and_then(|id| i32::try_from(id).ok())
190 .and_then(rustix::process::Pid::from_raw)
191 && let Err(e) = rustix::process::kill_process_group(pid, rustix::process::Signal::KILL)
192 {
193 warn!("Failed to kill docs build process group: {e}");
194 }
195 if let Err(e) = child.kill().await {
196 warn!("Failed to kill docs build process: {e}");
197 }
198}
199
200fn allowed_env(
201 vars: impl Iterator<Item = (OsString, OsString)>,
202) -> impl Iterator<Item = (OsString, OsString)> {
203 vars.filter(|(key, _)| key.to_str().is_some_and(is_allowed_env_var))
204}
205
206fn is_allowed_env_var(key: &str) -> bool {
207 ALLOWED_ENV_VARS.contains(&key) || key.starts_with(CARGO_ENV_PREFIX)
208}
209
210async fn read_tail(mut reader: impl AsyncRead + Unpin, max_bytes: usize) -> String {
212 let mut tail = Vec::new();
213 let mut buf = [0u8; 8192];
214 while let Ok(n) = reader.read(&mut buf).await {
215 if n == 0 {
216 break;
217 }
218 tail.extend_from_slice(&buf[..n]);
219 if tail.len() > max_bytes {
220 tail.drain(..tail.len() - max_bytes);
221 }
222 }
223 String::from_utf8_lossy(&tail).into_owned()
224}
225
226#[cfg(test)]
227mod tests {
228 use super::*;
229
230 fn vars(pairs: &[(&str, &str)]) -> Vec<(OsString, OsString)> {
231 pairs
232 .iter()
233 .map(|(k, v)| (OsString::from(k), OsString::from(v)))
234 .collect()
235 }
236
237 #[test]
238 fn allowed_env_drops_kellnr_settings_and_unknown_vars() {
239 let env = vars(&[
240 ("PATH", "/usr/bin"),
241 ("CARGO_HOME", "/cargo"),
242 ("RUSTUP_TOOLCHAIN", "stable"),
243 ("CARGO_REGISTRIES_KELLNR_TOKEN", "doc-token"),
244 ("KELLNR_POSTGRESQL__PWD", "secret"),
245 ("KELLNR_SETUP__ADMIN_TOKEN", "token"),
246 ("AWS_SECRET_ACCESS_KEY", "aws"),
247 ]);
248
249 let allowed: Vec<_> = allowed_env(env.into_iter()).collect();
250
251 assert_eq!(
252 allowed,
253 vars(&[
254 ("PATH", "/usr/bin"),
255 ("CARGO_HOME", "/cargo"),
256 ("RUSTUP_TOOLCHAIN", "stable"),
257 ("CARGO_REGISTRIES_KELLNR_TOKEN", "doc-token"),
258 ])
259 );
260 }
261
262 #[tokio::test]
263 async fn read_tail_keeps_only_the_end_of_the_output() {
264 let output = read_tail(&b"0123456789"[..], 4).await;
265
266 assert_eq!(output, "6789");
267 }
268
269 #[cfg(unix)]
272 fn recording_wrapper(out: &Path) -> Vec<String> {
273 vec![
274 "/bin/sh".to_string(),
275 "-c".to_string(),
276 "env > \"$0\"; echo \"args: $*\" >> \"$0\"".to_string(),
277 out.display().to_string(),
278 ]
279 }
280
281 #[cfg(unix)]
282 #[tokio::test]
283 async fn build_runs_wrapper_with_kellnr_binary_and_cleared_env() {
284 let dir = tempfile::tempdir().unwrap();
285 let out = dir.path().join("out");
286 let builder = DocBuilder::new(
287 PathBuf::from("/opt/kellnr/kellnr"),
288 recording_wrapper(&out),
289 Some("https://proxy/index/".to_string()),
290 );
291
292 builder.build(Path::new("/tmp/my-crate")).await.unwrap();
293
294 let recorded = std::fs::read_to_string(&out).unwrap();
295 assert!(recorded.contains(
296 "args: /opt/kellnr/kellnr doc-build /tmp/my-crate --cratesio-index https://proxy/index/"
297 ));
298 for line in recorded.lines().filter(|l| !l.starts_with("args: ")) {
299 let key = line.split('=').next().unwrap();
300 assert!(
302 is_allowed_env_var(key) || ["PWD", "SHLVL", "_", "OLDPWD"].contains(&key),
303 "unexpected variable {key} passed to the docs build"
304 );
305 }
306 }
307
308 #[cfg(unix)]
309 #[tokio::test]
310 async fn build_fails_with_output_of_failed_process() {
311 let builder = DocBuilder::new(
312 PathBuf::from("unused"),
313 vec![
314 "/bin/sh".to_string(),
315 "-c".to_string(),
316 "echo 'error: could not compile' >&2; exit 101".to_string(),
317 ],
318 None,
319 );
320
321 let err = builder.build(Path::new("/tmp/my-crate")).await.unwrap_err();
322
323 assert!(err.to_string().contains("error: could not compile"));
324 }
325
326 #[cfg(unix)]
327 fn process_is_alive(pid: &str) -> bool {
328 std::process::Command::new("kill")
329 .args(["-0", pid])
330 .status()
331 .unwrap()
332 .success()
333 }
334
335 #[cfg(unix)]
336 #[tokio::test]
337 async fn build_timeout_kills_build_and_its_child_processes() {
338 let dir = tempfile::tempdir().unwrap();
339 let pid_file = dir.path().join("pid");
340 let builder = DocBuilder::new(
341 PathBuf::from("unused"),
342 vec![
343 "/bin/sh".to_string(),
344 "-c".to_string(),
345 "sleep 30 & echo $! > \"$0\"; wait".to_string(),
346 pid_file.display().to_string(),
347 ],
348 None,
349 )
350 .with_timeout(Duration::from_millis(500));
351
352 let started = std::time::Instant::now();
353 let err = builder.build(Path::new("/tmp/my-crate")).await.unwrap_err();
354
355 assert!(matches!(err, DocsError::BuildTimeout(_)));
356 assert!(started.elapsed() < Duration::from_secs(10));
357 assert_process_stops(&pid_file).await;
358 }
359
360 #[cfg(unix)]
361 async fn assert_process_stops(pid_file: &Path) {
362 let pid = std::fs::read_to_string(pid_file).unwrap();
363 let deadline = std::time::Instant::now() + Duration::from_secs(5);
365 while process_is_alive(pid.trim()) && std::time::Instant::now() < deadline {
366 tokio::time::sleep(Duration::from_millis(50)).await;
367 }
368 assert!(!process_is_alive(pid.trim()));
369 }
370
371 #[cfg(unix)]
372 #[tokio::test]
373 async fn shutdown_kills_build_and_its_child_processes() {
374 let dir = tempfile::tempdir().unwrap();
375 let pid_file = dir.path().join("pid");
376 let shutdown = CancellationToken::new();
377 let builder = DocBuilder::new(
378 PathBuf::from("unused"),
379 vec![
380 "/bin/sh".to_string(),
381 "-c".to_string(),
382 "sleep 30 & echo $! > \"$0\"; wait".to_string(),
383 pid_file.display().to_string(),
384 ],
385 None,
386 )
387 .with_shutdown(shutdown.clone());
388
389 let started = std::time::Instant::now();
390 let build = tokio::spawn(async move { builder.build(Path::new("/tmp/my-crate")).await });
391 while std::fs::read_to_string(&pid_file).map_or(true, |pid| pid.trim().is_empty()) {
392 tokio::time::sleep(Duration::from_millis(50)).await;
393 }
394 shutdown.cancel();
395 let err = build.await.unwrap().unwrap_err();
396
397 assert!(matches!(err, DocsError::BuildStopped));
398 assert!(started.elapsed() < Duration::from_secs(10));
399 assert_process_stops(&pid_file).await;
400 }
401
402 #[test]
403 fn zero_timeout_disables_the_timeout() {
404 let builder =
405 DocBuilder::new(PathBuf::from("unused"), vec![], None).with_timeout(Duration::ZERO);
406
407 assert_eq!(builder.timeout, None);
408 }
409
410 #[tokio::test]
411 async fn build_fails_when_program_does_not_exist() {
412 let builder = DocBuilder::new(PathBuf::from("/does/not/exist/kellnr"), vec![], None);
413
414 let err = builder.build(Path::new("/tmp/my-crate")).await.unwrap_err();
415
416 assert!(
417 err.to_string()
418 .contains("failed to start docs build process")
419 );
420 }
421}