chadtmiller.com/knotless

//! A read accelerator for knotless repositories.
//!
//! A PDS can hand out records and blobs and nothing else: no partial clone,
//! no shallow fetch, no "what is in this file at this commit". This server
//! keeps a bare mirror per repository and answers those from it:
//!
//! - git smart HTTP (protocol v0 and v2, `--filter`, `--depth`) at
//!   `/<handle or did>/<name>`, read only;
//! - a JSON and raw API at `/api/<handle or did>/<name>/...` for agents that
//!   want a file, a tree, a log or a diff without a clone.
//!
//! It holds no state of its own: every mirror is a cache of the owner's
//! records, revalidated against them (one round trip) when older than the
//! TTL, or kept warm by the firehose or Jetstream. Every object it serves is named by its
//! hash, so a client can check it.

mod api;
mod appview;
pub mod card;
pub mod diff;
pub mod firehose;
pub mod gateway;
pub mod highlight;
pub mod index;
pub mod jetstream;
pub mod moderation;
pub mod objects;
mod smart;
pub mod web;

use anyhow::Result;
use axum::Router;
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use axum::routing::{get, post};
use knotless::env::Ctx;
use knotless::git::Git;
use knotless::mirror::Mirror;
use knotless::records::RepoRecord;
use knotless::remote::Remote;
use knotless::url::RepoUrl;
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

pub struct App {
    pub ctx: Ctx,
    /// How long a mirror is trusted without asking the PDS, when the
    /// firehose is not watching.
    pub ttl: Duration,
    repos: Mutex<HashMap<(String, String), Shared<Entry>>>,
    /// Firehose events seen per DID.
    events: Mutex<HashMap<String, u64>>,
    pub watching: AtomicBool,
    /// Set when this server also takes writes for owners who delegate them.
    pub gateway: Option<gateway::Gateway>,
    /// Set when this server is an appview, indexing what the firehose says.
    pub index: Option<Arc<index::Index>>,
    /// Set when this server also reads private repositories, as whoever is
    /// logged in: for that person alone, on their own machine.
    pub private: bool,
    /// The labels this server gives out, to those following them.
    pub labels: tokio::sync::broadcast::Sender<(i64, serde_json::Value)>,
}

type Shared<T> = Arc<tokio::sync::Mutex<T>>;

#[derive(Default)]
struct Entry {
    git: Option<Git>,
    /// The repository's owner, key and name now, which differ from what
    /// was asked for after a move or a rename.
    found: Option<(String, String, String)>,
    record: RepoRecord,
    space: bool,
    synced: Option<Instant>,
    events: u64,
}

/// A repository found by name, and its mirror.
pub struct Located {
    pub git: Git,
    pub did: String,
    pub key: String,
    /// Its name now: after a rename, not the name it was asked for by.
    pub name: String,
    /// Whether it is in another account than the one asked for: it moved.
    pub moved: bool,
    /// Its record, as the PDS had it when the mirror last synced.
    pub record: RepoRecord,
    /// Whether it is private, in a space.
    pub space: bool,
}

impl App {
    pub fn new(ctx: Ctx, ttl: Duration) -> Arc<App> {
        App::build(ctx, ttl, None, None, false)
    }

    pub fn build(
        ctx: Ctx,
        ttl: Duration,
        gateway: Option<gateway::Gateway>,
        index: Option<index::Index>,
        private: bool,
    ) -> Arc<App> {
        Arc::new(App {
            ctx,
            ttl,
            repos: Mutex::default(),
            events: Mutex::default(),
            watching: AtomicBool::new(false),
            gateway,
            index: index.map(Arc::new),
            private,
            labels: tokio::sync::broadcast::channel(256).0,
        })
    }

