use std::collections::HashSet;
use std::fs::{self, OpenOptions};
use std::io::Write;
use std::path::{Path, PathBuf};
use futures_util::stream::{FuturesUnordered, StreamExt};
use globset::{Glob, GlobSet, GlobSetBuilder};
use regex::Regex;
use serde::{Deserialize, Serialize};
use tokio::sync::mpsc;
use url::Url;
use crate::archive::{
archive_id_from_seed, ensure_archive_size_can_grow, new_manifest, read_manifest, stable_id,
validate_archive_id, write_checksums, write_manifest, ArchiveLayout, ArchiveManifest,
CrawlStats,
};
use crate::config::Config;
use crate::errors::{Result, SiteforgeError};
use crate::fetch::{bytes_to_string, is_html, FetchResult, HttpFetcher};
use crate::ocr::{backend_from_config, likely_contains_code, OcrBackend};
use crate::parse::{normalize_link, parse_html, AssetRef, ImageDimensions, ParsedPage};
use crate::render::{render_js_dom, BrowserRenderOptions};
use crate::storage::{FetchedAssetRecord, FrontierItem, Storage};
#[derive(Debug, Clone)]
pub struct ArchiveRunOptions {
pub seeds: Vec<Url>,
pub full_site: bool,
pub same_domain: bool,
pub max_depth: usize,
pub max_pages: usize,
pub concurrency: usize,
pub delay_ms: u64,
pub headers: Vec<(String, String)>,
pub cookie: Option<String>,
pub archive_id: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ArchiveSummary {
pub archive_id: String,
pub archive_path: PathBuf,
pub stats: CrawlStats,
}
pub type ArchiveProgressSender = mpsc::UnboundedSender<ArchiveProgress>;
#[derive(Debug, Clone)]
pub struct ArchiveProgress {
pub archive_id: String,
pub phase: ArchiveProgressPhase,
pub active_url: Option<String>,
pub stats: CrawlStats,
pub max_pages: usize,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ArchiveProgressPhase {
Starting,
Discovering,
Fetching,
Processing,
Completed,
Interrupted,
}
impl ArchiveProgressPhase {
pub fn as_str(self) -> &'static str {
match self {
ArchiveProgressPhase::Starting => "starting",
ArchiveProgressPhase::Discovering => "discovering",
ArchiveProgressPhase::Fetching => "fetching",
ArchiveProgressPhase::Processing => "processing",
ArchiveProgressPhase::Completed => "completed",
ArchiveProgressPhase::Interrupted => "interrupted",
}
}
}
#[derive(Debug, Clone)]
pub struct CrawlScope {
seeds: Vec<SeedScope>,
same_domain: bool,
max_depth: usize,
include: GlobSet,
exclude: GlobSet,
}
#[derive(Debug, Clone)]
struct SeedScope {
host: String,
path_prefix: String,
}
impl CrawlScope {
pub fn new(
seeds: &[Url],
same_domain: bool,
max_depth: usize,
config: &Config,
) -> Result<Self> {
Ok(Self {
seeds: seeds.iter().map(seed_scope).collect(),
same_domain,
max_depth,
include: globset_from_patterns(&config.include_url_patterns)?,
exclude: globset_from_patterns(&config.exclude_url_patterns)?,
})
}
pub fn allows(&self, url: &Url, depth: usize) -> bool {
if depth > self.max_depth || !matches!(url.scheme(), "http" | "https") {
return false;
}
let url_string = url.as_str();
if self.exclude.is_match(url_string) {
return false;
}
if !self.include.is_empty() && !self.include.is_match(url_string) {
return false;
}
if !self.same_domain {
return true;
}
let Some(host) = url.host_str() else {
return false;
};
let host = host.trim_start_matches("www.");
self.seeds.iter().any(|seed| {
seed.host == host
&& (url.path().starts_with(&seed.path_prefix) || seed.path_prefix == "/")
})
}
}
pub async fn run_archive(config: Config, options: ArchiveRunOptions) -> Result<ArchiveSummary> {
run_archive_inner(config, options, None).await
}
pub async fn run_archive_with_progress(
config: Config,
options: ArchiveRunOptions,
progress: ArchiveProgressSender,
) -> Result<ArchiveSummary> {
run_archive_inner(config, options, Some(progress)).await
}
async fn run_archive_inner(
config: Config,
options: ArchiveRunOptions,
progress: Option<ArchiveProgressSender>,
) -> Result<ArchiveSummary> {
if options.seeds.is_empty() {
return Err(SiteforgeError::MissingInput("URL or --input"));
}
let root = config.resolved_archive_root()?;
let archive_id = match options.archive_id.as_deref() {
Some(archive_id) => {
validate_archive_id(archive_id)?;
archive_id.to_string()
}
None => archive_id_from_seed(&options.seeds[0]),
};
let layout = ArchiveLayout::new(root, archive_id.clone());
if layout.base().exists() {
return Err(SiteforgeError::message(format!(
"archive {archive_id:?} already exists; use `siteforge resume {archive_id}` or choose a different --archive-id"
)));
}
layout.ensure()?;
let mut manifest = new_manifest(
archive_id.clone(),
&options.seeds,
options.full_site,
options.same_domain,
options.max_depth,
&config,
);
manifest.config.concurrency = options.concurrency;
manifest.config.delay_ms = options.delay_ms;
manifest.config.max_pages = options.max_pages;
write_manifest(&layout, &manifest)?;
let storage = Storage::open_compact(&layout.database())?;
storage.create_or_update_job(&archive_id, &archive_id, "running")?;
storage.log("info", "archive job started")?;
for seed in &options.seeds {
storage.enqueue_url(seed.as_str(), 0, None)?;
}
fs::write(layout.chunks_jsonl(), "")?;
emit_progress(
progress.as_ref(),
&manifest,
&storage,
options.max_pages,
ArchiveProgressPhase::Starting,
None,
);
run_loop(
&config,
&layout,
&storage,
&mut manifest,
&options,
progress.as_ref(),
)
.await?;
let stats = storage.crawl_stats()?;
storage.update_job(&archive_id, "completed", None)?;
storage.log("info", "archive job completed")?;
manifest.stats = stats.clone();
write_manifest(&layout, &manifest)?;
write_readable_index(&config, &layout, &storage, &manifest)?;
storage.checkpoint()?;
write_checksums(&layout)?;
emit_progress(
progress.as_ref(),
&manifest,
&storage,
options.max_pages,
ArchiveProgressPhase::Completed,
None,
);
Ok(ArchiveSummary {
archive_id,
archive_path: layout.base(),
stats,
})
}
pub async fn resume_archive(config: Config, archive_id: &str) -> Result<ArchiveSummary> {
resume_archive_inner(config, archive_id, None).await
}
pub async fn resume_archive_with_progress(
config: Config,
archive_id: &str,
progress: ArchiveProgressSender,
) -> Result<ArchiveSummary> {
resume_archive_inner(config, archive_id, Some(progress)).await
}
async fn resume_archive_inner(
config: Config,
archive_id: &str,
progress: Option<ArchiveProgressSender>,
) -> Result<ArchiveSummary> {
let root = config.resolved_archive_root()?;
let layout = ArchiveLayout::new(root, archive_id.to_string());
if !layout.manifest().exists() {
return Err(SiteforgeError::ArchiveNotFound(archive_id.to_string()));
}
let mut manifest = read_manifest(&layout.manifest())?;
let seeds = manifest
.seeds
.iter()
.map(|seed| Url::parse(seed))
.collect::<std::result::Result<Vec<_>, _>>()?;
let options = ArchiveRunOptions {
seeds,
full_site: manifest.scope.full_site,
same_domain: manifest.scope.same_domain,
max_depth: manifest.scope.max_depth,
max_pages: manifest.config.max_pages,
concurrency: manifest.config.concurrency.max(1),
delay_ms: manifest.config.delay_ms,
headers: Vec::new(),
cookie: None,
archive_id: Some(archive_id.to_string()),
};
let storage = Storage::open_compact(&layout.database())?;
storage.reset_in_progress_frontier()?;
storage.create_or_update_job(archive_id, archive_id, "running")?;
storage.log("info", "archive job resumed")?;
emit_progress(
progress.as_ref(),
&manifest,
&storage,
options.max_pages,
ArchiveProgressPhase::Starting,
None,
);
run_loop(
&config,
&layout,
&storage,
&mut manifest,
&options,
progress.as_ref(),
)
.await?;
let stats = storage.crawl_stats()?;
storage.update_job(archive_id, "completed", None)?;
manifest.stats = stats.clone();
write_manifest(&layout, &manifest)?;
write_readable_index(&config, &layout, &storage, &manifest)?;
storage.checkpoint()?;
write_checksums(&layout)?;
emit_progress(
progress.as_ref(),
&manifest,
&storage,
options.max_pages,
ArchiveProgressPhase::Completed,
None,
);
Ok(ArchiveSummary {
archive_id: archive_id.to_string(),
archive_path: layout.base(),
stats,
})
}
async fn run_loop(
config: &Config,
layout: &ArchiveLayout,
storage: &Storage,
manifest: &mut ArchiveManifest,
options: &ArchiveRunOptions,
progress: Option<&ArchiveProgressSender>,
) -> Result<()> {
let fetcher = HttpFetcher::new(
config,
options.delay_ms,
&options.headers,
options.cookie.as_deref(),
)?;
let scope = CrawlScope::new(
&options.seeds,
options.same_domain,
options.max_depth,
config,
)?;
let ocr_backend = backend_from_config(config);
if options.full_site {
emit_progress(
progress,
manifest,
storage,
options.max_pages,
ArchiveProgressPhase::Discovering,
None,
);
enqueue_sitemap_urls(storage, &fetcher, &scope, &options.seeds, options.max_pages).await?;
}
let mut interrupt = Box::pin(tokio::signal::ctrl_c());
loop {
let stats = storage.crawl_stats()?;
if stats.pages_fetched >= options.max_pages {
storage.log("warn", "max page limit reached")?;
break;
}
let remaining = options.max_pages.saturating_sub(stats.pages_fetched).max(1);
let mut pending = storage.pending_frontier(options.concurrency.min(remaining).max(1))?;
if pending.is_empty() {
break;
}
emit_progress(
progress,
manifest,
storage,
options.max_pages,
ArchiveProgressPhase::Fetching,
None,
);
for item in &pending {
storage.mark_frontier(&item.url, "in_progress", None)?;
}
let mut futures = FuturesUnordered::new();
for item in pending.drain(..) {
let fetcher = fetcher.clone();
futures.push(async move {
let result = match Url::parse(&item.url) {
Ok(url) => fetcher.fetch(url).await,
Err(err) => Err(SiteforgeError::from(err)),
};
(item, result)
});
}
loop {
let next = tokio::select! {
biased;
_ = &mut interrupt => {
finish_interrupted(config, storage, layout, manifest, options, progress)?;
return Err(SiteforgeError::message(format!(
"archive interrupted; resume with `siteforge resume {}`",
manifest.archive_id
)));
}
next = futures.next() => next,
};
let Some((item, result)) = next else {
break;
};
storage.update_job(&manifest.archive_id, "running", Some(&item.url))?;
emit_progress(
progress,
manifest,
storage,
options.max_pages,
ArchiveProgressPhase::Processing,
Some(&item.url),
);
match result {
Ok(fetch_result) => {
if let Err(err) = handle_page(
config,
layout,
storage,
manifest,
options,
&scope,
&fetcher,
ocr_backend
.as_ref()
.map(|backend| backend as &dyn OcrBackend),
item,
fetch_result,
)
.await
{
storage.log("error", &format!("page handling failed: {err}"))?;
}
}
Err(err @ SiteforgeError::RobotsDenied(_)) => {
storage.mark_frontier(&item.url, "skipped", Some(&err.to_string()))?;
storage.log("warn", &err.to_string())?;
}
Err(err) => {
let retries = storage.frontier_retry_count(&item.url)?;
if retries < config.retry_count {
storage.increment_frontier_retry(&item.url, &err.to_string())?;
} else {
let page_id = format!("page-{}", stable_id(&item.url));
storage.insert_failed_page(
&manifest.archive_id,
&page_id,
&item.url,
&item.url,
0,
item.depth,
&err.to_string(),
)?;
storage.mark_frontier(&item.url, "failed", Some(&err.to_string()))?;
storage.log("error", &format!("fetch failed for {}: {err}", item.url))?;
}
}
}
emit_progress(
progress,
manifest,
storage,
options.max_pages,
ArchiveProgressPhase::Fetching,
None,
);
}
manifest.stats = storage.crawl_stats()?;
write_manifest(layout, manifest)?;
}
storage.update_job(&manifest.archive_id, "running", None)?;
Ok(())
}
fn finish_interrupted(
config: &Config,
storage: &Storage,
layout: &ArchiveLayout,
manifest: &mut ArchiveManifest,
options: &ArchiveRunOptions,
progress: Option<&ArchiveProgressSender>,
) -> Result<()> {
storage.update_job(&manifest.archive_id, "interrupted", None)?;
storage.log("warn", "archive job interrupted by signal")?;
manifest.stats = storage.crawl_stats()?;
write_manifest(layout, manifest)?;
write_readable_index(config, layout, storage, manifest)?;
storage.checkpoint()?;
write_checksums(layout)?;
emit_progress(
progress,
manifest,
storage,
options.max_pages,
ArchiveProgressPhase::Interrupted,
None,
);
Ok(())
}
fn emit_progress(
progress: Option<&ArchiveProgressSender>,
manifest: &ArchiveManifest,
storage: &Storage,
max_pages: usize,
phase: ArchiveProgressPhase,
active_url: Option<&str>,
) {
let Some(progress) = progress else {
return;
};
let stats = storage
.crawl_stats()
.unwrap_or_else(|_| manifest.stats.clone());
let _ = progress.send(ArchiveProgress {
archive_id: manifest.archive_id.clone(),
phase,
active_url: active_url.map(str::to_string),
stats,
max_pages,
});
}
fn write_readable_index(
config: &Config,
layout: &ArchiveLayout,
storage: &Storage,
manifest: &ArchiveManifest,
) -> Result<()> {
let pages = storage.list_pages()?;
let assets = storage.list_assets()?;
let mut markdown = String::new();
markdown.push_str(&format!(
"# Siteforge Archive `{}`\n\n",
manifest.archive_id
));
markdown.push_str("This is the file-first index for humans and AI agents. The cleaned Markdown, structured JSON, text, chunks, and packs are the primary artifacts. SQLite is supporting crawl state/search metadata only.\n\n");
markdown.push_str("## Seeds\n\n");
for seed in &manifest.seeds {
markdown.push_str(&format!("- {seed}\n"));
}
markdown.push_str("\n## Stats\n\n");
markdown.push_str(&format!(
"- Pages discovered: {}\n- Pages parsed: {}\n- Pages failed: {}\n- Assets downloaded: {}\n\n",
manifest.stats.pages_discovered,
manifest.stats.pages_parsed,
manifest.stats.pages_failed,
manifest.stats.assets_downloaded
));
markdown.push_str("## AI-Readable Artifacts\n\n");
markdown
.push_str("- `readable/markdown/`: cleaned Markdown per page with source attribution.\n");
markdown.push_str("- `readable/basic_markdown/`: plain Markdown per page with no YAML frontmatter, plus page data, links, and assets.\n");
markdown.push_str("- `readable/json/`: structured page JSON.\n");
markdown.push_str("- `readable/text/`: plain text for quick search/read.\n");
markdown.push_str("- `chunks/chunks.jsonl`: semantic chunks for RAG/AI ingestion.\n");
markdown.push_str("- `packs/`: token-budgeted AI context packs.\n\n");
markdown.push_str("## Agent Entrypoints\n\n");
markdown.push_str("- `AGENTS.md`: start here for quick ingestion guidance.\n");
markdown.push_str("- `agent-index.json`: machine-readable page/asset index.\n\n");
markdown.push_str("## Pages\n\n");
for page in pages.iter().filter(|page| page.markdown_path.is_some()) {
let markdown_link = page
.markdown_path
.as_deref()
.unwrap_or_default()
.strip_prefix("readable/")
.unwrap_or(page.markdown_path.as_deref().unwrap_or_default());
let basic_link = format!("basic_markdown/{}.md", page.page_id);
markdown.push_str(&format!(
"- [{}]({}) - {}\n - Basic: [{}]({})\n - Source: {}\n",
page.title, markdown_link, page.page_id, page.title, basic_link, page.source_url
));
}
if !assets.is_empty() {
markdown.push_str("\n## Assets\n\n");
for asset in assets.iter().take(200) {
markdown.push_str(&format!(
"- {} - {}\n - Local: {}\n",
asset.asset_id,
asset.source_url,
asset.local_path.as_deref().unwrap_or("not downloaded")
));
}
if assets.len() > 200 {
markdown.push_str(&format!(
"\n_{} additional assets omitted from this index._\n",
assets.len() - 200
));
}
}
write_archive_file(
config,
layout,
&layout.readable_dir().join("README.md"),
markdown.as_bytes(),
)?;
write_agent_entrypoints(config, layout, manifest, &pages, &assets)
}
fn write_agent_entrypoints(
config: &Config,
layout: &ArchiveLayout,
manifest: &ArchiveManifest,
pages: &[crate::storage::StoredPage],
assets: &[crate::storage::StoredAsset],
) -> Result<()> {
#[derive(Serialize)]
struct AgentPathIndex<'a> {
agents_md: &'a str,
readable_readme: &'a str,
basic_markdown: &'a str,
markdown: &'a str,
json: &'a str,
text: &'a str,
chunks_jsonl: &'a str,
packs: &'a str,
raw_pages: &'a str,
rendered_pages: &'a str,
raw_assets: &'a str,
}
#[derive(Serialize)]
struct AgentPageIndex<'a> {
page_id: &'a str,
title: &'a str,
source_url: &'a str,
final_url: &'a str,
basic_markdown_path: Option<String>,
markdown_path: Option<&'a str>,
json_path: Option<&'a str>,
text_path: Option<&'a str>,
raw_path: Option<&'a str>,
status: i64,
}
#[derive(Serialize)]
struct AgentAssetIndex<'a> {
asset_id: &'a str,
page_id: &'a str,
source_url: &'a str,
local_path: Option<&'a str>,
mime_type: Option<&'a str>,
size_bytes: Option<u64>,
ocr_status: Option<&'a str>,
}
#[derive(Serialize)]
struct AgentIndex<'a> {
archive_id: &'a str,
updated_at: String,
seeds: &'a [String],
paths: AgentPathIndex<'a>,
pages: Vec<AgentPageIndex<'a>>,
assets: Vec<AgentAssetIndex<'a>>,
}
let mut agents = String::new();
agents.push_str(&format!("# AGENTS.md for `{}`\n\n", manifest.archive_id));
agents.push_str("This archive is optimized for direct agent ingestion. Prefer file artifacts over SQLite; the database is crawl/search state only.\n\n");
agents.push_str("## Fast Paths\n\n");
agents.push_str("- Start with `agent-index.json` for machine-readable page and asset paths.\n");
agents.push_str("- Use `readable/basic_markdown/` for simple Markdown with content, metadata, links, and assets.\n");
agents.push_str(
"- Use `readable/markdown/` when YAML frontmatter and source comments are useful.\n",
);
agents
.push_str("- Use `chunks/chunks.jsonl` for RAG ingestion; code chunks are never split.\n");
agents.push_str("- Use `packs/` for token-budgeted context packs.\n");
agents.push_str("- Use `raw/pages/`, `raw/rendered/`, and `raw/assets/` when exact source payloads or local assets are needed.\n\n");
agents.push_str("## Suggested Agent Order\n\n");
agents.push_str("1. Read `agent-index.json`.\n");
agents.push_str("2. Load `readable/basic_markdown/*.md` for broad understanding.\n");
agents.push_str("3. Load `chunks/chunks.jsonl` for precise retrieval.\n");
agents.push_str("4. Open referenced raw/assets files only when needed.\n\n");
agents.push_str("## Archive Stats\n\n");
agents.push_str(&format!(
"- Pages parsed: {}\n- Pages failed: {}\n- Assets downloaded: {}\n",
manifest.stats.pages_parsed, manifest.stats.pages_failed, manifest.stats.assets_downloaded
));
let readable_agents = layout.readable_dir().join("AGENTS.md");
write_archive_file(config, layout, &layout.agents_md(), agents.as_bytes())?;
write_archive_file(config, layout, &readable_agents, agents.as_bytes())?;
let index = AgentIndex {
archive_id: &manifest.archive_id,
updated_at: manifest.updated_at.to_rfc3339(),
seeds: &manifest.seeds,
paths: AgentPathIndex {
agents_md: "AGENTS.md",
readable_readme: "readable/README.md",
basic_markdown: "readable/basic_markdown/",
markdown: "readable/markdown/",
json: "readable/json/",
text: "readable/text/",
chunks_jsonl: "chunks/chunks.jsonl",
packs: "packs/",
raw_pages: "raw/pages/",
rendered_pages: "raw/rendered/",
raw_assets: "raw/assets/",
},
pages: pages
.iter()
.map(|page| AgentPageIndex {
page_id: &page.page_id,
title: &page.title,
source_url: &page.source_url,
final_url: &page.final_url,
basic_markdown_path: page
.markdown_path
.as_ref()
.map(|_| format!("readable/basic_markdown/{}.md", page.page_id)),
markdown_path: page.markdown_path.as_deref(),
json_path: page.json_path.as_deref(),
text_path: page.text_path.as_deref(),
raw_path: page.raw_path.as_deref(),
status: page.status,
})
.collect(),
assets: assets
.iter()
.map(|asset| AgentAssetIndex {
asset_id: &asset.asset_id,
page_id: &asset.page_id,
source_url: &asset.source_url,
local_path: asset.local_path.as_deref(),
mime_type: asset.mime_type.as_deref(),
size_bytes: asset.size_bytes,
ocr_status: asset.ocr_status.as_deref(),
})
.collect(),
};
write_archive_file(
config,
layout,
&layout.agent_index_json(),
serde_json::to_string_pretty(&index)?.as_bytes(),
)
}
#[allow(clippy::too_many_arguments)]
async fn handle_page(
config: &Config,
layout: &ArchiveLayout,
storage: &Storage,
manifest: &ArchiveManifest,
options: &ArchiveRunOptions,
scope: &CrawlScope,
fetcher: &HttpFetcher,
ocr_backend: Option<&dyn OcrBackend>,
item: FrontierItem,
result: FetchResult,
) -> Result<()> {
let page_id = format!("page-{}", stable_id(&item.url));
if result.status >= 400 {
storage.insert_failed_page(
&manifest.archive_id,
&page_id,
&item.url,
result.final_url.as_str(),
result.status as i64,
item.depth,
&format!("HTTP {}", result.status),
)?;
storage.mark_frontier(
&item.url,
"failed",
Some(&format!("HTTP {}", result.status)),
)?;
return Ok(());
}
if !is_html(&result) {
storage.insert_failed_page(
&manifest.archive_id,
&page_id,
&item.url,
result.final_url.as_str(),
result.status as i64,
item.depth,
"unsupported non-HTML page response",
)?;
storage.mark_frontier(
&item.url,
"skipped",
Some("unsupported non-HTML page response"),
)?;
return Ok(());
}
let raw_path = layout.raw_pages_dir().join(format!("{page_id}.html"));
write_archive_file(config, layout, &raw_path, &result.bytes)?;
let mut html = bytes_to_string(&result.bytes);
let mut parse_hash = result.content_hash.clone();
if config.render_js {
match render_js_dom(
BrowserRenderOptions::from_config(config),
result.final_url.clone(),
)
.await
{
Ok(rendered_html) => {
parse_hash = blake3::hash(rendered_html.as_bytes()).to_hex().to_string();
let rendered_path = layout.rendered_pages_dir().join(format!("{page_id}.html"));
write_archive_file(config, layout, &rendered_path, rendered_html.as_bytes())?;
html = rendered_html;
storage.log(
"info",
&format!("captured rendered DOM for {}", result.final_url),
)?;
}
Err(err) => {
storage.log(
"warn",
&format!(
"JS render failed for {}: {err}; using fetched HTML",
result.final_url
),
)?;
}
}
}
let mut parsed = parse_html(
&manifest.archive_id,
&page_id,
&result.final_url,
&html,
&parse_hash,
item.depth,
);
add_downloadable_link_assets(&mut parsed);
download_assets(config, layout, storage, fetcher, ocr_backend, &mut parsed).await?;
let markdown = parsed.to_markdown()?;
let basic_markdown = parsed.to_basic_markdown();
let json = serde_json::to_string_pretty(&parsed.to_page_json())?;
let text = parsed.to_plain_text();
let markdown_path = layout.readable_markdown_dir().join(format!("{page_id}.md"));
let basic_markdown_path = layout
.readable_basic_markdown_dir()
.join(format!("{page_id}.md"));
let json_path = layout.readable_json_dir().join(format!("{page_id}.json"));
let text_path = layout.readable_text_dir().join(format!("{page_id}.txt"));
write_archive_file(config, layout, &markdown_path, markdown.as_bytes())?;
write_archive_file(
config,
layout,
&basic_markdown_path,
basic_markdown.as_bytes(),
)?;
write_archive_file(config, layout, &json_path, json.as_bytes())?;
write_archive_file(config, layout, &text_path, text.as_bytes())?;
append_chunks(config, layout, &parsed)?;
let raw_rel = relative_to_base(layout, &raw_path);
let markdown_rel = relative_to_base(layout, &markdown_path);
let json_rel = relative_to_base(layout, &json_path);
let text_rel = relative_to_base(layout, &text_path);
let headers_json = serde_json::to_string(&result.headers)?;
storage.upsert_page(
&parsed,
result.final_url.as_str(),
result.status as i64,
result.mime_type.as_deref(),
Some(&raw_rel),
Some(&markdown_rel),
Some(&json_rel),
Some(&text_rel),
Some(&headers_json),
None,
)?;
storage.mark_frontier(&item.url, "done", None)?;
storage.log("info", &format!("archived {}", result.final_url))?;
if options.full_site && item.depth < options.max_depth {
enqueue_discovered_links(storage, scope, &parsed, item.depth + 1)?;
}
Ok(())
}
async fn download_assets(
config: &Config,
layout: &ArchiveLayout,
storage: &Storage,
fetcher: &HttpFetcher,
ocr_backend: Option<&dyn OcrBackend>,
parsed: &mut ParsedPage,
) -> Result<()> {
let mut seen = HashSet::new();
for index in 0..parsed.assets.len() {
let asset_id = parsed.assets[index].asset_id.clone();
if !seen.insert(asset_id.clone()) {
continue;
}
let url = match Url::parse(&parsed.assets[index].url) {
Ok(url) => url,
Err(err) => {
parsed.assets[index].ocr_status = Some(format!("asset_fetch_failed: {err}"));
storage.upsert_asset(&parsed.page_id, &parsed.assets[index], None)?;
continue;
}
};
match fetcher.fetch(url).await {
Ok(result) => {
if result.bytes.len() as u64 > config.max_asset_size_bytes {
parsed.assets[index].ocr_status = Some("asset_fetch_failed".to_string());
storage.upsert_asset(&parsed.page_id, &parsed.assets[index], None)?;
continue;
}
let extension = extension_for(&result.final_url, result.mime_type.as_deref());
let file_name = format!("{asset_id}.{extension}");
let asset_path = layout.raw_assets_dir().join(file_name);
write_archive_file(config, layout, &asset_path, &result.bytes)?;
let rel = relative_to_base(layout, &asset_path);
let markdown_rel = format!("../../{rel}");
parsed.assets[index].local_path = Some(markdown_rel);
parsed.assets[index].hash = Some(result.content_hash.clone());
parsed.assets[index].dimensions =
image_dimensions(&result.bytes, result.mime_type.as_deref());
let source_url = parsed.assets[index].url.clone();
if is_source_text_asset(&result.final_url, result.mime_type.as_deref()) {
let text = bytes_to_string(&result.bytes);
if !text.trim().is_empty() {
parsed.add_source_asset_code(
&asset_id,
&source_url,
text,
language_for_source_asset(
&result.final_url,
result.mime_type.as_deref(),
),
);
}
}
let fetched = FetchedAssetRecord {
final_url: result.final_url.to_string(),
mime_type: result.mime_type.clone(),
size_bytes: result.bytes.len() as u64,
content_hash: Some(result.content_hash.clone()),
};
if let Some(backend) = ocr_backend {
if likely_contains_code(&parsed.assets[index]) {
match backend.extract(&asset_path) {
Ok(output) if !output.text.trim().is_empty() => {
parsed.set_asset_ocr(
&asset_id,
"ok".to_string(),
Some(output.text),
);
}
Ok(_) => {
parsed.set_asset_ocr(&asset_id, "empty".to_string(), None);
}
Err(SiteforgeError::OcrUnavailable(message)) => {
parsed.set_asset_ocr(&asset_id, "unavailable".to_string(), None);
storage.log("warn", &format!("OCR unavailable: {message}"))?;
}
Err(err) => {
parsed.set_asset_ocr(&asset_id, format!("failed: {err}"), None);
}
}
}
}
if let Some(asset) = parsed
.assets
.iter()
.find(|asset| asset.asset_id == asset_id)
{
storage.upsert_asset(&parsed.page_id, asset, Some(&fetched))?;
}
}
Err(err) => {
parsed.assets[index].ocr_status = Some("asset_fetch_failed".to_string());
storage.upsert_asset(&parsed.page_id, &parsed.assets[index], None)?;
storage.log(
"warn",
&format!("asset fetch failed for {}: {err}", parsed.assets[index].url),
)?;
}
}
}
Ok(())
}
fn enqueue_discovered_links(
storage: &Storage,
scope: &CrawlScope,
parsed: &ParsedPage,
depth: usize,
) -> Result<()> {
for link in &parsed.links {
let Ok(url) = Url::parse(&link.url) else {
continue;
};
if scope.allows(&url, depth) {
storage.enqueue_url(url.as_str(), depth, Some(&parsed.source_url))?;
}
}
Ok(())
}
async fn enqueue_sitemap_urls(
storage: &Storage,
fetcher: &HttpFetcher,
scope: &CrawlScope,
seeds: &[Url],
max_pages: usize,
) -> Result<()> {
let mut sitemap_urls = Vec::new();
for seed in seeds {
match fetcher.robots_sitemaps(seed).await {
Ok(urls) => sitemap_urls.extend(urls),
Err(err) => storage.log("warn", &format!("robots sitemap discovery failed: {err}"))?,
}
if let Ok(default_sitemap) = seed.join("/sitemap.xml") {
sitemap_urls.push(default_sitemap);
}
}
sitemap_urls.sort();
sitemap_urls.dedup();
let mut visited = HashSet::new();
let mut enqueued = 0usize;
for sitemap_url in sitemap_urls {
enqueue_sitemap_url(
storage,
fetcher,
scope,
sitemap_url,
&mut visited,
&mut enqueued,
max_pages,
0,
)
.await?;
if enqueued >= max_pages {
break;
}
}
if enqueued > 0 {
storage.log("info", &format!("queued {enqueued} URL(s) from sitemaps"))?;
}
Ok(())
}
#[allow(clippy::too_many_arguments)]
async fn enqueue_sitemap_url(
storage: &Storage,
fetcher: &HttpFetcher,
scope: &CrawlScope,
sitemap_url: Url,
visited: &mut HashSet<String>,
enqueued: &mut usize,
max_pages: usize,
sitemap_depth: usize,
) -> Result<()> {
const MAX_SITEMAP_DEPTH: usize = 3;
if sitemap_depth > MAX_SITEMAP_DEPTH
|| *enqueued >= max_pages
|| !visited.insert(sitemap_url.to_string())
{
return Ok(());
}
let result = match fetcher.fetch(sitemap_url.clone()).await {
Ok(result) => result,
Err(err) => {
storage.log(
"warn",
&format!("sitemap fetch failed for {sitemap_url}: {err}"),
)?;
return Ok(());
}
};
if result.status >= 400 {
storage.log(
"warn",
&format!(
"sitemap fetch returned HTTP {} for {sitemap_url}",
result.status
),
)?;
return Ok(());
}
let body = sitemap_body(
&result.bytes,
result.mime_type.as_deref(),
sitemap_url.path(),
)?;
for loc in sitemap_locs(&result.final_url, &body) {
if *enqueued >= max_pages {
break;
}
if is_sitemap_like(&loc) && sitemap_depth < MAX_SITEMAP_DEPTH {
Box::pin(enqueue_sitemap_url(
storage,
fetcher,
scope,
loc,
visited,
enqueued,
max_pages,
sitemap_depth + 1,
))
.await?;
} else if scope.allows(&loc, 1) && storage.enqueue_url(loc.as_str(), 1, Some("sitemap"))? {
*enqueued += 1;
}
}
Ok(())
}
fn sitemap_body(bytes: &[u8], mime_type: Option<&str>, path: &str) -> Result<String> {
if path.ends_with(".gz")
|| mime_type == Some("application/x-gzip")
|| mime_type == Some("application/gzip")
{
let mut decoder = flate2::read::GzDecoder::new(bytes);
let mut body = String::new();
std::io::Read::read_to_string(&mut decoder, &mut body)?;
Ok(body)
} else {
Ok(bytes_to_string(bytes))
}
}
fn sitemap_locs(base: &Url, body: &str) -> Vec<Url> {
let regex = Regex::new(r"(?is)<loc>\s*([^<]+)\s*</loc>").expect("static sitemap regex");
let mut urls = regex
.captures_iter(body)
.filter_map(|capture| capture.get(1).map(|value| value.as_str()))
.map(|value| html_escape::decode_html_entities(value).trim().to_string())
.filter_map(|value| Url::parse(&value).or_else(|_| base.join(&value)).ok())
.collect::<Vec<_>>();
urls.sort();
urls.dedup();
urls
}
fn is_sitemap_like(url: &Url) -> bool {
let path = url.path().to_ascii_lowercase();
path.ends_with(".xml") || path.ends_with(".xml.gz") || path.contains("sitemap")
}
fn add_downloadable_link_assets(parsed: &mut ParsedPage) {
let mut existing = parsed
.assets
.iter()
.map(|asset| asset.url.clone())
.collect::<HashSet<_>>();
for link in &parsed.links {
if !is_downloadable_asset(&link.url) || !existing.insert(link.url.clone()) {
continue;
}
let asset_id = format!("asset-{}", stable_id(&link.url));
parsed.assets.push(AssetRef {
asset_id,
url: link.url.clone(),
local_path: None,
alt_text: Some(link.text.clone()).filter(|text| !text.trim().is_empty()),
caption: None,
dimensions: None,
hash: None,
appears_in: vec!["download-link".to_string()],
ocr_text: None,
ocr_status: None,
});
}
}
fn is_downloadable_asset(url: &str) -> bool {
let lower = url.to_ascii_lowercase();
[
".zip",
".tar",
".gz",
".pdf",
".cs",
".rs",
".py",
".js",
".mjs",
".jsx",
".ts",
".tsx",
".css",
".scss",
".sass",
".less",
".html",
".htm",
".md",
".txt",
".csv",
".json",
".jsonl",
".xml",
".svg",
".ipynb",
".wasm",
".shader",
".unitypackage",
]
.iter()
.any(|extension| {
lower
.split('?')
.next()
.unwrap_or(&lower)
.ends_with(extension)
})
}
fn append_chunks(config: &Config, layout: &ArchiveLayout, parsed: &ParsedPage) -> Result<()> {
let mut payload = String::new();
for chunk in parsed.chunks() {
payload.push_str(&serde_json::to_string(&chunk)?);
payload.push('\n');
}
ensure_archive_size_can_grow(config, layout, payload.len() as u64, &layout.chunks_jsonl())?;
let mut file = OpenOptions::new()
.create(true)
.append(true)
.open(layout.chunks_jsonl())?;
file.write_all(payload.as_bytes())?;
Ok(())
}
fn write_archive_file(
config: &Config,
layout: &ArchiveLayout,
path: &Path,
bytes: &[u8],
) -> Result<()> {
let existing = fs::metadata(path)
.map(|metadata| metadata.len())
.unwrap_or(0);
let growth = (bytes.len() as u64).saturating_sub(existing);
ensure_archive_size_can_grow(config, layout, growth, path)?;
fs::write(path, bytes)?;
Ok(())
}
fn relative_to_base(layout: &ArchiveLayout, path: &Path) -> String {
path.strip_prefix(layout.base())
.unwrap_or(path)
.to_string_lossy()
.replace('\\', "/")
}
fn extension_for(url: &Url, mime_type: Option<&str>) -> String {
if let Some(extension) = Path::new(url.path())
.extension()
.and_then(|ext| ext.to_str())
{
let clean = extension
.chars()
.filter(|ch| ch.is_ascii_alphanumeric())
.collect::<String>();
if !clean.is_empty() && clean.len() <= 8 {
return clean;
}
}
match mime_type.unwrap_or_default() {
"image/png" => "png",
"image/jpeg" => "jpg",
"image/gif" => "gif",
"image/webp" => "webp",
"image/svg+xml" => "svg",
"application/pdf" => "pdf",
"application/zip" => "zip",
"text/css" => "css",
"application/javascript" | "text/javascript" => "js",
"application/json" => "json",
"text/plain" => "txt",
_ => "bin",
}
.to_string()
}
fn image_dimensions(bytes: &[u8], mime_type: Option<&str>) -> Option<ImageDimensions> {
if !mime_type
.map(|mime| mime.starts_with("image/"))
.unwrap_or(false)
{
return None;
}
imagesize::blob_size(bytes)
.ok()
.map(|size| ImageDimensions {
width: size.width as u32,
height: size.height as u32,
})
}
fn is_source_text_asset(url: &Url, mime_type: Option<&str>) -> bool {
if mime_type
.map(|mime| {
mime.starts_with("text/")
|| matches!(
mime,
"application/json"
| "application/javascript"
| "application/x-javascript"
| "application/xml"
| "application/toml"
| "application/yaml"
)
})
.unwrap_or(false)
{
return true;
}
source_extension(url).is_some()
}
fn language_for_source_asset(url: &Url, mime_type: Option<&str>) -> Option<String> {
source_extension(url)
.map(|extension| match extension.as_str() {
"cs" => "csharp".to_string(),
"rs" => "rust".to_string(),
"py" => "python".to_string(),
"js" => "javascript".to_string(),
"mjs" => "javascript".to_string(),
"jsx" => "javascript".to_string(),
"ts" => "typescript".to_string(),
"tsx" => "typescript".to_string(),
"yml" => "yaml".to_string(),
"md" => "markdown".to_string(),
other => other.to_string(),
})
.or_else(|| {
mime_type.and_then(|mime| match mime {
"application/json" => Some("json".to_string()),
"application/javascript" | "application/x-javascript" | "text/javascript" => {
Some("javascript".to_string())
}
"application/xml" | "text/xml" => Some("xml".to_string()),
"text/css" => Some("css".to_string()),
_ => None,
})
})
}
fn source_extension(url: &Url) -> Option<String> {
Path::new(url.path())
.extension()
.and_then(|extension| extension.to_str())
.map(|extension| extension.to_ascii_lowercase())
.and_then(|extension| {
matches!(
extension.as_str(),
"cs" | "rs"
| "py"
| "js"
| "mjs"
| "jsx"
| "ts"
| "tsx"
| "css"
| "scss"
| "sass"
| "less"
| "html"
| "htm"
| "md"
| "markdown"
| "txt"
| "csv"
| "json"
| "jsonl"
| "toml"
| "yaml"
| "yml"
| "xml"
| "svg"
| "sh"
| "bash"
| "ipynb"
| "wasm"
| "shader"
| "cginc"
| "hlsl"
| "glsl"
)
.then_some(extension)
})
}
fn seed_scope(seed: &Url) -> SeedScope {
let host = seed
.host_str()
.unwrap_or_default()
.trim_start_matches("www.")
.to_string();
let path = seed.path();
let path_prefix = if path.is_empty() || path == "/" {
"/".to_string()
} else if path.ends_with('/') {
path.to_string()
} else {
path.rsplit_once('/')
.map(|(prefix, _)| format!("{prefix}/"))
.unwrap_or_else(|| "/".to_string())
};
SeedScope { host, path_prefix }
}
fn globset_from_patterns(patterns: &[String]) -> Result<GlobSet> {
let mut builder = GlobSetBuilder::new();
for pattern in patterns {
builder.add(Glob::new(pattern)?);
}
Ok(builder.build()?)
}
pub fn normalize_url(base: &Url, raw: &str) -> Option<Url> {
normalize_link(base, raw)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn scope_respects_domain_and_prefix() {
let config = Config::default();
let seed = Url::parse("https://catlikecoding.com/unity/tutorials/").unwrap();
let scope = CrawlScope::new(&[seed], true, 2, &config).unwrap();
assert!(scope.allows(
&Url::parse("https://catlikecoding.com/unity/tutorials/mesh").unwrap(),
1
));
assert!(!scope.allows(&Url::parse("https://catlikecoding.com/other/").unwrap(), 1));
assert!(!scope.allows(
&Url::parse("https://example.com/unity/tutorials/").unwrap(),
1
));
}
#[test]
fn scope_respects_depth() {
let config = Config::default();
let seed = Url::parse("https://example.com/").unwrap();
let scope = CrawlScope::new(&[seed], true, 1, &config).unwrap();
assert!(scope.allows(&Url::parse("https://example.com/a").unwrap(), 1));
assert!(!scope.allows(&Url::parse("https://example.com/b").unwrap(), 2));
}
#[test]
fn captures_image_dimensions_from_bytes() {
let png_1x1 = [
137, 80, 78, 71, 13, 10, 26, 10, 0, 0, 0, 13, 73, 72, 68, 82, 0, 0, 0, 1, 0, 0, 0, 1,
8, 6, 0, 0, 0, 31, 21, 196, 137, 0, 0, 0, 13, 73, 68, 65, 84, 120, 156, 99, 248, 15, 4,
0, 9, 251, 3, 253, 167, 161, 176, 219, 0, 0, 0, 0, 73, 69, 78, 68, 174, 66, 96, 130,
];
let dimensions = image_dimensions(&png_1x1, Some("image/png")).unwrap();
assert_eq!(dimensions.width, 1);
assert_eq!(dimensions.height, 1);
}
#[test]
fn extracts_urls_from_sitemap_xml() {
let base = Url::parse("https://example.test/sitemap.xml").unwrap();
let urls = sitemap_locs(
&base,
r#"<urlset><url><loc>https://example.test/a</loc></url><url><loc>/b</loc></url></urlset>"#,
);
assert_eq!(urls.len(), 2);
assert!(urls
.iter()
.any(|url| url.as_str() == "https://example.test/a"));
assert!(urls
.iter()
.any(|url| url.as_str() == "https://example.test/b"));
}
#[test]
fn reads_gzip_sitemap_body() {
use flate2::{write::GzEncoder, Compression};
use std::io::Write;
let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
encoder.write_all(b"<urlset></urlset>").unwrap();
let bytes = encoder.finish().unwrap();
let body = sitemap_body(&bytes, Some("application/gzip"), "/sitemap.xml.gz").unwrap();
assert_eq!(body, "<urlset></urlset>");
}
}