chadtmiller.com/knotless

//! The gateway: writes on owners' behalf, for collaborators and for batches.
//!
//! The gateway is an OAuth client of its owners' PDSes. An owner delegates
//! (`dev.knotless.putDelegation`) by approving the gateway's authorization
//! request, which grants it the knotless records and blobs and nothing else
//! of the account, and the gateway keeps the session. After that, anyone on a repository's collaborator list can push
//! (`dev.knotless.pushRefs`) or commit (`dev.knotless.createCommit`) with a service
//! token from their own PDS, and the gateway writes as the owner.
//!
//! Commits to one branch queue up while a write to it is in flight. The
//! next write takes the whole queue: each commit is built on the one before
//! it, and the chain goes out as one ref update. Sixteen agents committing
//! to one branch cost about two writes, not sixteen rounds of losing a race.

use crate::{App, AppError};
use axum::Json;
use axum::body::Bytes;
use axum::extract::{Query, State};
use axum::http::{HeaderMap, StatusCode, header};
use knotless::commit::{self, Committed, Spec};
use knotless::gateway::{
    self, AWAIT_DELEGATION, COMMIT, CommitRequest, CommitResponse, DELEGATE, PUSH, PushResults,
    RefResult,
};
use knotless::mirror::Mirror;
use knotless::push::{self, Update};
use knotless::remote::{Remote, Snapshot};
use knotless::session::{self, Session};
use knotless::url::RepoUrl;
use knotless::{identity, oauth};
use serde::Deserialize;
use serde_json::{Value, json};
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use tokio::sync::{oneshot, watch};

/// How many times a batch is rebuilt on a branch that moved under it.
const RETRIES: u32 = 8;

pub struct Gateway {
    pub did: String,
    pub endpoint: String,
    /// Batches written, each one ref update however many commits it holds.
    pub writes: AtomicU64,
    /// Batches rebuilt because the branch moved under them.
    pub rebuilt: AtomicU64,
    /// Where browsers and authorization servers reach the gateway as an
    /// OAuth client: the endpoint, unless that is served elsewhere too.
    pub oauth_url: String,
    /// How long `getDelegation` waits for an approval before it says to
    /// ask again.
    pub wait: Duration,
    repos: Mutex<HashMap<(String, String), Lock>>,
    branches: Mutex<HashMap<(String, String, String), Arc<Branch>>>,
    /// Delegations waiting for their owners' approval, by state.
    delegations: Mutex<HashMap<String, Delegation>>,
}

type Outcome = watch::Sender<Option<Result<(), String>>>;

struct Delegation {
    did: String,
    started: Instant,
    /// Taken by the browser's return.
    request: Option<oauth::Request>,
    done: Arc<Outcome>,
}

/// Where the gateway's OAuth client metadata is, its client ID.
pub const METADATA: &str = "/oauth-client-metadata.json";
const CALLBACK: &str = "/oauth/callback";
/// How long a delegation waits for approval before it is dropped.
const ABANDONED: Duration = Duration::from_secs(15 * 60);

/// Held while a repository is written.
type Lock = Arc<tokio::sync::Mutex<()>>;

#[derive(Default)]
struct Branch {
    queue: Mutex<Vec<Pending>>,
    turn: tokio::sync::Mutex<()>,
}

struct Pending {
    spec: Spec,
    reply: oneshot::Sender<Result<CommitResponse, AppError>>,
}

impl Gateway {
    /// A gateway reached at `endpoint`, an `https://` or `http://` URL, as
    /// the `did:web` of its host.
    pub fn new(endpoint: &str) -> Result<Gateway, String> {
        let endpoint = endpoint.trim_end_matches('/');
        let host = endpoint
            .strip_prefix("https://")
            .or_else(|| endpoint.strip_prefix("http://"))
            .filter(|h| !h.is_empty() && !h.contains('/'))
            .ok_or(format!(
                "{endpoint}: a gateway URL is http(s)://host[:port]"
            ))?;
        Ok(Gateway {
            did: format!("did:web:{}", host.replace(':', "%3A")),
            endpoint: endpoint.to_string(),
            writes: Default::default(),
            rebuilt: Default::default(),
            oauth_url: endpoint.to_string(),
            wait: Duration::from_secs(25),
            repos: Mutex::default(),
            branches: Mutex::default(),
            delegations: Mutex::default(),
        })
    }

    /// The gateway as an OAuth client, signing with the key in its config
    /// directory.
    fn client(&self, ctx: &knotless::env::Ctx) -> Result<oauth::Client, AppError> {
        let key = oauth::client_key(ctx).map_err(AppError::internal)?;
        Ok(oauth::Client::confidential(
            &format!("{}{METADATA}", self.oauth_url),
            key,
        ))
    }
}