    /// Takes in a firehose or Jetstream event: refreshes the mirrors of an account whose
    /// repositories moved, and indexes what it wrote.
    pub fn observe(self: &Arc<App>, event: firehose::Event) {
        let Some(did) = event.did else {
            return;
        };
        let repository = event
            .ops
            .iter()
            .any(|op| knotless::records::nsid::ALL.contains(&op.collection.as_str()));
        if event.ops.is_empty() || repository {
            self.touch(&did);
        }
        if let Some(index) = self.index.clone()
            && !event.ops.is_empty()
        {
            // What the event carries in full is indexed now, in order, so a
            // later write cannot land before an earlier one.
            if event
                .ops
                .iter()
                .all(|op| op.action == "delete" || op.record.is_some())
            {
                for op in &event.ops {
                    let _ = index_inline(&index, &did, op);
                }
                return;
            }
            let ctx = self.ctx.clone();
            tokio::spawn(async move {
                for op in event.ops {
                    let _ = index_op(&ctx, &index, &did, &op).await;
                }
            });
        }
    }

    /// Firehose events seen for `did`.
    pub fn events_for(&self, did: &str) -> u64 {
        self.events.lock().unwrap().get(did).copied().unwrap_or(0)
    }

    /// Records that `did` wrote to a repository collection, and refreshes its
    /// mirrors in the background.
    pub fn touch(self: &Arc<App>, did: &str) {
        *self
            .events
            .lock()
            .unwrap()
            .entry(did.to_string())
            .or_default() += 1;
        let names: Vec<String> = self
            .repos
            .lock()
            .unwrap()
            .keys()
            .filter(|(d, _)| d == did)
            .map(|(_, n)| n.clone())
            .collect();
        for name in names {
            let app = self.clone();
            let did = did.to_string();
            tokio::spawn(async move {
                let _ = app.mirror(&did, &name).await;
            });
        }
    }

    /// The mirror of `authority/name`, synced if it may be stale.
    pub async fn mirror(&self, authority: &str, name: &str) -> Result<Git, AppError> {
        Ok(self.located(authority, name).await?.git)
    }

    /// The mirror of `authority/name`, with the repository's key and the
    /// name it has now.
    pub async fn located(&self, authority: &str, name: &str) -> Result<Located, AppError> {
        let url = RepoUrl::parse(&format!("knotless://{authority}/{name}"))
            .map_err(|e| AppError(StatusCode::BAD_REQUEST, e.to_string()))?;
        let remote = if self.private {
            Remote::open(&self.ctx, &url).await
        } else {
            Remote::connect(&self.ctx, &url).await
        }
        .map_err(AppError::upstream)?;
        let did = remote.did().to_string();
        let entry = self
            .repos
            .lock()
            .unwrap()
            .entry((did.clone(), url.name.clone()))
            .or_default()
            .clone();
        let mut entry = entry.lock().await;
        let events = self.events_for(&did);
        let fresh = match (&entry.git, entry.synced) {
            (Some(_), Some(_)) if self.watching.load(Ordering::Relaxed) => entry.events == events,
            (Some(_), Some(at)) => at.elapsed() < self.ttl,
            _ => false,
        };
        if !fresh {
            let snap = remote
                .snapshot(&self.ctx)
                .await
                .map_err(AppError::upstream)?;
            let (remote, snap) = remote
                .follow(&self.ctx, snap)
                .await
                .map_err(AppError::upstream)?;
            let Some(repo) = &snap.repo else {
                return Err(AppError(
                    StatusCode::NOT_FOUND,
                    format!("{url}: no such repository"),
                ));
            };
            let key = repo.value.key.as_str();
            let mirror = Mirror::open(&self.ctx, remote.did(), key)
                .await
                .map_err(AppError::internal)?;
            mirror
                .sync(&remote, &snap)
                .await
                .map_err(AppError::upstream)?;
            entry.found = Some((remote.did().to_string(), key.to_string(), repo.rkey.clone()));
            entry.record = repo.value.clone();
            entry.space = remote.space.is_some();
            entry.git = Some(mirror.git);
            entry.synced = Some(Instant::now());
            entry.events = events;
        }
        let (now, key, name) = entry.found.clone().expect("a fresh entry has a mirror");
        let uri = knotless::forge::repo_uri(&now, &name);
        if let Some(index) = &self.index
            && index
                .down(&[&now, &uri])
                .map_err(anyhow::Error::from)
                .map_err(AppError::internal)?
        {
            return Err(AppError(
                StatusCode::NOT_FOUND,
                format!("{url}: taken down"),
            ));
        }
        Ok(Located {
            git: entry.git.clone().expect("a fresh entry has a mirror"),
            moved: now != did,
            did: now,
            key,
            name,
            record: entry.record.clone(),
            space: entry.space,
        })
    }
}

