117 lines
3.4 KiB
Rust
117 lines
3.4 KiB
Rust
use crate::parse_seeds;
|
|
use anyhow::{anyhow, bail};
|
|
use log::warn;
|
|
use smartstring::alias::CompactString;
|
|
use std::collections::BTreeSet;
|
|
use std::ffi::OsStr;
|
|
use std::path::PathBuf;
|
|
use tokio::sync::mpsc::Sender;
|
|
|
|
pub const SEED_EXTENSION: &'static str = ".seed";
|
|
pub const WEED_EXTENSION: &'static str = ".weed";
|
|
|
|
pub struct Seed {
|
|
pub url: UrlOrUrlPattern,
|
|
pub tags: BTreeSet<CompactString>,
|
|
}
|
|
|
|
/// Either a URL or a URL prefix.
|
|
#[derive(Clone, Debug)]
|
|
pub enum UrlOrUrlPattern {
|
|
Url(String),
|
|
UrlPrefix(String),
|
|
}
|
|
|
|
impl UrlOrUrlPattern {
|
|
pub fn as_str(&self) -> &str {
|
|
match self {
|
|
UrlOrUrlPattern::Url(url) => url.as_str(),
|
|
UrlOrUrlPattern::UrlPrefix(url_prefix) => url_prefix.as_str(),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Task that loads seeds from the filesystem
|
|
pub async fn seed_loader(seed_files: Vec<PathBuf>, send: &Sender<Seed>) -> anyhow::Result<()> {
|
|
for seed_file in seed_files {
|
|
// Parse the seed file and send out the seeds.
|
|
let seed_file_text = tokio::fs::read_to_string(&seed_file).await?;
|
|
match parse_seeds(&seed_file_text) {
|
|
Ok(seedblocks) => {
|
|
for seedblock in seedblocks {
|
|
for seed in seedblock.seeds {
|
|
let tags: BTreeSet<CompactString> = seedblock
|
|
.tags
|
|
.iter()
|
|
.chain(seed.extra_tags.iter())
|
|
.cloned()
|
|
.collect();
|
|
|
|
send.send(Seed {
|
|
url: seed_url_parse_pattern(seed.url),
|
|
tags,
|
|
})
|
|
.await
|
|
.map_err(|_| anyhow!("Seed receiver shut down prematurely"))?;
|
|
}
|
|
}
|
|
}
|
|
Err(err) => {
|
|
eprintln!(
|
|
"~~~~~ Error in seed file ({:?}):\n{:?}\n~~~~~",
|
|
seed_file, err
|
|
);
|
|
bail!("Failed to parse {:?}; see error above.", seed_file);
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
pub fn seed_url_parse_pattern(mut url: String) -> UrlOrUrlPattern {
|
|
if url.ends_with('*') {
|
|
url.pop();
|
|
UrlOrUrlPattern::UrlPrefix(url)
|
|
} else {
|
|
UrlOrUrlPattern::Url(url)
|
|
}
|
|
}
|
|
|
|
pub async fn find_seed_files(seed_dir: PathBuf, extension: &str) -> anyhow::Result<Vec<PathBuf>> {
|
|
let mut dirs = vec![seed_dir];
|
|
let mut seedfiles = Vec::new();
|
|
|
|
while let Some(dir_to_scan) = dirs.pop() {
|
|
let mut dir = tokio::fs::read_dir(&dir_to_scan).await?;
|
|
|
|
while let Some(entry) = dir.next_entry().await? {
|
|
let path = entry.path();
|
|
let file_name = match path
|
|
.file_name()
|
|
.map(|osstr: &OsStr| osstr.to_str())
|
|
.flatten()
|
|
{
|
|
None => {
|
|
warn!("Skipping non-UTF-8 name.");
|
|
continue;
|
|
}
|
|
Some(file_name) => file_name,
|
|
};
|
|
if file_name.starts_with(".") {
|
|
continue;
|
|
}
|
|
if file_name.ends_with(extension) {
|
|
seedfiles.push(path);
|
|
continue;
|
|
}
|
|
if path.is_dir() {
|
|
// Recurse into this directory later.
|
|
dirs.push(path);
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(seedfiles)
|
|
}
|