fn forbidden(msg: impl Into<String>) -> AppError {
    AppError(StatusCode::FORBIDDEN, msg.into())
}

fn bad(msg: impl Into<String>) -> AppError {
    AppError(StatusCode::BAD_REQUEST, msg.into())
}

pub(crate) fn gateway(app: &App) -> Result<&Gateway, AppError> {
    app.gateway
        .as_ref()
        .ok_or_else(|| AppError(StatusCode::NOT_FOUND, "this server is not a gateway".into()))
}

/// The DID that signed the request's service token for `lxm`.
pub(crate) async fn caller(app: &App, headers: &HeaderMap, lxm: &str) -> Result<String, AppError> {
    let gw = gateway(app)?;
    let token = headers
        .get(header::AUTHORIZATION)
        .and_then(|v| v.to_str().ok())
        .and_then(|v| v.strip_prefix("Bearer "))
        .ok_or_else(|| {
            AppError(
                StatusCode::UNAUTHORIZED,
                "a service token is required".into(),
            )
        })?;
    knotless::auth::verify(&app.ctx, token, &gw.did, lxm)
        .await
        .map_err(|e| AppError(StatusCode::UNAUTHORIZED, format!("{e:#}")))
}

pub async fn did_document(State(app): State<Arc<App>>) -> Result<Json<Value>, AppError> {
    let gw = gateway(&app)?;
    Ok(Json(json!({
        "@context": ["https://www.w3.org/ns/did/v1"],
        "id": gw.did,
        "service": [{
            "id": gateway::SERVICE,
            "type": "KnotlessGateway",
            "serviceEndpoint": gw.endpoint,
        }, {
            "id": knotless::moderation::LABELER,
            "type": "AtprotoLabeler",
            "serviceEndpoint": gw.endpoint,
        }],
    })))
}

/// The gateway's OAuth client metadata, which its client ID names.
pub async fn client_metadata(State(app): State<Arc<App>>) -> Result<Json<Value>, AppError> {
    let gw = gateway(&app)?;
    let client = gw.client(&app.ctx)?;
    Ok(Json(json!({
        "client_id": client.id,
        "client_name": "knotless gateway",
        "application_type": "web",
        "grant_types": ["authorization_code", "refresh_token"],
        "response_types": ["code"],
        "redirect_uris": [format!("{}{CALLBACK}", gw.oauth_url)],
        "scope": gateway::SCOPE,
        "token_endpoint_auth_method": "private_key_jwt",
        "token_endpoint_auth_signing_alg": "ES256",
        "dpop_bound_access_tokens": true,
        "jwks": client.jwks(),
    })))
}

/// Starts delegating the caller's writes here: asks their PDS for a grant,
/// and answers with the page where they approve it.
pub async fn delegate(
    State(app): State<Arc<App>>,
    headers: HeaderMap,
) -> Result<Json<Value>, AppError> {
    let did = caller(&app, &headers, DELEGATE).await?;
    let gw = gateway(&app)?;
    let id = identity::resolve(&app.ctx, &did)
        .await
        .map_err(AppError::upstream)?;
    let back = format!("{}{CALLBACK}", gw.oauth_url);
    let client = gw.client(&app.ctx)?;
    let request = oauth::authorize(&app.ctx, &client, id, &did, &back, gateway::SCOPE)
        .await
        .map_err(AppError::upstream)?;
    let out = json!({ "url": request.url, "state": request.state });
    let mut delegations = gw.delegations.lock().unwrap();
    delegations.retain(|_, d| d.started.elapsed() < ABANDONED);
    delegations.insert(
        request.state.clone(),
        Delegation {
            did,
            started: Instant::now(),
            request: Some(request),
            done: Arc::new(watch::channel(None).0),
        },
    );
    Ok(Json(out))
}

/// Where the owner's browser comes back to: trades the code for the
/// session the gateway writes with.
pub async fn callback(
    State(app): State<Arc<App>>,
    Query(params): Query<HashMap<String, String>>,
) -> Result<String, AppError> {
    let gw = gateway(&app)?;
    let state = params.get("state").map_or("", String::as_str);
    let waiting = gw
        .delegations
        .lock()
        .unwrap()
        .get_mut(state)
        .and_then(|d| Some((d.request.take()?, d.done.clone())));
    let Some((request, done)) = waiting else {
        return Err(bad("no delegation is waiting for this"));
    };
    let result = match request
        .redeem(&app.ctx, &gw.client(&app.ctx)?, &params)
        .await
    {
        Ok(session) => session::save(&app.ctx, &session),
        Err(err) => Err(err),
    };
    let page = match &result {
        Ok(()) => "The gateway now writes for you. You can close this tab.".to_string(),
        Err(err) => format!("The gateway was not delegated to: {err:#}"),
    };
    done.send_replace(Some(result.map_err(|e| format!("{e:#}"))));
    Ok(page)
}