/// Indexes a write the event carried in full: a deletion, or a record
/// delivered inline. `false` when the record has to be fetched.
fn index_inline(index: &index::Index, did: &str, op: &firehose::Op) -> rusqlite::Result<bool> {
    match (&op.record, op.action.as_str()) {
        (_, "delete") => index.delete(did, &op.collection, &op.rkey)?,
        (Some(record), _) => index.put(did, &op.collection, &op.rkey, record)?,
        (None, _) => return Ok(false),
    }
    Ok(true)
}

/// Indexes one write: the record the event carried, else the record as its
/// PDS now has it, or its absence.
pub async fn index_op(ctx: &Ctx, index: &index::Index, did: &str, op: &firehose::Op) -> Result<()> {
    if index_inline(index, did, op)? {
        return Ok(());
    }
    let id = knotless::identity::resolve(ctx, did).await?;
    let out: serde_json::Value = knotless::xrpc::Xrpc::new(ctx.http.clone(), &id.pds)
        .query(
            "com.atproto.repo.getRecord",
            &[
                ("repo", did),
                ("collection", &op.collection),
                ("rkey", &op.rkey),
            ],
            None,
        )
        .await?;
    index.put(did, &op.collection, &op.rkey, &out["value"])?;
    Ok(())
}

/// Backfills from `service`, saying how it went.
pub async fn backfilled(ctx: &Ctx, index: &index::Index, service: &str) -> String {
    match backfill(ctx, index, service).await {
        Ok(n) => format!("backfilled {n} records from {service}"),
        Err(err) => format!("backfilling from {service}: {err:#}"),
    }
}

/// Indexes every account on `service`, a PDS or relay: each account's
/// repositories and forge records.
pub async fn backfill(ctx: &Ctx, index: &index::Index, service: &str) -> Result<usize> {
    #[derive(serde::Deserialize)]
    struct Repos {
        repos: Vec<Account>,
        cursor: Option<String>,
    }
    #[derive(serde::Deserialize)]
    struct Account {
        did: String,
    }
    let xrpc = knotless::xrpc::Xrpc::new(ctx.http.clone(), service);
    let mut dids = Vec::new();
    let mut cursor: Option<String> = None;
    loop {
        let mut params = vec![("limit", "1000")];
        if let Some(c) = &cursor {
            params.push(("cursor", c));
        }
        let page: Repos = xrpc
            .query("com.atproto.sync.listRepos", &params, None)
            .await?;
        let done = page.repos.is_empty();
        dids.extend(page.repos.into_iter().map(|r| r.did));
        match page.cursor {
            Some(next) if !done => cursor = Some(next),
            _ => break,
        }
    }
    let mut indexed = 0;
    for did in dids {
        let url = RepoUrl::parse(&format!("knotless://{did}/_"))?;
        let remote = Remote::connect(ctx, &url).await?;
        let collections = [
            knotless::records::nsid::REPO,
            knotless::records::nsid::CLAIM,
        ]
        .into_iter()
        .chain(knotless::records::nsid::FORGE);
        for collection in collections {
            for r in remote.list_all::<serde_json::Value>(collection).await? {
                index.put(&did, collection, &r.rkey, &r.value)?;
                indexed += 1;
            }
        }
    }
    Ok(indexed)
}

/// A status and a message, as the response.
#[derive(Debug)]
pub struct AppError(pub StatusCode, pub String);

impl AppError {
    pub(crate) fn upstream(err: anyhow::Error) -> AppError {
        AppError(StatusCode::BAD_GATEWAY, format!("{err:#}"))
    }
    /// 400 for input the request got wrong, else 500.
    pub(crate) fn input(err: anyhow::Error) -> AppError {
        let status = match err.downcast_ref::<knotless::read::BadInput>() {
            Some(_) => StatusCode::BAD_REQUEST,
            None => StatusCode::INTERNAL_SERVER_ERROR,
        };
        AppError(status, format!("{err:#}"))
    }

