1use std::collections::HashSet;
2use std::path::Path;
3use std::sync::Arc;
4use std::time::Duration;
5
6use cargo::GlobalContext;
7use cargo::core::Workspace;
8use cargo::core::compiler::UserIntent;
9use cargo::core::resolver::CliFeatures;
10use cargo::ops::{self, CompileOptions, DocOptions, OutputFormat};
11use flate2::read::GzDecoder;
12use kellnr_common::original_name::OriginalName;
13use kellnr_common::version::Version;
14use kellnr_db::{DbProvider, DocQueueEntry};
15use kellnr_settings::Settings;
16use kellnr_storage::docs_storage::DocsStorage;
17use kellnr_storage::kellnr_crate_storage::KellnrCrateStorage;
18use tokio::fs::remove_dir_all;
19use tracing::error;
20
21use crate::build_process::DocBuilder;
22use crate::compute_doc_url;
23use crate::docs_error::DocsError;
24use crate::unpack::{ExtractLimits, unpack_crate};
25use crate::upload::upload_dir_and_prune;
26
27pub fn doc_extraction_queue(
28 db: Arc<dyn DbProvider>,
29 cs: Arc<KellnrCrateStorage>,
30 docs_storage: Arc<DocsStorage>,
31 settings: &Settings,
32) {
33 let program = match std::env::current_exe() {
34 Ok(program) => program,
35 Err(e) => {
36 error!("Rustdoc generation disabled, cannot locate the kellnr binary: {e}");
37 return;
38 }
39 };
40 let builder = DocBuilder::new(
41 program,
42 settings.docs.build_wrapper.clone(),
43 settings
44 .proxy
45 .cratesio_index_override()
46 .map(ToString::to_string),
47 )
48 .with_timeout(Duration::from_secs(settings.docs.build_timeout_seconds));
49 let worker = DocQueueWorker {
50 db,
51 cs,
52 docs_storage,
53 path_prefix: settings.origin.path.clone(),
54 builder,
55 limits: ExtractLimits::from_max_docs_size(settings.docs.max_size),
56 };
57
58 tokio::spawn(async move {
59 loop {
60 tokio::time::sleep(Duration::from_secs(10)).await;
61 if let Err(e) = worker.process_queue().await {
62 error!("Rustdoc generation loop failed: {e}");
63 }
64 }
65 });
66}
67
68struct DocQueueWorker {
70 db: Arc<dyn DbProvider>,
71 cs: Arc<KellnrCrateStorage>,
72 docs_storage: Arc<DocsStorage>,
73 path_prefix: String,
74 builder: DocBuilder,
75 limits: ExtractLimits,
76}
77
78impl DocQueueWorker {
79 async fn process_queue(&self) -> Result<(), DocsError> {
80 let entries = self.db.get_doc_queue().await?;
81 let mut built = HashSet::new();
82
83 for entry in entries {
84 let is_first = built.insert((entry.normalized_name.to_string(), entry.version.clone()));
86 let result = if is_first {
87 Some(self.extract_docs(&entry).await)
88 } else {
89 None
90 };
91
92 if let Err(e) = clean_up(&entry.path).await {
94 error!("Failed to delete temporary rustdoc queue folder: {e}");
95 }
96
97 match result {
98 Some(Ok(())) => {
99 let version = Version::from_unchecked_str(&entry.version);
100 let docs_link =
101 compute_doc_url(&entry.normalized_name, &version, &self.path_prefix);
102 self.db
103 .update_docs_link(&entry.normalized_name, &version, &docs_link)
104 .await?;
105 }
106 Some(Err(e)) => error!("Failed to extract docs from crate: {e}"),
107 None => {}
108 }
109 self.db.delete_doc_queue(entry.id).await?;
110 }
111
112 Ok(())
113 }
114
115 async fn extract_docs(&self, doc: &DocQueueEntry) -> Result<(), DocsError> {
116 let orig_name = OriginalName::from_unchecked(doc.normalized_name.to_string());
118 let version = Version::from_unchecked_str(&doc.version);
119 let contents = self.cs.get(&orig_name, &version).await.ok_or_else(|| {
120 error!("Failed to get crate from storage");
121 DocsError::CrateDoesNotExist(doc.normalized_name.to_string(), doc.version.clone())
122 })?;
123 let crate_root = format!("{}-{}", doc.normalized_name, doc.version);
124 let tar = GzDecoder::new(std::io::Cursor::new(contents));
125 let dest = doc.path.clone();
126 let root = crate_root.clone();
127 let limits = self.limits;
128 tokio::task::spawn_blocking(move || unpack_crate(tar, &dest, &root, limits))
129 .await
130 .map_err(|_| DocsError::ExtractFailed)??;
131
132 let generated_docs_path = &doc.path.join(crate_root);
134 strip_rust_toolchain_files(generated_docs_path).await?;
135
136 self.builder.build(generated_docs_path).await?;
139
140 let from = generated_docs_path.join("target").join("doc");
145 upload_dir_and_prune(
146 &from,
147 "doc",
148 &doc.normalized_name.to_string(),
149 &doc.version,
150 &self.docs_storage,
151 )
152 .await?;
153
154 Ok(())
155 }
156}
157
158async fn clean_up(path: &Path) -> Result<(), DocsError> {
159 match remove_dir_all(path).await {
160 Err(e) if e.kind() != std::io::ErrorKind::NotFound => Err(e.into()),
161 _ => Ok(()),
162 }
163}
164
165async fn strip_rust_toolchain_files(crate_path: &Path) -> Result<(), DocsError> {
188 for name in ["rust-toolchain.toml", "rust-toolchain"] {
189 let path = crate_path.join(name);
190 match tokio::fs::remove_file(&path).await {
191 Ok(()) => {}
192 Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
193 Err(e) => return Err(e.into()),
194 }
195 }
196 Ok(())
197}
198
199fn build_doc_context(cratesio_index: Option<&str>) -> Result<GlobalContext, DocsError> {
206 let mut ctx = GlobalContext::default().map_err(|e| DocsError::CargoError(e.to_string()))?;
207
208 if let Some(index) = cratesio_index {
209 let cli_config = vec![
210 "source.crates-io.replace-with=\"kellnr-proxy\"".to_string(),
211 format!("source.kellnr-proxy.registry=\"sparse+{index}\""),
212 ];
213 ctx.configure(0, false, None, false, false, false, &None, &[], &cli_config)
214 .map_err(|e| DocsError::CargoError(e.to_string()))?;
215 }
216
217 Ok(ctx)
218}
219
220fn doc_features(workspace: &Workspace) -> Result<CliFeatures, DocsError> {
229 let package = workspace
230 .current()
231 .map_err(|e| DocsError::CargoError(e.to_string()))?;
232 let docs_rs = package
234 .manifest()
235 .custom_metadata()
236 .and_then(|m| m.get("docs"))
237 .and_then(|m| m.get("rs"));
238
239 let flag = |key: &str| docs_rs.and_then(|m| m.get(key)?.as_bool());
240 let all_features = flag("all-features").unwrap_or(true);
241 let no_default_features = flag("no-default-features").unwrap_or(false);
242 let features: Vec<String> = docs_rs
243 .and_then(|m| m.get("features"))
244 .and_then(|v| v.as_array())
245 .map(|a| {
246 a.iter()
247 .filter_map(|f| f.as_str().map(str::to_string))
248 .collect()
249 })
250 .unwrap_or_default();
251
252 CliFeatures::from_command_line(&features, all_features, !no_default_features)
253 .map_err(|e| DocsError::CargoError(e.to_string()))
254}
255
256pub fn generate_docs(
262 crate_path: impl AsRef<Path>,
263 cratesio_index: Option<&str>,
264) -> Result<(), DocsError> {
265 let manifest_path = crate_path.as_ref().join("Cargo.toml").canonicalize()?;
266 let ctx = build_doc_context(cratesio_index)?;
267 let workspace =
268 Workspace::new(&manifest_path, &ctx).map_err(|e| DocsError::CargoError(e.to_string()))?;
269 let compile_opts = CompileOptions {
270 cli_features: doc_features(&workspace)?,
271 ..CompileOptions::new(
272 &ctx,
273 UserIntent::Doc {
274 deps: false,
275 json: false,
276 },
277 )
278 .map_err(|e| DocsError::CargoError(e.to_string()))?
279 };
280 let options = DocOptions {
281 open_result: false,
282 compile_opts,
283 output_format: OutputFormat::Html,
284 };
285 ops::doc(&workspace, &options).map_err(|e| DocsError::CargoError(e.to_string()))?;
286 Ok(())
287}
288
289#[cfg(test)]
290mod tests {
291 use std::path::PathBuf;
292
293 use cargo::core::SourceId;
294 use cargo::sources::SourceConfigMap;
295 use kellnr_common::normalized_name::NormalizedName;
296 use kellnr_db::mock::MockDb;
297 use kellnr_storage::fs_storage::FSStorage;
298
299 use super::*;
300
301 fn crate_tarball(name: &str, version: &str) -> Vec<u8> {
302 let mut builder = tar::Builder::new(flate2::write::GzEncoder::new(
303 Vec::new(),
304 flate2::Compression::default(),
305 ));
306 let manifest = format!("[package]\nname = \"{name}\"\nversion = \"{version}\"\n");
307 let mut header = tar::Header::new_gnu();
308 header.set_size(manifest.len() as u64);
309 header.set_mode(0o644);
310 header.set_cksum();
311 builder
312 .append_data(
313 &mut header,
314 format!("{name}-{version}/Cargo.toml"),
315 manifest.as_bytes(),
316 )
317 .unwrap();
318 builder.into_inner().unwrap().finish().unwrap()
319 }
320
321 fn queue_entry(id: i64, name: &str, path: PathBuf) -> DocQueueEntry {
322 DocQueueEntry {
323 id,
324 normalized_name: NormalizedName::from_unchecked(name.to_string()),
325 version: "1.0.0".to_string(),
326 path,
327 }
328 }
329
330 async fn failing_worker(db: MockDb, data_dir: &Path, build_log: &Path) -> DocQueueWorker {
332 let mut settings = kellnr_settings::test_settings();
333 settings.registry.data_dir = data_dir.display().to_string();
334 let storage = FSStorage::new(data_dir.join("crates").to_str().unwrap()).unwrap();
335 let cs = KellnrCrateStorage::new(&settings, Box::new(storage));
336 for name in ["a", "b"] {
337 cs.put(
338 &OriginalName::from_unchecked(name.to_string()),
339 &Version::from_unchecked_str("1.0.0"),
340 crate_tarball(name, "1.0.0").into(),
341 )
342 .await
343 .unwrap();
344 }
345 let docs_storage = FSStorage::new(data_dir.join("docs").to_str().unwrap()).unwrap();
346 DocQueueWorker {
347 db: Arc::new(db),
348 cs: Arc::new(cs),
349 docs_storage: Arc::new(DocsStorage::new(Box::new(docs_storage))),
350 path_prefix: String::new(),
351 builder: DocBuilder::new(
352 PathBuf::from("unused"),
353 vec![
354 "/bin/sh".to_string(),
355 "-c".to_string(),
356 "echo \"$3\" >> \"$0\"; exit 1".to_string(),
357 build_log.display().to_string(),
358 ],
359 None,
360 ),
361 limits: ExtractLimits::from_max_docs_size(1),
362 }
363 }
364
365 #[cfg(unix)]
366 #[tokio::test]
367 async fn process_queue_builds_duplicates_once_and_always_cleans_up() {
368 let data_dir = tempfile::tempdir().unwrap();
369 let build_log = data_dir.path().join("builds.log");
370 let paths: Vec<PathBuf> = (0..3)
371 .map(|i| {
372 let path = data_dir.path().join(format!("queue{i}"));
373 std::fs::create_dir(&path).unwrap();
374 path
375 })
376 .collect();
377 let queued = paths.clone();
378 let mut db = MockDb::new();
379 db.expect_get_doc_queue().returning(move || {
380 Ok(vec![
381 queue_entry(0, "a", queued[0].clone()),
382 queue_entry(1, "a", queued[1].clone()),
383 queue_entry(2, "b", queued[2].clone()),
384 ])
385 });
386 db.expect_delete_doc_queue().times(3).returning(|_| Ok(()));
387 db.expect_update_docs_link().never();
388 let worker = failing_worker(db, data_dir.path(), &build_log).await;
389
390 worker.process_queue().await.unwrap();
391
392 let builds = std::fs::read_to_string(&build_log).unwrap();
393 assert_eq!(builds.lines().count(), 2, "builds: {builds}");
394 assert!(paths.iter().all(|p| !p.exists()));
395 }
396
397 #[tokio::test]
398 async fn clean_up_ignores_missing_directory() {
399 let dir = tempfile::tempdir().unwrap();
400
401 clean_up(&dir.path().join("missing")).await.unwrap();
402 }
403
404 #[test]
405 fn no_index_override_does_not_inject_proxy_source() {
406 let ctx = build_doc_context(None).unwrap();
407 assert!(
413 ctx.get_string("source.kellnr-proxy.registry")
414 .unwrap()
415 .is_none()
416 );
417 }
418
419 #[test]
420 fn index_override_replaces_crates_io_source_with_proxy() {
421 let ctx = build_doc_context(Some("https://rsproxy.cn/index/")).unwrap();
422
423 assert_eq!(
425 ctx.get_string("source.crates-io.replace-with")
426 .unwrap()
427 .map(|v| v.val),
428 Some("kellnr-proxy".to_string())
429 );
430
431 let map = SourceConfigMap::new(&ctx).unwrap();
434 let crates_io = SourceId::crates_io(&ctx).unwrap();
435 let source = map.load(crates_io, &HashSet::new()).unwrap();
436 let replaced = source.replaced_source_id();
437 assert!(!replaced.is_crates_io());
438 assert!(
439 replaced.url().as_str().contains("rsproxy.cn/index/"),
440 "expected replacement source to point at the proxy index, got {}",
441 replaced.url()
442 );
443 }
444
445 fn features_for(manifest: &str) -> CliFeatures {
446 let dir = tempfile::tempdir().unwrap();
447 std::fs::write(dir.path().join("Cargo.toml"), manifest).unwrap();
448 std::fs::create_dir(dir.path().join("src")).unwrap();
449 std::fs::write(dir.path().join("src").join("lib.rs"), "").unwrap();
450 let ctx = build_doc_context(None).unwrap();
451 let manifest_path = dir.path().join("Cargo.toml").canonicalize().unwrap();
452 let workspace = Workspace::new(&manifest_path, &ctx).unwrap();
453 doc_features(&workspace).unwrap()
454 }
455
456 const PACKAGE: &str = "[package]\nname = \"x\"\nversion = \"0.1.0\"\nedition = \"2021\"\n\n[features]\na = []\nb = []\n";
457
458 fn feature_names(features: &CliFeatures) -> Vec<String> {
459 let mut names: Vec<String> = features.features.iter().map(ToString::to_string).collect();
460 names.sort();
461 names
462 }
463
464 #[test]
465 fn doc_features_default_to_all_features_without_metadata() {
466 let features = features_for(PACKAGE);
467
468 assert!(features.all_features);
469 assert!(features.uses_default_features);
470 assert!(features.features.is_empty());
471 }
472
473 #[test]
474 fn doc_features_honor_docs_rs_metadata() {
475 let features = features_for(&format!(
476 "{PACKAGE}\n[package.metadata.docs.rs]\nall-features = false\nno-default-features = true\nfeatures = [\"a\", \"b\"]\n"
477 ));
478
479 assert!(!features.all_features);
480 assert!(!features.uses_default_features);
481 assert_eq!(feature_names(&features), vec!["a", "b"]);
482 }
483
484 #[test]
485 fn doc_features_keep_all_features_when_only_features_are_listed() {
486 let features = features_for(&format!(
487 "{PACKAGE}\n[package.metadata.docs.rs]\nfeatures = [\"a\"]\n"
488 ));
489
490 assert!(features.all_features);
491 assert!(features.uses_default_features);
492 assert_eq!(feature_names(&features), vec!["a"]);
493 }
494
495 #[test]
496 fn doc_features_ignore_unrelated_metadata() {
497 let features = features_for(&format!(
498 "{PACKAGE}\n[package.metadata.other]\nall-features = false\n"
499 ));
500
501 assert!(features.all_features);
502 }
503
504 #[test]
505 fn generate_docs_builds_crate_with_mutually_exclusive_features() {
506 let dir = tempfile::tempdir().unwrap();
507 std::fs::write(
508 dir.path().join("Cargo.toml"),
509 format!("{PACKAGE}\n[package.metadata.docs.rs]\nall-features = false\n"),
510 )
511 .unwrap();
512 std::fs::create_dir(dir.path().join("src")).unwrap();
513 std::fs::write(
514 dir.path().join("src").join("lib.rs"),
515 "#[cfg(all(feature = \"a\", feature = \"b\"))]\ncompile_error!(\"a and b are exclusive\");\n",
516 )
517 .unwrap();
518
519 generate_docs(dir.path(), None).unwrap();
520
521 assert!(dir.path().join("target/doc/x/index.html").exists());
522 }
523
524 #[tokio::test]
525 async fn strip_rust_toolchain_files_removes_both_variants() {
526 let dir = tempfile::tempdir().unwrap();
527 let toml_path = dir.path().join("rust-toolchain.toml");
528 let legacy_path = dir.path().join("rust-toolchain");
529 tokio::fs::write(&toml_path, "[toolchain]\nchannel = \"1.65.0\"\n")
530 .await
531 .unwrap();
532 tokio::fs::write(&legacy_path, "1.65.0\n").await.unwrap();
533
534 strip_rust_toolchain_files(dir.path()).await.unwrap();
535
536 assert!(!toml_path.exists());
537 assert!(!legacy_path.exists());
538 }
539
540 #[tokio::test]
541 async fn strip_rust_toolchain_files_noop_when_absent() {
542 let dir = tempfile::tempdir().unwrap();
543 let unrelated = dir.path().join("Cargo.toml");
544 tokio::fs::write(&unrelated, "[package]\nname = \"x\"\n")
545 .await
546 .unwrap();
547
548 strip_rust_toolchain_files(dir.path()).await.unwrap();
549
550 assert!(unrelated.exists());
551 }
552
553 #[tokio::test]
554 async fn strip_rust_toolchain_files_only_removes_named_files() {
555 let dir = tempfile::tempdir().unwrap();
556 let toml_path = dir.path().join("rust-toolchain.toml");
557 let cargo_toml = dir.path().join("Cargo.toml");
558 let src = dir.path().join("src");
559 tokio::fs::write(&toml_path, "").await.unwrap();
560 tokio::fs::write(&cargo_toml, "").await.unwrap();
561 tokio::fs::create_dir(&src).await.unwrap();
562
563 strip_rust_toolchain_files(dir.path()).await.unwrap();
564
565 assert!(!toml_path.exists());
566 assert!(cargo_toml.exists());
567 assert!(src.exists());
568 }
569}