#[derive(Deserialize)]
pub struct Awaited {
    state: String,
}

/// Waits a while for the caller's delegation to be approved.
pub async fn await_delegation(
    State(app): State<Arc<App>>,
    headers: HeaderMap,
    Query(query): Query<Awaited>,
) -> Result<Json<Value>, AppError> {
    let did = caller(&app, &headers, AWAIT_DELEGATION).await?;
    let gw = gateway(&app)?;
    let mut done = gw
        .delegations
        .lock()
        .unwrap()
        .get(&query.state)
        .filter(|d| d.did == did)
        .map(|d| d.done.subscribe())
        .ok_or_else(|| AppError(StatusCode::NOT_FOUND, "no such delegation".into()))?;
    let waited = tokio::time::timeout(gw.wait, done.wait_for(Option::is_some)).await;
    let Ok(Ok(outcome)) = waited else {
        return Ok(Json(json!({ "done": false })));
    };
    let outcome = outcome.clone().expect("waited for an outcome");
    gw.delegations.lock().unwrap().remove(&query.state);
    outcome.map_err(forbidden)?;
    Ok(Json(json!({ "done": true })))
}

/// The repository, as its owner, for a caller who may write to it.
async fn writable(
    app: &App,
    caller: &str,
    owner: &str,
    name: &str,
) -> Result<(Remote, Session, Snapshot), AppError> {
    let url =
        RepoUrl::parse(&format!("knotless://{owner}/{name}")).map_err(|e| bad(e.to_string()))?;
    let remote = Remote::open(&app.ctx, &url)
        .await
        .map_err(AppError::upstream)?;
    let session = session::ensure(&app.ctx, &remote.identity)
        .await
        .map_err(|_| forbidden(format!("{owner} has not delegated writes to this gateway")))?;
    let snap = remote
        .snapshot(&app.ctx)
        .await
        .map_err(AppError::upstream)?;
    let Some(repo) = &snap.repo else {
        return Err(AppError(
            StatusCode::NOT_FOUND,
            format!("{url}: no such repository"),
        ));
    };
    if caller != remote.did() && !repo.value.is_collaborator(caller) {
        return Err(forbidden(format!(
            "{caller} is not a collaborator on {url}"
        )));
    }
    Ok((remote, session, snap))
}

async fn mirror(app: &App, remote: &Remote, snap: &Snapshot) -> Result<Mirror, AppError> {
    let key = snap.key().expect("a writable repository exists");
    let mirror = Mirror::open(&app.ctx, remote.did(), key)
        .await
        .map_err(AppError::internal)?;
    mirror
        .sync(remote, snap)
        .await
        .map_err(AppError::upstream)?;
    Ok(mirror)
}

pub async fn push(
    State(app): State<Arc<App>>,
    headers: HeaderMap,
    body: Bytes,
) -> Result<Json<PushResults>, AppError> {
    let caller = caller(&app, &headers, PUSH).await?;
    let (header, pack) = gateway::split_push(&body).map_err(|e| bad(format!("{e:#}")))?;
    let mut updates = Vec::new();
    for u in &header.updates {
        if u.src.as_deref().is_some_and(|s| !gateway::is_sha(s)) || !u.dst.starts_with("refs/") {
            return Err(bad(format!("bad update of {}", u.dst)));
        }
        updates.push(Update {
            src: u.src.clone(),
            dst: u.dst.clone(),
            force: u.force,
        });
    }
    let lock = gateway(&app)?
        .repos
        .lock()
        .unwrap()
        .entry((header.repo.clone(), header.name.clone()))
        .or_default()
        .clone();
    let _turn = lock.lock().await;
    let (remote, session, snap) = writable(&app, &caller, &header.repo, &header.name).await?;
    let mirror = mirror(&app, &remote, &snap).await?;
    if !pack.is_empty() {
        mirror
            .git
            .index_pack(pack)
            .await
            .map_err(|e| bad(format!("{e:#}")))?;
    }
    let tips: Vec<String> = updates.iter().filter_map(|u| u.src.clone()).collect();
    if let Some(sha) = mirror
        .git
        .missing(&tips)
        .await
        .map_err(AppError::internal)?
        .into_iter()
        .next()
    {
        return Err(bad(format!("the pack does not hold {sha}")));
    }
    // Protected branches hold whoever pushes here; the owner can write
    // past the gateway.
    let (mut allowed, mut refused) = (Vec::new(), Vec::new());
    for update in updates {
        let refusal = gateway::refusal(&mirror.git, &snap, &update).await;
        match refusal.map_err(AppError::internal)? {
            Some(reason) => refused.push((update.dst, Err(reason))),
            None => allowed.push(update),
        }
    }
    let updates = allowed;
    let mut report = push::push(
        &app.ctx,
        &remote,
        &snap,
        &mirror.git,
        &session.auth(),
        &updates,
    )
    .await
    .map_err(AppError::upstream)?;
    report.extend(refused);
    Ok(Json(PushResults {
        results: report
            .into_iter()
            .map(|(name, res)| RefResult {
                name,
                error: res.err(),
            })
            .collect(),
    }))
}

