1const DEFAULT_EXEC_SQLITE: &str = "sqlite3 -json /tmp/test.db";
7const DEFAULT_EXEC_OSQUERY: &str = "osqueryi --json";
8const DEFAULT_EXEC_FALLBACK: &str = "cat";
9
10use std::collections::hash_map::Entry;
11use std::collections::HashMap;
12use std::fmt::Write;
13use std::path::Path;
14
15use mdbook_preprocessor::book::{Book, BookItem, Chapter};
16use mdbook_preprocessor::errors::Error;
17use mdbook_preprocessor::{Preprocessor, PreprocessorContext};
18use pulldown_cmark::{CodeBlockKind, Event, Parser, Tag, TagEnd};
19
20use crate::command::RealCommandRunner;
21use crate::config::{Config, ValidatorConfig};
22use crate::container::ValidatorContainer;
23use crate::error::ValidatorError;
24use crate::host_validator;
25use crate::parser::{extract_markers, parse_info_string, ExtractedMarkers};
26use crate::transpiler::strip_markers;
27
28pub struct ValidatorPreprocessor;
30
31impl ValidatorPreprocessor {
32 #[must_use]
34 pub fn new() -> Self {
35 Self
36 }
37}
38
39impl Default for ValidatorPreprocessor {
40 fn default() -> Self {
41 Self::new()
42 }
43}
44
45impl Preprocessor for ValidatorPreprocessor {
46 fn name(&self) -> &'static str {
47 "validator"
48 }
49
50 fn run(&self, ctx: &PreprocessorContext, mut book: Book) -> Result<Book, Error> {
51 let config = Config::from_context(ctx)
53 .map_err(|e| Error::msg(format!("Failed to parse config: {e}")))?;
54
55 let rt = tokio::runtime::Builder::new_current_thread()
57 .enable_all()
58 .build()
59 .map_err(|e| Error::msg(format!("Failed to create tokio runtime: {e}")))?;
60
61 rt.block_on(async {
62 self.run_async_with_config(&mut book, &config, &ctx.root)
63 .await
64 })?;
65
66 Ok(book)
67 }
68
69 fn supports_renderer(&self, renderer: &str) -> Result<bool, anyhow::Error> {
70 let _ = renderer;
73 Ok(true)
74 }
75}
76
77impl ValidatorPreprocessor {
78 pub fn process_book_with_script(
83 &self,
84 mut book: Book,
85 validator_script: &[u8],
86 ) -> Result<Book, Error> {
87 let rt = tokio::runtime::Builder::new_current_thread()
88 .enable_all()
89 .build()
90 .map_err(|e| Error::msg(format!("Failed to create tokio runtime: {e}")))?;
91
92 rt.block_on(async {
93 self.run_async_with_script(&mut book, validator_script)
94 .await
95 })?;
96
97 Ok(book)
98 }
99
100 pub fn process_book_with_config(
104 &self,
105 mut book: Book,
106 config: &Config,
107 book_root: &Path,
108 ) -> Result<Book, Error> {
109 let rt = tokio::runtime::Builder::new_current_thread()
110 .enable_all()
111 .build()
112 .map_err(|e| Error::msg(format!("Failed to create tokio runtime: {e}")))?;
113
114 rt.block_on(async {
115 self.run_async_with_config(&mut book, config, book_root)
116 .await
117 })?;
118
119 Ok(book)
120 }
121
122 async fn run_async_with_config(
124 &self,
125 book: &mut Book,
126 config: &Config,
127 book_root: &Path,
128 ) -> Result<(), Error> {
129 let mut containers: HashMap<String, ValidatorContainer> = HashMap::new();
131
132 for item in &mut book.items {
133 self.process_book_item_with_config(item, config, book_root, &mut containers)
134 .await?;
135 }
136
137 Ok(())
138 }
139
140 async fn run_async_with_script(
142 &self,
143 book: &mut Book,
144 validator_script: &[u8],
145 ) -> Result<(), Error> {
146 let container = ValidatorContainer::start(validator_script)
147 .await
148 .map_err(|e| Error::msg(format!("Failed to start container: {e}")))?;
149
150 for item in &mut book.items {
151 self.process_book_item(item, &container).await?;
152 }
153
154 Ok(())
155 }
156
157 async fn process_book_item(
158 &self,
159 item: &mut BookItem,
160 container: &ValidatorContainer,
161 ) -> Result<(), Error> {
162 if let BookItem::Chapter(chapter) = item {
163 self.process_chapter(chapter, container).await?;
164
165 for sub_item in &mut chapter.sub_items {
167 Box::pin(self.process_book_item(sub_item, container)).await?;
168 }
169 }
170 Ok(())
171 }
172
173 async fn process_book_item_with_config(
174 &self,
175 item: &mut BookItem,
176 config: &Config,
177 book_root: &Path,
178 containers: &mut HashMap<String, ValidatorContainer>,
179 ) -> Result<(), Error> {
180 if let BookItem::Chapter(chapter) = item {
181 self.process_chapter_with_config(chapter, config, book_root, containers)
182 .await?;
183
184 for sub_item in &mut chapter.sub_items {
186 Box::pin(
187 self.process_book_item_with_config(sub_item, config, book_root, containers),
188 )
189 .await?;
190 }
191 }
192 Ok(())
193 }
194
195 async fn process_chapter(
196 &self,
197 chapter: &mut Chapter,
198 container: &ValidatorContainer,
199 ) -> Result<(), Error> {
200 if chapter.content.is_empty() {
201 return Ok(());
202 }
203
204 let blocks = Self::find_validator_blocks(&chapter.content);
206
207 if blocks.is_empty() {
208 return Ok(());
209 }
210
211 for block in &blocks {
213 if block.skip {
214 continue;
215 }
216
217 let validation_content = block.markers.validation_content();
218 let result = container
219 .exec_with_env(
220 block.markers.setup.as_deref(),
221 &validation_content,
222 block.markers.assertions.as_deref(),
223 block.markers.expect.as_deref(),
224 )
225 .await
226 .map_err(|e| {
227 Error::msg(format!(
228 "Validation exec failed in '{}': {}",
229 chapter.name, e
230 ))
231 })?;
232
233 if result.exit_code != 0 {
234 let mut error_msg = format!(
235 "Validation failed in '{}' (exit code {}):\n\nCode:\n{}\n",
236 chapter.name, result.exit_code, block.markers.visible_content
237 );
238 if !result.stderr.is_empty() {
239 let _ = write!(error_msg, "\nValidator stderr:\n{}", result.stderr);
240 }
241 if !result.stdout.is_empty() {
242 let _ = write!(error_msg, "\nValidator stdout:\n{}", result.stdout);
243 }
244 return Err(Error::msg(error_msg));
245 }
246 }
247
248 chapter.content = Self::strip_markers_from_chapter(&chapter.content);
250
251 Ok(())
252 }
253
254 async fn process_chapter_with_config(
255 &self,
256 chapter: &mut Chapter,
257 config: &Config,
258 book_root: &Path,
259 containers: &mut HashMap<String, ValidatorContainer>,
260 ) -> Result<(), Error> {
261 if chapter.content.is_empty() {
262 return Ok(());
263 }
264
265 let blocks = Self::find_validator_blocks(&chapter.content);
267
268 if blocks.is_empty() {
269 return Ok(());
270 }
271
272 for block in &blocks {
274 if block.skip && block.hidden {
275 return Err(Error::new(ValidatorError::MutuallyExclusiveAttributes));
276 }
277 }
278
279 for block in &blocks {
281 if block.skip {
282 continue;
283 }
284
285 let validator_config = config.get_validator(&block.validator_name).map_err(|e| {
287 Error::msg(format!(
288 "Unknown validator '{}': {}",
289 block.validator_name, e
290 ))
291 })?;
292
293 let container = self
295 .get_or_start_container(&block.validator_name, config, book_root, containers)
296 .await?;
297
298 self.validate_block_host_based(
300 container,
301 validator_config,
302 block,
303 &chapter.name,
304 book_root,
305 )
306 .await?;
307 }
308
309 chapter.content = Self::strip_markers_from_chapter(&chapter.content);
311
312 Ok(())
313 }
314
315 async fn validate_block_host_based(
319 &self,
320 container: &ValidatorContainer,
321 validator_config: &ValidatorConfig,
322 block: &ValidatorBlock,
323 chapter_name: &str,
324 book_root: &Path,
325 ) -> Result<(), Error> {
326 let script_path = book_root.join(&validator_config.script);
328 if !script_path.exists() {
329 return Err(Error::msg(format!(
330 "Failed to read validator script '{}': file not found",
331 script_path.display()
332 )));
333 }
334
335 let exec_cmd = Self::get_exec_command(&block.validator_name, validator_config);
337
338 if let Some(setup) = &block.markers.setup {
341 let setup_script = setup.trim();
342 if !setup_script.is_empty() {
343 let setup_result = container
344 .exec_raw(&["sh", "-c", setup_script])
345 .await
346 .map_err(|e| Error::msg(format!("Setup exec failed: {e}")))?;
347
348 if setup_result.exit_code != 0 {
349 #[allow(clippy::cast_possible_truncation)]
350 return Err(ValidatorError::SetupFailed {
351 exit_code: setup_result.exit_code as i32,
352 message: format!(
353 "in '{}' (validator: {}):\n\nScript:\n{}\n\nError:\n{}",
354 chapter_name, block.validator_name, setup_script, setup_result.stderr
355 ),
356 }
357 .into());
358 }
359 }
360 }
361
362 let query_sql = block.markers.validation_content();
366 let query_sql = query_sql.trim();
367 if query_sql.is_empty() {
368 return Err(Error::msg(format!(
369 "Validation failed in '{}' (validator: {}): Query content is empty",
370 chapter_name, block.validator_name
371 )));
372 }
373
374 let query_result = container
376 .exec_with_stdin(&["sh", "-c", &exec_cmd], query_sql)
377 .await
378 .map_err(|e| Error::msg(format!("Query exec failed: {e}")))?;
379
380 if query_result.exit_code != 0 {
381 return Err(Error::msg(format!(
382 "Query failed in '{}' (validator: {}):\n\nSQL:\n{}\n\nError:\n{}",
383 chapter_name, block.validator_name, query_sql, query_result.stderr
384 )));
385 }
386
387 let script_path_str = script_path
390 .to_str()
391 .ok_or_else(|| Error::msg(format!("Invalid script path: {}", script_path.display())))?;
392
393 let validation_result = host_validator::run_validator(
394 &RealCommandRunner,
395 script_path_str,
396 &query_result.stdout,
397 block.markers.assertions.as_deref(),
398 block.markers.expect.as_deref(),
399 Some(&query_result.stderr), )
401 .map_err(|e| {
402 Error::msg(format!(
403 "Host validator failed in '{}' (validator: {}): {}",
404 chapter_name, block.validator_name, e
405 ))
406 })?;
407
408 if validation_result.exit_code != 0 {
409 let mut error_msg = format!(
410 "in '{}' (validator: {}):\n\nCode:\n{}\n",
411 chapter_name, block.validator_name, block.markers.visible_content
412 );
413 if !validation_result.stderr.is_empty() {
414 let _ = write!(
415 error_msg,
416 "\nValidator stderr:\n{}",
417 validation_result.stderr
418 );
419 }
420 if !validation_result.stdout.is_empty() {
421 let _ = write!(
422 error_msg,
423 "\nValidator stdout:\n{}",
424 validation_result.stdout
425 );
426 }
427 return Err(ValidatorError::ValidationFailed {
428 exit_code: validation_result.exit_code,
429 message: error_msg,
430 }
431 .into());
432 }
433
434 Ok(())
435 }
436
437 fn get_exec_command(validator_name: &str, config: &ValidatorConfig) -> String {
441 config
442 .exec_command
443 .clone()
444 .unwrap_or_else(|| match validator_name {
445 "sqlite" => DEFAULT_EXEC_SQLITE.to_owned(),
446 "osquery" => DEFAULT_EXEC_OSQUERY.to_owned(),
447 _ => DEFAULT_EXEC_FALLBACK.to_owned(),
448 })
449 }
450
451 async fn get_or_start_container<'a>(
453 &self,
454 validator_name: &str,
455 config: &Config,
456 book_root: &Path,
457 containers: &'a mut HashMap<String, ValidatorContainer>,
458 ) -> Result<&'a ValidatorContainer, Error> {
459 match containers.entry(validator_name.to_owned()) {
460 Entry::Occupied(entry) => Ok(entry.into_mut()),
461 Entry::Vacant(entry) => {
462 let validator_config = config.get_validator(validator_name).map_err(|e| {
464 Error::msg(format!("Unknown validator '{validator_name}': {e}"))
465 })?;
466
467 validator_config.validate(validator_name)?;
469
470 let mount = if let Some(ref fixtures_dir) = config.fixtures_dir {
472 let fixtures_path = if fixtures_dir.is_absolute() {
474 fixtures_dir.clone()
475 } else {
476 book_root.join(fixtures_dir)
477 };
478
479 if !fixtures_path.exists() {
481 return Err(Error::msg(format!(
482 "fixtures_dir '{}' does not exist",
483 fixtures_path.display()
484 )));
485 }
486 if !fixtures_path.is_dir() {
487 return Err(Error::msg(format!(
488 "fixtures_dir '{}' is not a directory",
489 fixtures_path.display()
490 )));
491 }
492
493 Some((fixtures_path, "/fixtures"))
494 } else {
495 None
496 };
497
498 let container = ValidatorContainer::start_raw_with_mount(
500 &validator_config.container,
501 mount.as_ref().map(|(p, c)| (p.as_path(), *c)),
502 )
503 .await
504 .map_err(|e| {
505 Error::msg(format!(
506 "Failed to start container '{}': {}",
507 validator_config.container, e
508 ))
509 })?;
510
511 Ok(entry.insert(container))
512 }
513 }
514 }
515
516 fn find_validator_blocks(content: &str) -> Vec<ValidatorBlock> {
518 let mut blocks = Vec::new();
519 let parser = Parser::new(content);
520
521 let mut in_code_block = false;
522 let mut current_info = String::new();
523 let mut current_content = String::new();
524
525 for event in parser {
526 match event {
527 Event::Start(Tag::CodeBlock(CodeBlockKind::Fenced(info))) => {
528 in_code_block = true;
529 current_info = info.to_string();
530 current_content.clear();
531 }
532 Event::Text(text) if in_code_block => {
533 current_content.push_str(&text);
534 }
535 Event::End(TagEnd::CodeBlock) if in_code_block => {
536 in_code_block = false;
537
538 let (_language, validator, skip, hidden) = parse_info_string(¤t_info);
539
540 if let Some(validator_name) = validator {
542 if !validator_name.is_empty() {
544 let markers = extract_markers(¤t_content);
545 blocks.push(ValidatorBlock {
546 validator_name,
547 markers,
548 skip,
549 hidden,
550 });
551 }
552 }
553 }
554 _ => {}
555 }
556 }
557
558 blocks
559 }
560
561 fn strip_markers_from_chapter(content: &str) -> String {
565 let mut result = String::new();
566 let parser = Parser::new(content);
567
568 let mut in_code_block = false;
569 let mut current_info = String::new();
570 let mut current_hidden = false;
571
572 for event in parser {
573 match &event {
574 Event::Start(Tag::CodeBlock(CodeBlockKind::Fenced(info))) => {
575 in_code_block = true;
576 current_info = info.to_string();
577 let (_language, _validator, _skip, hidden) = parse_info_string(¤t_info);
578 current_hidden = hidden;
579
580 if !current_hidden {
582 result.push_str("```");
583 result.push_str(¤t_info);
584 result.push('\n');
585 }
586 }
587 Event::Text(text) if in_code_block => {
588 if current_hidden {
590 continue;
591 }
592
593 let (_language, validator, _skip, _hidden) = parse_info_string(¤t_info);
594
595 if validator.is_some() {
597 let stripped = strip_markers(text);
598 let trimmed = stripped.trim();
600 if !trimmed.is_empty() {
601 result.push_str(trimmed);
602 result.push('\n');
603 }
604 } else {
605 result.push_str(text);
606 }
607 }
608 Event::End(TagEnd::CodeBlock) if in_code_block => {
609 in_code_block = false;
610 if !current_hidden {
612 result.push_str("```\n");
613 }
614 current_hidden = false;
615 }
616 Event::Start(Tag::CodeBlock(CodeBlockKind::Indented)) => {
617 in_code_block = true;
619 current_info.clear();
620 current_hidden = false;
621 }
622 Event::End(TagEnd::CodeBlock) => {
623 in_code_block = false;
624 current_hidden = false;
625 }
626 Event::SoftBreak | Event::HardBreak => {
627 if !in_code_block {
628 result.push('\n');
629 }
630 }
631 Event::Text(text) if !in_code_block => {
632 result.push_str(text);
633 }
634 Event::End(TagEnd::Paragraph | TagEnd::Heading(_)) => {
635 result.push_str("\n\n");
636 }
637 Event::Start(Tag::Heading { level, .. }) => {
638 result.push_str(&"#".repeat(*level as usize));
639 result.push(' ');
640 }
641 _ => {}
642 }
643 }
644
645 result.trim().to_owned()
646 }
647}
648
649struct ValidatorBlock {
651 validator_name: String,
653 markers: ExtractedMarkers,
655 skip: bool,
657 hidden: bool,
659}
660
661#[cfg(test)]
662#[allow(clippy::needless_raw_string_hashes)]
663mod tests {
664 use super::*;
665
666 #[test]
669 fn strip_markers_from_chapter_removes_hidden_block() {
670 let content = r#"Some text
671
672```sql validator=sqlite hidden
673SELECT 1;
674```
675
676More text"#;
677 let result = ValidatorPreprocessor::strip_markers_from_chapter(content);
678 assert!(!result.contains("SELECT 1"));
680 assert!(!result.contains("```sql"));
681 assert!(result.contains("Some text"));
682 assert!(result.contains("More text"));
683 }
684
685 #[test]
686 fn strip_markers_from_chapter_keeps_non_hidden_block() {
687 let content = r#"Some text
688
689```sql validator=sqlite
690SELECT 1;
691```
692
693More text"#;
694 let result = ValidatorPreprocessor::strip_markers_from_chapter(content);
695 assert!(result.contains("SELECT 1"));
697 assert!(result.contains("```sql"));
698 assert!(result.contains("Some text"));
699 assert!(result.contains("More text"));
700 }
701
702 #[test]
703 fn strip_markers_from_chapter_mixed_hidden_and_non_hidden() {
704 let content = r#"Start
705
706```sql validator=sqlite hidden
707HIDDEN QUERY;
708```
709
710Middle
711
712```sql validator=sqlite
713VISIBLE QUERY;
714```
715
716End"#;
717 let result = ValidatorPreprocessor::strip_markers_from_chapter(content);
718 assert!(!result.contains("HIDDEN QUERY"));
720 assert!(result.contains("VISIBLE QUERY"));
721 assert!(result.contains("Start"));
722 assert!(result.contains("Middle"));
723 assert!(result.contains("End"));
724 }
725
726 #[test]
727 fn strip_markers_from_chapter_adjacent_hidden_blocks() {
728 let content = r#"Start
729
730```sql validator=sqlite hidden
731HIDDEN 1;
732```
733
734```sql validator=sqlite hidden
735HIDDEN 2;
736```
737
738End"#;
739 let result = ValidatorPreprocessor::strip_markers_from_chapter(content);
740 assert!(!result.contains("HIDDEN 1"));
742 assert!(!result.contains("HIDDEN 2"));
743 assert!(result.contains("Start"));
744 assert!(result.contains("End"));
745 }
746
747 #[test]
748 fn strip_markers_from_chapter_hidden_block_at_start() {
749 let content = r#"```sql validator=sqlite hidden
750HIDDEN;
751```
752
753Visible content"#;
754 let result = ValidatorPreprocessor::strip_markers_from_chapter(content);
755 assert!(!result.contains("HIDDEN"));
757 assert!(result.contains("Visible content"));
758 assert!(!result.starts_with('\n'));
760 }
761
762 #[test]
763 fn strip_markers_from_chapter_hidden_block_at_end() {
764 let content = r#"Visible content
765
766```sql validator=sqlite hidden
767HIDDEN;
768```"#;
769 let result = ValidatorPreprocessor::strip_markers_from_chapter(content);
770 assert!(!result.contains("HIDDEN"));
772 assert!(result.contains("Visible content"));
773 assert!(!result.ends_with("\n\n"));
775 }
776
777 #[test]
778 fn strip_markers_from_chapter_only_hidden_block() {
779 let content = r#"```sql validator=sqlite hidden
780HIDDEN;
781```"#;
782 let result = ValidatorPreprocessor::strip_markers_from_chapter(content);
783 assert!(!result.contains("HIDDEN"));
785 assert!(result.is_empty() || result.trim().is_empty());
786 }
787
788 #[test]
789 fn strip_markers_from_chapter_hidden_with_markers() {
790 let content = r#"Text
791
792```sql validator=sqlite hidden
793<!--SETUP
794CREATE TABLE t;
795-->
796SELECT * FROM t;
797<!--ASSERT
798rows >= 1
799-->
800```
801
802More text"#;
803 let result = ValidatorPreprocessor::strip_markers_from_chapter(content);
804 assert!(!result.contains("SETUP"));
806 assert!(!result.contains("ASSERT"));
807 assert!(!result.contains("CREATE TABLE"));
808 assert!(!result.contains("SELECT"));
809 assert!(result.contains("Text"));
810 assert!(result.contains("More text"));
811 }
812}