pdsfs / src / main.rs

#![feature(let_chains)]
mod client;
mod error;
mod fs;
mod resolver;

use atrium_api::{client::AtpServiceClient, com, types};
use atrium_common::resolver::Resolver;
use atrium_identity::identity_resolver::ResolvedIdentity;
use atrium_repo::{Repository, blockstore::CarStore};
use atrium_xrpc_client::isahc::IsahcClient;
use fuser::MountOption;
use futures::{Stream, StreamExt, TryStream, TryStreamExt, stream};
use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
use std::{
    collections::HashMap,
    io::{Cursor, Write},
    path::PathBuf,
    sync::Arc,
};
use xdg::BaseDirectories;

fn main() {
    let rt = tokio::runtime::Runtime::new().unwrap();
    let matches = clap::command!()
        .arg(
            clap::Arg::new("handles")
                .index(1)
                .required(true)
                .num_args(1..)
                .help("One or more handles to download and mount"),
        )
        .arg(
            clap::Arg::new("mountpoint")
                .short('m')
                .action(clap::ArgAction::Set)
                .value_parser(clap::value_parser!(PathBuf)),
        )
        .get_matches();
    let handles = matches
        .get_many::<String>("handles")
        .unwrap()
        .cloned()
        .collect::<Vec<_>>();
    let mountpoint = matches
        .get_one::<PathBuf>("mountpoint")
        .map(ToOwned::to_owned)
        .unwrap_or(PathBuf::from("mnt"));
    let _ = std::fs::create_dir_all(&mountpoint);

    let resolver = Arc::new(resolver::id_resolver());
    let bars = Arc::new(MultiProgress::new());
    let repos = rt.block_on(
        stream::iter(handles)
            .then(|handle| {
                let h = handle.clone();
                let r = Arc::clone(&resolver);
                let b = Arc::clone(&bars);
                async move {
                    let id = r.resolve(&h).await?;
                    let bytes = cached_download(&id, &b).await?;
                    let repo = build_repo(bytes).await?;
                    Ok::<_, error::Error>((id.did, repo))
                }
            })
            .collect::<Vec<_>>(),
    );
    let (success, errors): (Vec<_>, Vec<_>) = repos.into_iter().partition(|r| r.is_ok());
    for e in errors {
        eprintln!("{:?}", e.as_ref().unwrap_err());
    }
    let repos = success
        .into_iter()
        .map(|s| s.unwrap())
        .collect::<HashMap<_, _>>();

    // construct the fs
    let mut fs = fs::PdsFs::new();
    for (did, repo) in repos {
        rt.block_on(fs.add(did, repo))
    }

    // mount
    let options = vec![MountOption::RO, MountOption::FSName("pdsfs".to_string())];
    let join_handle = fuser::spawn_mount2(fs, &mountpoint, &options).unwrap();

    println!("mounted at {mountpoint:?}");
    print!("hit enter to unmount and exit...");
    std::io::stdout().flush().unwrap();

    // Wait for user input
    let mut input = String::new();
    std::io::stdin().read_line(&mut input).unwrap();

    join_handle.join();
    std::fs::remove_dir(&mountpoint).unwrap();

    println!("unmounted {mountpoint:?}");
}

async fn cached_download(
    id: &ResolvedIdentity,
    m: &MultiProgress,
) -> Result<Vec<u8>, error::Error> {
    let mut pb = ProgressBar::new_spinner();
    pb.set_style(
        ProgressStyle::default_spinner()
            .template("{spinner:.green} [{elapsed_precise}] {msg}")
            .unwrap()
            .tick_strings(&["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]),
    );
    pb.enable_steady_tick(std::time::Duration::from_millis(100));
    pb = m.add(pb);

    let dirs = BaseDirectories::new();

    let dir = dirs.get_data_home().expect("$HOME is absent").join("pdsfs");
    tokio::fs::create_dir_all(&dir).await?;

    let file = dir.join(&id.did);
    let exists = std::fs::exists(&file)?;

    let bytes = if !exists {
        pb.set_message(format!("downloading CAR file for...{}", id.did));
        download_car_file(id, &pb).await?
    } else {
        pb.set_message(format!("using cached CAR file for...{}", id.did));
        tokio::fs::read(&file).await?
    };

    // write to disk
    if !exists {
        tokio::fs::write(&file, &bytes).await?;
    }

    pb.finish();
    Ok(bytes)
}

async fn download_car_file(
    id: &ResolvedIdentity,
    pb: &ProgressBar,
) -> Result<Vec<u8>, error::Error> {
    // download the entire car file first before mounting it as a fusefs
    let client = AtpServiceClient::new(IsahcClient::new(&id.pds));
    let did = types::string::Did::new(id.did.clone()).unwrap();

    let bytes = client
        .service
        .com
        .atproto
        .sync
        .get_repo(com::atproto::sync::get_repo::Parameters::from(
            com::atproto::sync::get_repo::ParametersData { did, since: None },
        ))
        .await?;

    pb.finish_with_message(format!("download complete for \t...\t{}", id.did));

    Ok(bytes)
}

async fn build_repo(bytes: Vec<u8>) -> Result<Repository<CarStore<Cursor<Vec<u8>>>>, error::Error> {
    let store = CarStore::open(Cursor::new(bytes)).await?;
    let root = store.roots().next().unwrap();
    let repo = Repository::open(store, root).await?;
    Ok(repo)
}