pub async fn commit(
    State(app): State<Arc<App>>,
    headers: HeaderMap,
    Json(req): Json<CommitRequest>,
) -> Result<Json<CommitResponse>, AppError> {
    let caller = caller(&app, &headers, COMMIT).await?;
    if req.branch.is_empty() || req.message.is_empty() {
        return Err(bad("a commit needs a branch and a message"));
    }
    // Checked before queueing, so a batch holds only commits that may land.
    writable(&app, &caller, &req.repo, &req.name).await?;
    let branch = gateway(&app)?
        .branches
        .lock()
        .unwrap()
        .entry((req.repo.clone(), req.name.clone(), req.branch.clone()))
        .or_default()
        .clone();
    let (reply, answer) = oneshot::channel();
    branch.queue.lock().unwrap().push(Pending {
        spec: req.spec(),
        reply,
    });
    {
        let _turn = branch.turn.lock().await;
        let batch = std::mem::take(&mut *branch.queue.lock().unwrap());
        // Empty when the write before this one took this commit along.
        if !batch.is_empty() {
            land(&app, &caller, &req, batch).await;
        }
    }
    answer
        .await
        .expect("every queued commit is answered")
        .map(Json)
}

/// Builds the batch on the branch and writes it, rebuilding on a branch
/// that moved, and answers every commit in it.
async fn land(app: &App, caller: &str, req: &CommitRequest, batch: Vec<Pending>) {
    match attempt(app, caller, req, &batch).await {
        Ok(results) => {
            for (pending, result) in batch.into_iter().zip(results) {
                let _ = pending.reply.send(result);
            }
        }
        Err(err) => {
            for pending in batch {
                let _ = pending.reply.send(Err(AppError(err.0, err.1.clone())));
            }
        }
    }
}

type Answers = Vec<Result<CommitResponse, AppError>>;

async fn attempt(
    app: &App,
    caller: &str,
    req: &CommitRequest,
    batch: &[Pending],
) -> Result<Answers, AppError> {
    let refname = format!("refs/heads/{}", req.branch);
    let mut round = 0;
    loop {
        let (remote, session, snap) = writable(app, caller, &req.repo, &req.name).await?;
        let mirror = mirror(app, &remote, &snap).await?;
        let start = commit::start(&snap, &batch[0].spec);
        let mut tip = start.clone();
        let mut built: Vec<Result<Committed, String>> = Vec::new();
        for pending in batch {
            if let Err(err) = commit::conflicts(&mirror.git, &pending.spec, tip.as_deref()).await {
                built.push(Err(format!("{err:#}")));
                continue;
            }
            let sha = commit::build(
                &app.ctx,
                &mirror.git,
                tip.as_deref(),
                &pending.spec,
                remote.did(),
            )
            .await
            .map_err(AppError::input)?;
            let parent = tip.clone();
            if let Some(sha) = &sha {
                tip = Some(sha.clone());
            }
            built.push(Ok(Committed {
                sha: tip.clone().unwrap_or_default(),
                parent,
                attempts: 1,
                changed: sha.is_some(),
            }));
        }
        let landed = built.iter().flatten().filter(|c| c.changed).count();
        if tip != start {
            let update = Update {
                src: tip.clone(),
                dst: refname.clone(),
                force: true,
            };
            let report = push::push(
                &app.ctx,
                &remote,
                &snap,
                &mirror.git,
                &session.auth(),
                &[update],
            )
            .await
            .map_err(AppError::upstream)?;
            gateway(app)?.writes.fetch_add(1, Ordering::Relaxed);
            // Forced against the generation read above, so the one refusal
            // is that the branch moved: rebuild on where it went.
            if let Err(reason) = &report[0].1 {
                let gave_up = AppError(StatusCode::CONFLICT, format!("{refname}: {reason}"));
                (round < RETRIES).then_some(()).ok_or(gave_up)?;
                gateway(app)?.rebuilt.fetch_add(1, Ordering::Relaxed);
                round += 1;
                continue;
            }
        }
        let answers = built
            .into_iter()
            .map(|c| match c {
                Ok(c) => Ok(CommitResponse {
                    sha: c.sha,
                    parent: c.parent,
                    changed: c.changed,
                    batched: landed,
                }),
                Err(msg) => Err(AppError(StatusCode::CONFLICT, msg)),
            })
            .collect();
        return Ok(answers);
    }
}