    pub(crate) fn internal(err: anyhow::Error) -> AppError {
        AppError(StatusCode::INTERNAL_SERVER_ERROR, format!("{err:#}"))
    }
}

impl IntoResponse for AppError {
    fn into_response(self) -> Response {
        (self.0, self.1).into_response()
    }
}

pub fn router(app: Arc<App>) -> Router {
    Router::new()
        .route("/healthz", get(|| async { "ok" }))
        .route("/{authority}/{repo}/info/refs", get(smart::info_refs))
        .route(
            "/{authority}/{repo}/git-upload-pack",
            post(smart::upload_pack),
        )
        .route("/api/{authority}/{repo}/refs", get(api::refs))
        .route("/api/{authority}/{repo}/tree", get(api::tree))
        .route("/api/{authority}/{repo}/blob", get(api::blob))
        .route("/api/{authority}/{repo}/log", get(api::log))
        .route("/api/{authority}/{repo}/diff", get(api::diff))
        .route("/api/{authority}/{repo}/object/{sha}", get(api::object))
        .route("/.well-known/did.json", get(gateway::did_document))
        .route("/xrpc/dev.knotless.putDelegation", post(gateway::delegate))
        .route(
            "/xrpc/dev.knotless.getDelegation",
            get(gateway::await_delegation),
        )
        .route(gateway::METADATA, get(gateway::client_metadata))
        .route("/oauth/callback", get(gateway::callback))
        .route(
            "/xrpc/dev.knotless.pushRefs",
            post(gateway::push).layer(axum::extract::DefaultBodyLimit::max(PUSH_LIMIT)),
        )
        .route("/xrpc/dev.knotless.createCommit", post(gateway::commit))
        .route("/xrpc/dev.knotless.getForge", get(appview::get_forge))
        .route("/xrpc/dev.knotless.listRepos", get(appview::list_repos))
        .route("/xrpc/dev.knotless.getFeed", get(web::get_feed))
        .route(
            "/xrpc/com.atproto.moderation.createReport",
            post(moderation::create_report),
        )
        .route(
            "/xrpc/com.atproto.admin.updateSubjectStatus",
            post(moderation::update_subject_status),
        )
        .route(
            "/xrpc/com.atproto.admin.getSubjectStatus",
            get(moderation::get_subject_status),
        )
        .route(
            "/xrpc/dev.knotless.listReports",
            get(moderation::list_reports),
        )
        .route(
            "/xrpc/com.atproto.label.queryLabels",
            get(moderation::query_labels),
        )
        .route(
            "/xrpc/com.atproto.label.subscribeLabels",
            get(moderation::subscribe_labels),
        )
        .merge(web::routes())
        .with_state(app)
}

/// The largest push a gateway takes.
const PUSH_LIMIT: usize = 512 << 20;

/// `name` with a trailing `.git` dropped, as clone URLs often carry one.
fn repo_name(name: &str) -> &str {
    name.strip_suffix(".git").unwrap_or(name)
}

/// The arguments `knotless-serve` takes.
#[derive(Debug, PartialEq)]
pub struct Options {
    pub listen: String,
    pub ttl: Duration,
    pub firehose: Option<String>,
    /// A Jetstream v2 host, followed like the firehose.
    pub jetstream: Option<String>,
    /// The URL this server is reached at, which makes it a gateway.
    pub public_url: Option<String>,
    /// Where to keep the appview index, which makes it an appview.
    pub index: Option<std::path::PathBuf>,
    /// A PDS or relay to index every account of at startup.
    pub backfill: Option<String>,
    /// Labelers, by DID, whose takedowns hold here too.
    pub labelers: Vec<String>,
    /// Read private repositories too, as whoever is logged in.
    pub private: bool,
    /// The SvelteKit app that draws the web pages, which this server hands
    /// every request it does not answer itself.
    pub web: Option<String>,
}

