1use crate::auth::{AuthProvider, Credential};
22use crate::error::{Error, Result};
23use crate::fetch::{FetchOptions, Fetcher};
24use crate::payload::DocPayload;
25use crate::resource::{Page, Resource, ResourceKind, SourceRef};
26use crate::source::Source;
27use crate::url_key::UrlKey;
28use crate::url_key::UrlKeyExt as _;
29use compact_str::CompactString;
30use futures::stream::{FuturesUnordered, StreamExt};
31use secrecy::{ExposeSecret, SecretString};
32use std::borrow::Cow;
33use std::collections::HashMap;
34use std::future::Ready;
35use url::Url;
36
37const DEFAULT_MAX_BYTES: u64 = 2 * 1024 * 1024;
39const DEFAULT_CONCURRENCY: usize = 8;
41
42#[derive(Clone, Debug, PartialEq, Eq)]
58pub struct GithubSpec {
59 pub owner: CompactString,
61 pub repo: CompactString,
63 pub ref_name: CompactString,
65 pub subdir: CompactString,
67}
68
69impl GithubSpec {
70 pub fn parse(s: &str) -> Result<Self> {
72 let s = s.trim();
73 if let Some(rest) = s.strip_prefix("github:") {
74 return Self::parse_ref_grammar(rest);
75 }
76 if s.starts_with("https://github.com/") || s.starts_with("http://github.com/") {
77 return Self::parse_url(s);
78 }
79 Err(Error::invalid_url(format!(
80 "not a GitHub spec: {s:?} (expected https://github.com/owner/repo[/tree/ref[/dir]] \
81 or github:owner/repo@ref[//dir])"
82 )))
83 }
84
85 #[must_use]
88 pub fn matches(s: &str) -> bool {
89 let s = s.trim();
90 s.starts_with("github:")
91 || s.starts_with("https://github.com/")
92 || s.starts_with("http://github.com/")
93 }
94
95 fn parse_ref_grammar(rest: &str) -> Result<Self> {
96 let bad = |m: &str| Error::invalid_url(format!("github:{rest}: {m}"));
99 let (locator, subdir) = match rest.split_once("//") {
100 Some((l, s)) => (l, s.trim_matches('/')),
101 None => (rest, ""),
102 };
103 let (path, ref_name) = match locator.split_once('@') {
104 Some((p, r)) if !r.is_empty() => (p, r),
105 Some(_) => return Err(bad("empty ref after '@'")),
106 None => (locator, ""),
107 };
108 let (owner, repo) = path
109 .split_once('/')
110 .ok_or_else(|| bad("expected owner/repo"))?;
111 if owner.is_empty() || repo.is_empty() || repo.contains('/') {
112 return Err(bad("expected exactly owner/repo"));
113 }
114 Ok(Self {
115 owner: owner.into(),
116 repo: repo.into(),
117 ref_name: ref_name.into(),
118 subdir: subdir.into(),
119 })
120 }
121
122 fn parse_url(s: &str) -> Result<Self> {
123 let url = Url::parse(s).map_err(|e| Error::invalid_url(format!("{s:?}: {e}")))?;
124 let bad = |m: &str| Error::invalid_url(format!("{s:?}: {m}"));
125 let mut segs = url
126 .path_segments()
127 .ok_or_else(|| bad("no path"))?
128 .filter(|p| !p.is_empty());
129 let owner = segs.next().ok_or_else(|| bad("missing owner"))?;
130 let repo_raw = segs.next().ok_or_else(|| bad("missing repo"))?;
131 let repo = repo_raw.strip_suffix(".git").unwrap_or(repo_raw);
132 if repo.is_empty() {
133 return Err(bad("missing repo"));
134 }
135 let (ref_name, subdir) = match segs.next() {
136 None => (String::new(), String::new()),
137 Some("tree" | "blob") => {
138 let r = segs.next().ok_or_else(|| bad("missing ref after /tree/"))?;
139 let rest: Vec<&str> = segs.collect();
140 (r.to_owned(), rest.join("/"))
141 }
142 Some(other) => {
143 return Err(bad(&format!(
144 "unexpected path segment {other:?} (only /tree/<ref>/<dir> is understood)"
145 )));
146 }
147 };
148 Ok(Self {
149 owner: owner.into(),
150 repo: repo.into(),
151 ref_name: ref_name.into(),
152 subdir: subdir.trim_matches('/').into(),
153 })
154 }
155}
156
157#[derive(Debug)]
165pub struct GithubAuth {
166 token: Option<SecretString>,
167 api_host: CompactString,
168 raw_host: CompactString,
169}
170
171impl GithubAuth {
172 #[must_use]
176 pub fn new(token: Option<String>) -> Self {
177 Self {
178 token: token.map(|t| SecretString::new(t.into_boxed_str())),
179 api_host: "api.github.com".into(),
180 raw_host: "raw.githubusercontent.com".into(),
181 }
182 }
183
184 #[must_use]
186 pub fn with_hosts(
187 mut self,
188 api_host: impl Into<CompactString>,
189 raw_host: impl Into<CompactString>,
190 ) -> Self {
191 self.api_host = api_host.into();
192 self.raw_host = raw_host.into();
193 self
194 }
195}
196
197impl AuthProvider for GithubAuth {
198 type RefreshFuture<'a> = Ready<Result<()>>;
199
200 fn credential(&self, resource: &Resource) -> Credential<'_> {
201 let Some(token) = &self.token else {
202 return Credential::Anonymous;
203 };
204 let trusted = resource
205 .url
206 .host_str()
207 .is_some_and(|h| h == self.api_host || h == self.raw_host);
208 if trusted {
209 Credential::Bearer(Cow::Borrowed(token.expose_secret()))
210 } else {
211 Credential::Anonymous
212 }
213 }
214
215 fn refresh(&self) -> Self::RefreshFuture<'_> {
216 std::future::ready(Ok(()))
217 }
218}
219
220#[derive(Debug, serde::Deserialize)]
222struct TreeEntry {
223 path: String,
224 #[serde(rename = "type")]
225 kind: String,
226 sha: String,
227 #[serde(default)]
228 size: Option<u64>,
229}
230
231#[derive(Debug, serde::Deserialize)]
232struct TreeResponse {
233 #[serde(default)]
234 truncated: bool,
235 tree: Vec<TreeEntry>,
236}
237
238#[derive(Debug, serde::Deserialize)]
239struct RepoResponse {
240 default_branch: String,
241}
242
243#[derive(Debug)]
249pub struct GitHubSource<F: Fetcher> {
250 fetcher: F,
251 spec: GithubSpec,
252 api_base: CompactString,
253 raw_base: CompactString,
254 depth: Option<u16>,
256 max_bytes: u64,
257 concurrency: usize,
258 extensions: Vec<CompactString>,
261 validators: HashMap<UrlKey, CompactString>,
264 revalidated: std::sync::Mutex<Vec<Url>>,
266 failed: std::sync::atomic::AtomicU32,
268}
269
270impl<F: Fetcher> GitHubSource<F> {
271 #[must_use]
273 pub fn new(fetcher: F, spec: GithubSpec) -> Self {
274 Self {
275 fetcher,
276 spec,
277 api_base: "https://api.github.com".into(),
278 raw_base: "https://raw.githubusercontent.com".into(),
279 depth: None,
280 max_bytes: DEFAULT_MAX_BYTES,
281 concurrency: DEFAULT_CONCURRENCY,
282 extensions: Vec::new(),
283 validators: HashMap::new(),
284 revalidated: std::sync::Mutex::new(Vec::new()),
285 failed: std::sync::atomic::AtomicU32::new(0),
286 }
287 }
288
289 #[must_use]
291 pub fn bases(
292 mut self,
293 api_base: impl Into<CompactString>,
294 raw_base: impl Into<CompactString>,
295 ) -> Self {
296 self.api_base = api_base.into();
297 self.raw_base = raw_base.into();
298 self
299 }
300
301 #[must_use]
304 pub fn depth(mut self, depth: u16) -> Self {
305 self.depth = Some(depth);
306 self
307 }
308
309 #[must_use]
312 pub fn max_bytes(mut self, max_bytes: u64) -> Self {
313 self.max_bytes = max_bytes;
314 self
315 }
316
317 #[must_use]
319 pub fn concurrency(mut self, concurrency: usize) -> Self {
320 self.concurrency = concurrency.max(1);
321 self
322 }
323
324 #[must_use]
326 pub fn accept_extension(mut self, extension: impl AsRef<str>) -> Self {
327 let e = extension
328 .as_ref()
329 .trim_start_matches('.')
330 .to_ascii_lowercase();
331 self.extensions.push(CompactString::from(e));
332 self
333 }
334
335 #[must_use]
338 pub fn validators(
339 mut self,
340 entries: impl IntoIterator<Item = (UrlKey, CompactString)>,
341 ) -> Self {
342 self.validators = entries.into_iter().collect();
343 self
344 }
345
346 #[must_use]
350 pub fn take_revalidated(&self) -> Vec<Url> {
351 std::mem::take(&mut self.revalidated.lock().expect("revalidated lock poisoned"))
352 }
353
354 #[must_use]
356 pub fn failed_count(&self) -> u32 {
357 self.failed.load(std::sync::atomic::Ordering::Relaxed)
358 }
359
360 fn raw_url(&self, ref_name: &str, path: &str) -> Result<Url> {
362 let s = format!(
363 "{}/{}/{}/{ref_name}/{path}",
364 self.raw_base, self.spec.owner, self.spec.repo
365 );
366 Url::parse(&s).map_err(|e| Error::invalid_url(format!("{s:?}: {e}")))
367 }
368
369 async fn api_json(&self, url: &str) -> Result<Vec<u8>> {
371 let url = Url::parse(url).map_err(|e| Error::invalid_url(format!("{url:?}: {e}")))?;
372 let resource = Resource::new(url);
373 let opts = FetchOptions {
374 if_none_match: None,
375 max_bytes: Some(64 * 1024 * 1024),
378 user_agent: None,
379 };
380 let got = self.fetcher.fetch(&resource, opts).await?;
381 Ok(got.payload.into_bytes().await?.to_vec())
382 }
383
384 fn admit<'e>(&self, entry: &'e TreeEntry) -> Option<&'e str> {
387 if entry.kind != "blob" {
388 return None;
389 }
390 let sub = self.spec.subdir.as_str();
391 if !sub.is_empty() {
392 let rest = entry.path.strip_prefix(sub)?;
393 if !rest.starts_with('/') {
394 return None; }
396 }
397 if let Some(depth) = self.depth {
398 let rel_start = if sub.is_empty() { 0 } else { sub.len() + 1 };
399 let rel = &entry.path[rel_start..];
400 let dir_levels = rel.matches('/').count();
401 if dir_levels > usize::from(depth) {
402 return None;
403 }
404 }
405 if entry.size.is_some_and(|s| s > self.max_bytes) {
406 return None; }
408 let kind = ResourceKind::from_path(&entry.path);
409 if !kind.is_indexable() {
410 return None;
411 }
412 if !self.extensions.is_empty() {
413 let ext = entry.path.rsplit('/').next().unwrap_or("").rsplit_once('.');
414 let matched = ext.is_some_and(|(_, e)| {
415 self.extensions
416 .iter()
417 .any(|a| a.as_str().eq_ignore_ascii_case(e))
418 });
419 if !matched {
420 return None;
421 }
422 }
423 Some(&entry.path)
424 }
425}
426
427async fn fetch_blob<F: Fetcher>(
430 fetcher: &F,
431 url: Url,
432 sha: CompactString,
433 max_bytes: u64,
434) -> (Url, CompactString, Result<bytes::Bytes>) {
435 let resource = Resource::new(url);
436 let opts = FetchOptions {
437 if_none_match: None,
438 max_bytes: Some(max_bytes),
439 user_agent: None,
440 };
441 let outcome = match fetcher.fetch(&resource, opts).await {
442 Ok(got) => got.payload.into_bytes().await,
443 Err(e) => Err(e),
444 };
445 (resource.url, sha, outcome)
446}
447
448impl<F: Fetcher> Source for GitHubSource<F> {
449 type Pages<'a>
450 = impl futures::Stream<Item = Result<Page<'a>>> + Send + 'a
451 where
452 Self: 'a;
453
454 fn kind(&self) -> &'static str {
455 "github"
456 }
457
458 #[allow(clippy::too_many_lines)] fn discover<'a>(&'a self, _root: &'a SourceRef) -> Self::Pages<'a> {
460 async_stream::try_stream! {
461 self.failed.store(0, std::sync::atomic::Ordering::Relaxed);
463 self.revalidated.lock().expect("revalidated lock poisoned").clear();
464
465 let ref_name = if self.spec.ref_name.is_empty() {
467 let url = format!(
468 "{}/repos/{}/{}",
469 self.api_base, self.spec.owner, self.spec.repo
470 );
471 let body = self.api_json(&url).await?;
472 let repo: RepoResponse = serde_json::from_slice(&body).map_err(|e| {
473 Error::format(format!("github repo response: {e}"))
474 })?;
475 CompactString::from(repo.default_branch)
476 } else {
477 self.spec.ref_name.clone()
478 };
479
480 let url = format!(
482 "{}/repos/{}/{}/git/trees/{}?recursive=1",
483 self.api_base, self.spec.owner, self.spec.repo, ref_name
484 );
485 let body = self.api_json(&url).await?;
486 let tree: TreeResponse = serde_json::from_slice(&body)
487 .map_err(|e| Error::format(format!("github tree response: {e}")))?;
488 if tree.truncated {
489 Err(Error::crawl(
490 "github",
491 "tree listing truncated by GitHub (repository too large); \
492 narrow the spec to a subdirectory",
493 ))?;
494 }
495
496 let mut to_fetch: Vec<(Url, CompactString)> = Vec::new();
499 for entry in &tree.tree {
500 let Some(path) = self.admit(entry) else { continue };
501 let url = self.raw_url(&ref_name, path)?;
502 let sha = CompactString::from(entry.sha.as_str());
503 if self
504 .validators
505 .get(&UrlKey::from_url(&url))
506 .is_some_and(|known| *known == sha)
507 {
508 self.revalidated
509 .lock()
510 .expect("revalidated lock poisoned")
511 .push(url);
512 continue;
513 }
514 to_fetch.push((url, sha));
515 }
516
517 let mut queue = to_fetch.into_iter();
518 let mut in_flight = FuturesUnordered::new();
519 loop {
520 while in_flight.len() < self.concurrency {
521 let Some((url, sha)) = queue.next() else { break };
522 in_flight.push(fetch_blob(&self.fetcher, url, sha, self.max_bytes));
523 }
524 let Some((url, sha, outcome)) = in_flight.next().await else { break };
525 match outcome {
526 Ok(bytes) => {
527 let kind = ResourceKind::from_path(url.path());
528 let resource = Resource::new(url).with_kind(kind).with_etag(sha);
529 yield Page::new(resource, DocPayload::Owned(bytes));
530 }
531 Err(_) => {
532 self.failed.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
535 }
536 }
537 }
538 }
539 }
540}
541
542#[cfg(test)]
543mod tests {
544 use super::*;
545
546 fn spec(s: &str) -> GithubSpec {
547 GithubSpec::parse(s).unwrap()
548 }
549
550 #[test]
551 fn url_form_parses_the_default_corpus_shape() {
552 let s = spec("https://github.com/example-org/docs-site/tree/main/stacks");
553 assert_eq!(s.owner, "example-org");
554 assert_eq!(s.repo, "docs-site");
555 assert_eq!(s.ref_name, "main");
556 assert_eq!(s.subdir, "stacks");
557 }
558
559 #[test]
560 fn url_form_variants() {
561 assert_eq!(
562 spec("https://github.com/o/r"),
563 GithubSpec {
564 owner: "o".into(),
565 repo: "r".into(),
566 ref_name: "".into(),
567 subdir: "".into()
568 }
569 );
570 assert_eq!(spec("https://github.com/o/r.git").repo, "r");
571 let deep = spec("https://github.com/o/r/tree/v1.2.3/a/b/c");
572 assert_eq!(deep.ref_name, "v1.2.3");
573 assert_eq!(deep.subdir, "a/b/c");
574 assert!(GithubSpec::parse("https://github.com/o").is_err());
575 assert!(GithubSpec::parse("https://github.com/o/r/pulls").is_err());
576 }
577
578 #[test]
579 fn ref_grammar_parses_both_shapes() {
580 let s = spec("github:example-org/docs-site@main//stacks");
581 assert_eq!(
582 (s.owner.as_str(), s.repo.as_str()),
583 ("example-org", "docs-site")
584 );
585 assert_eq!((s.ref_name.as_str(), s.subdir.as_str()), ("main", "stacks"));
586 let s = spec("github:o/r@feat/x//docs/examples");
588 assert_eq!(s.ref_name, "feat/x");
589 assert_eq!(s.subdir, "docs/examples");
590 assert_eq!(spec("github:o/r").ref_name, "");
591 assert!(GithubSpec::parse("github:r").is_err());
592 assert!(GithubSpec::parse("github:o/r@").is_err());
593 assert!(GithubSpec::parse("gitlab:o/r").is_err());
594 }
595
596 #[test]
597 fn matches_routes_specs_not_ordinary_urls() {
598 assert!(GithubSpec::matches("github:o/r@main"));
599 assert!(GithubSpec::matches("https://github.com/o/r/tree/main/x"));
600 assert!(!GithubSpec::matches("https://docs.example.com/"));
601 assert!(!GithubSpec::matches(
602 "https://raw.githubusercontent.com/o/r/main/a.md"
603 ));
604 }
605
606 #[test]
607 fn github_auth_confines_the_token_to_the_two_hosts() {
608 let auth = GithubAuth::new(Some("sekrit".into()));
609 let cred = |u: &str| auth.credential(&Resource::new(Url::parse(u).unwrap()));
610 assert!(matches!(
611 cred("https://api.github.com/repos/o/r"),
612 Credential::Bearer(_)
613 ));
614 assert!(matches!(
615 cred("https://raw.githubusercontent.com/o/r/main/a.md"),
616 Credential::Bearer(_)
617 ));
618 assert!(matches!(
620 cred("https://evil.example.com/"),
621 Credential::Anonymous
622 ));
623 assert!(matches!(
624 cred("https://github.com/o/r"),
625 Credential::Anonymous
626 ));
627 let anon = GithubAuth::new(None);
629 assert!(matches!(
630 anon.credential(&Resource::new(
631 Url::parse("https://api.github.com/x").unwrap()
632 )),
633 Credential::Anonymous
634 ));
635 assert!(!format!("{auth:?}").contains("sekrit"));
637 }
638
639 #[test]
640 fn admit_applies_every_filter() {
641 let src = GitHubSource::new(DummyFetcher, spec("github:o/r@main//stacks"))
642 .depth(1)
643 .max_bytes(1000)
644 .accept_extension("md")
645 .accept_extension(".YAML");
646 let entry = |path: &str, size: u64| TreeEntry {
647 path: path.into(),
648 kind: "blob".into(),
649 sha: "abc".into(),
650 size: Some(size),
651 };
652 assert!(src.admit(&entry("stacks/api/README.md", 10)).is_some());
654 assert!(src.admit(&entry("stacks/Pulumi.yaml", 10)).is_some());
656 assert!(src.admit(&entry("stacks/a/b/x.md", 10)).is_none());
658 assert!(src.admit(&entry("other/x.md", 10)).is_none());
660 assert!(src.admit(&entry("stacksfoo/x.md", 10)).is_none());
661 assert!(src.admit(&entry("stacks/big.md", 100_000)).is_none());
663 assert!(src.admit(&entry("stacks/query.sql", 10)).is_none());
665 assert!(src.admit(&entry("stacks/logo.png", 10)).is_none());
667 let dir = TreeEntry {
669 path: "stacks/api".into(),
670 kind: "tree".into(),
671 sha: "d".into(),
672 size: None,
673 };
674 assert!(src.admit(&dir).is_none());
675 }
676
677 struct DummyFetcher;
679 impl Fetcher for DummyFetcher {
680 type FetchFuture<'a> = std::future::Ready<Result<crate::fetch::Fetched<'a>>>;
681 fn fetch<'a>(
682 &'a self,
683 resource: &'a Resource,
684 _opts: FetchOptions<'a>,
685 ) -> Self::FetchFuture<'a> {
686 std::future::ready(Err(Error::not_found(resource.url.as_str())))
687 }
688 }
689}