pub fn parse_args(args: &[String]) -> Result<Options, String> {
    let mut opts = Options {
        listen: "127.0.0.1:2590".into(),
        ttl: Duration::from_secs(5),
        firehose: None,
        jetstream: None,
        public_url: None,
        index: None,
        backfill: None,
        labelers: Vec::new(),
        private: false,
        web: None,
    };
    let mut it = args.iter().skip(1);
    while let Some(flag) = it.next() {
        let mut value = || it.next().cloned().ok_or(format!("{flag} takes a value"));
        match flag.as_str() {
            "--listen" => opts.listen = value()?,
            "--ttl" => {
                opts.ttl = Duration::from_secs(
                    value()?
                        .parse()
                        .map_err(|_| "--ttl takes seconds".to_string())?,
                )
            }
            "--firehose" => opts.firehose = Some(value()?),
            "--jetstream" => opts.jetstream = Some(value()?),
            "--public-url" => opts.public_url = Some(value()?),
            "--index" => opts.index = Some(value()?.into()),
            "--backfill" => opts.backfill = Some(value()?),
            "--labeler" => opts.labelers.push(value()?),
            "--private" => opts.private = true,
            "--web" => opts.web = Some(value()?),
            other => {
                return Err(format!(
                    "unexpected {other}\nusage: knotless-serve [--listen <addr>] [--ttl <secs>] [--firehose <wss://relay or pds>] [--jetstream <wss://host>] [--public-url <https://host>] [--index <file>] [--backfill <https://pds or relay>] [--labeler <did>]... [--private] [--web <http://app>]"
                ));
            }
        }
    }
    Ok(opts)
}

/// Serves until SIGTERM or SIGINT. Prints the bound address first.
pub async fn run(opts: Options, ctx: Ctx) -> Result<()> {
    // Its pages show private repositories to anyone who reaches them.
    if opts.private
        && tokio::net::lookup_host(&opts.listen)
            .await?
            .any(|a| !a.ip().is_loopback())
    {
        anyhow::bail!("--private listens only on this machine: give --listen a loopback address");
    }
    let gateway = match &opts.public_url {
        Some(url) => Some(gateway::Gateway::new(url).map_err(anyhow::Error::msg)?),
        None => None,
    };
    let index = match &opts.index {
        Some(path) => Some(index::Index::open(path)?),
        None => None,
    };
    let app = App::build(ctx, opts.ttl, gateway, index, opts.private);
    if let (Some(service), Some(index)) = (opts.backfill, app.index.clone()) {
        let ctx = app.ctx.clone();
        tokio::spawn(async move { eprintln!("{}", backfilled(&ctx, &index, &service).await) });
    }
    for did in opts.labelers {
        tokio::spawn(moderation::follow(app.clone(), did));
    }
    if let Some(url) = opts.firehose {
        app.watching.store(true, Ordering::Relaxed);
        tokio::spawn(firehose::watch(app.clone(), url));
    }
    if let Some(url) = opts.jetstream {
        app.watching.store(true, Ordering::Relaxed);
        tokio::spawn(jetstream::watch(app.clone(), url));
    }
    let listener = tokio::net::TcpListener::bind(&opts.listen).await?;
    println!("listening on http://{}", listener.local_addr()?);
    // With no pages to hand on to, what this server does not answer is not found.
    let pages = opts.web.map(web::pages).unwrap_or_default();
    let routes = router(app).merge(pages);
    axum::serve(listener, routes)
        .with_graceful_shutdown(async {
            let mut term =
                tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
                    .expect("installing a SIGTERM handler");
            tokio::select! {
                _ = term.recv() => {}
                _ = tokio::signal::ctrl_c() => {}
            }
        })
        .await?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn input_errors() {
        let bad = AppError::input(knotless::read::BadInput("no".into()).into());
        assert_eq!(bad.0, StatusCode::BAD_REQUEST);
        let broken = AppError::input(anyhow::anyhow!("disk on fire"));
        assert_eq!(broken.0, StatusCode::INTERNAL_SERVER_ERROR);
    }
}