johnandersen777.bsky.social/zkwid

Switch repository

Loading…

use crate::AppState;
use crate::atp;
use crate::error::XrpcError;
use crate::store::{N_CLIENT, Session};
use atproto_client::client::{
    CallLimits, DPoPAuth, DpopBody, XrpcResponse, dpop_call_with_timeout,
};
use atproto_identity::key::{KeyType, generate_key, identify_key, to_public};
use atproto_identity::model::Document;
use atproto_oauth::refresh::{RefreshCoordinator, RefreshFailure, RefreshOutcome};
use atproto_oauth::resources::{oauth_authorization_server, pds_resources};
use atproto_oauth::workflow::{
    OAuthClient, OAuthRequest, OAuthRequestState, TokenResponse, generate_state, oauth_complete,
    oauth_init, oauth_refresh,
};
use axum::extract::{Path, Query, State};
use axum::http::{HeaderMap, Method, StatusCode, header};
use axum::response::{IntoResponse, Response};
use axum::routing::get;
use axum::{Json, Router};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use std::sync::LazyLock;
use std::time::Duration;

pub const N_SCOPE: &str = "atproto repo:at.zkwid.key";

const FORWARDED: [&str; 3] = ["atproto-proxy", "accept", "accept-language"];

const DROPPED: [&str; 11] = [
    "connection",
    "keep-alive",
    "proxy-authenticate",
    "proxy-authorization",
    "proxy-connection",
    "te",
    "trailer",
    "transfer-encoding",
    "upgrade",
    "content-length",
    "dpop-nonce",
];

static REFRESH: LazyLock<RefreshCoordinator> = LazyLock::new(RefreshCoordinator::new);

pub struct Upstream {
    pub status: StatusCode,
    pub headers: HeaderMap,
    pub body: Vec<u8>,
}

struct ClientSpec {
    client: String,
    client_id: String,
    redirect_uri: String,
    scope: String,
}

#[derive(Serialize, Deserialize)]
struct Pending {
    client: String,
    client_id: String,
    redirect_uri: String,
    did: String,
    pds: String,
    issuer: String,
    scope: String,
    pkce_verifier: String,
    signing_public_key: String,
    dpop_key: String,
    created_at: i64,
    expires_at: i64,
}

#[derive(Deserialize)]
struct LoginQuery {
    handle: Option<String>,
}

#[derive(Deserialize)]
struct CallbackQuery {
    code: Option<String>,
    state: Option<String>,
    iss: Option<String>,
    error: Option<String>,
    error_description: Option<String>,
}

pub fn routes() -> Router<AppState> {
    Router::new()
        .route("/oauth-client-metadata.json", get(n_metadata))
        .route("/login", get(n_login))
        .route("/callback", get(n_callback))
        .route(
            "/{handle}/{rkey}/oauth-client-metadata.json",
            get(workload_metadata),
        )
        .route("/{handle}/{rkey}/login", get(workload_login))
        .route("/{handle}/{rkey}/callback", get(workload_callback))
}

fn n_spec(state: &AppState) -> ClientSpec {
    let base = &state.config.base_url;
    ClientSpec {
        client: N_CLIENT.into(),
        client_id: zkwid::client::n_client_id(base),
        redirect_uri: zkwid::client::n_redirect_uri(base),
        scope: N_SCOPE.into(),
    }
}

fn workload_spec(state: &AppState, handle: &str, rkey: &str, scope: String) -> ClientSpec {
    let base = &state.config.base_url;
    ClientSpec {
        client: rkey.into(),
        client_id: zkwid::client::client_id(base, handle, rkey),
        redirect_uri: zkwid::client::redirect_uri(base, handle, rkey),
        scope,
    }
}

async fn workload(
    state: &AppState,
    handle: &str,
    rkey: &str,
) -> anyhow::Result<(Document, ClientSpec)> {
    let missing = |message: String| XrpcError::not_found("WorkloadNotFound", message);
    let document = atp::resolve(state, handle)
        .await
        .map_err(|error| missing(format!("{handle}: {error:#}")))?;
    let uri = format!("at://{}/{}/{rkey}", document.id, zkwid::nsid::WORKLOAD);
    let (value, _) = atp::find_record(state, &uri)
        .await
        .map_err(|error| XrpcError::upstream(format!("{error:#}")))?
        .ok_or_else(|| missing(format!("{uri} does not exist")))?;
    let record: zkwid::lexicon::Workload =
        serde_json::from_value(value).map_err(|error| missing(format!("{uri}: {error}")))?;
    let scope = zkwid::scope::client_scope(&record.scope)
        .map_err(|error| missing(format!("{uri}: {error:#}")))?;
    Ok((document, workload_spec(state, handle, rkey, scope)))
}

fn oauth_client(state: &AppState, spec: &ClientSpec) -> anyhow::Result<OAuthClient> {
    Ok(OAuthClient {
        redirect_uri: spec.redirect_uri.clone(),
        client_id: spec.client_id.clone(),
        private_signing_key_data: state.store.client_key(&spec.client_id)?,
    })
}

fn metadata(state: &AppState, spec: &ClientSpec, client_name: &str) -> anyhow::Result<Json<Value>> {
    let key = state.store.client_key(&spec.client_id)?;
    Ok(Json(zkwid::client::metadata(
        &spec.client_id,
        &spec.redirect_uri,
        &spec.scope,
        client_name,
        zkwid::client::public_jwk(&key)?,
    )))
}

fn first_pds(document: &Document) -> anyhow::Result<String> {
    Ok(document
        .pds_endpoints_validated()
        .first()
        .ok_or_else(|| XrpcError::invalid_request(format!("{} names no usable PDS", document.id)))?
        .trim_end_matches('/')
        .to_string())
}

fn upstream_error(error: impl std::fmt::Display) -> XrpcError {
    XrpcError::upstream(error.to_string())
}

async fn n_metadata(State(state): State<AppState>) -> Result<Json<Value>, XrpcError> {
    Ok(metadata(&state, &n_spec(&state), "zkwid.at")?)
}

async fn workload_metadata(
    State(state): State<AppState>,
    Path((handle, rkey)): Path<(String, String)>,
) -> Result<Json<Value>, XrpcError> {
    let (_, spec) = workload(&state, &handle, &rkey).await?;
    Ok(metadata(&state, &spec, &format!("zkwid.at {rkey}"))?)
}

async fn n_login(
    State(state): State<AppState>,
    Query(query): Query<LoginQuery>,
) -> Result<Response, XrpcError> {
    let handle = query
        .handle
        .filter(|handle| !handle.is_empty())
        .ok_or_else(|| XrpcError::invalid_request("missing handle"))?;
    let document = atp::resolve(&state, &handle)
        .await
        .map_err(|error| XrpcError::invalid_request(format!("{handle}: {error:#}")))?;
    Ok(start(&state, &handle, &document, &n_spec(&state)).await?)
}

async fn workload_login(
    State(state): State<AppState>,
    Path((handle, rkey)): Path<(String, String)>,
) -> Result<Response, XrpcError> {
    let (document, spec) = workload(&state, &handle, &rkey).await?;
    Ok(start(&state, &handle, &document, &spec).await?)
}

async fn n_callback(
    State(state): State<AppState>,
    Query(query): Query<CallbackQuery>,
) -> Result<Json<Value>, XrpcError> {
    let spec = n_spec(&state);
    Ok(complete(&state, &spec.redirect_uri, query).await?)
}

async fn workload_callback(
    State(state): State<AppState>,
    Path((handle, rkey)): Path<(String, String)>,
    Query(query): Query<CallbackQuery>,
) -> Result<Json<Value>, XrpcError> {
    let spec = workload_spec(&state, &handle, &rkey, String::new());
    Ok(complete(&state, &spec.redirect_uri, query).await?)
}

async fn start(
    state: &AppState,
    handle: &str,
    document: &Document,
    spec: &ClientSpec,
) -> anyhow::Result<Response> {
    let pds = first_pds(document)?;
    let (_, server) = pds_resources(&state.http, &pds)
        .await
        .map_err(upstream_error)?;
    let client = oauth_client(state, spec)?;
    let dpop_key = generate_key(KeyType::P256Private)?;
    let (pkce_verifier, code_challenge) = atproto_oauth::pkce::generate();
    let request_state = OAuthRequestState {
        state: generate_state(),
        code_challenge,
        scope: spec.scope.clone(),
    };
    let par = oauth_init(
        &state.http,
        &client,
        &dpop_key,
        Some(handle),
        &server,
        &request_state,
    )
    .await
    .map_err(upstream_error)?;
    let now = Utc::now().timestamp();
    let pending = Pending {
        client: spec.client.clone(),
        client_id: spec.client_id.clone(),
        redirect_uri: spec.redirect_uri.clone(),
        did: document.id.clone(),
        pds,
        issuer: server.issuer.clone(),
        scope: spec.scope.clone(),
        pkce_verifier,
        signing_public_key: to_public(&client.private_signing_key_data)?.to_string(),
        dpop_key: dpop_key.to_string(),
        created_at: now,
        expires_at: now + i64::try_from(par.expires_in).unwrap_or(0),
    };
    state
        .store
        .put_request(&request_state.state, &serde_json::to_string(&pending)?)?;
    let location = reqwest::Url::parse_with_params(
        &server.authorization_endpoint,
        [
            ("client_id", spec.client_id.as_str()),
            ("request_uri", par.request_uri.as_str()),
        ],
    )?;
    Ok((
        StatusCode::FOUND,
        [(header::LOCATION, location.to_string())],
    )
        .into_response())
}

async fn complete(
    state: &AppState,
    redirect_uri: &str,
    query: CallbackQuery,
) -> anyhow::Result<Json<Value>> {
    if let Some(error) = query.error {
        let description = query.error_description.unwrap_or_default();
        return Err(XrpcError::invalid_request(format!(
            "authorization failed: {error} {description}"
        ))
        .into());
    }
    let code = query
        .code
        .ok_or_else(|| XrpcError::invalid_request("missing code"))?;
    let oauth_state = query
        .state
        .ok_or_else(|| XrpcError::invalid_request("missing state"))?;
    let stored = state
        .store
        .take_request(&oauth_state)?
        .ok_or_else(|| XrpcError::invalid_request("unknown or already used state"))?;
    let pending: Pending = serde_json::from_str(&stored)?;
    if pending.redirect_uri != redirect_uri {
        return Err(XrpcError::invalid_request("state was issued for another client").into());
    }
    let server = oauth_authorization_server(&state.http, &pending.issuer)
        .await
        .map_err(upstream_error)?;
    let dpop_key = identify_key(&pending.dpop_key)?;
    let client = OAuthClient {
        redirect_uri: pending.redirect_uri.clone(),
        client_id: pending.client_id.clone(),
        private_signing_key_data: state.store.client_key(&pending.client_id)?,
    };
    let request = OAuthRequest {
        oauth_state,
        subject: pending.did.clone(),
        issuer: pending.issuer.clone(),
        pkce_verifier: pending.pkce_verifier,
        signing_public_key: pending.signing_public_key,
        dpop_private_key: pending.dpop_key.clone(),
        created_at: DateTime::from_timestamp(pending.created_at, 0).unwrap_or_default(),
        expires_at: DateTime::from_timestamp(pending.expires_at, 0).unwrap_or_default(),
    };
    let tokens = oauth_complete(
        &state.http,
        &client,
        &dpop_key,
        &code,
        &pending.did,
        query.iss.as_deref(),
        &request,
        &server,
    )
    .await
    .map_err(|error| XrpcError::invalid_request(format!("token exchange failed: {error}")))?;
    let session = Session {
        did: pending.did.clone(),
        client: pending.client.clone(),
        access_token: tokens.access_token.clone(),
        refresh_token: tokens.refresh_token.clone().unwrap_or_default(),
        dpop_key: pending.dpop_key,
        pds: pending.pds,
        issuer: pending.issuer,
        scope: tokens.scope.clone(),
        expires_at: expires_at(&tokens),
    };
    state.store.put_session(&session)?;
    REFRESH.forget(&refresh_key(state, &session));
    Ok(Json(
        json!({ "did": session.did, "client": session.client }),
    ))
}

fn expires_at(tokens: &TokenResponse) -> i64 {
    Utc::now().timestamp() + i64::from(tokens.expires_in)
}

fn refresh_key(state: &AppState, session: &Session) -> String {
    format!(
        "{}\n{}\n{}",
        state.config.db_path.display(),
        session.did,
        session.client
    )
}

async fn refresh(state: &AppState, session: Session) -> anyhow::Result<Session> {
    let current = state.store.session(&session.did, &session.client)?;
    if let Some(current) = current
        && current.access_token != session.access_token
    {
        return Ok(current);
    }
    let document = atp::resolve(state, &session.did)
        .await
        .map_err(upstream_error)?;
    let spec = if session.client == N_CLIENT {
        n_spec(state)
    } else {
        let handle = document.handles().ok_or_else(|| {
            XrpcError::forbidden("NotAuthorized", format!("{} has no handle", session.did))
        })?;
        workload_spec(state, handle, &session.client, String::new())
    };
    let client = oauth_client(state, &spec)?;
    let key = refresh_key(state, &session);
    let outcome = REFRESH
        .refresh(&key, async {
            let latest = state
                .store
                .session(&session.did, &session.client)?
                .unwrap_or_else(|| session.clone());
            let dpop_key = identify_key(&latest.dpop_key)?;
            let tokens = oauth_refresh(
                &state.http,
                &client,
                &dpop_key,
                &latest.refresh_token,
                &document,
            )
            .await?;
            state.store.put_session(&Session {
                access_token: tokens.access_token.clone(),
                refresh_token: tokens
                    .refresh_token
                    .clone()
                    .unwrap_or_else(|| latest.refresh_token.clone()),
                scope: tokens.scope.clone(),
                expires_at: expires_at(&tokens),
                ..latest
            })?;
            Ok::<_, anyhow::Error>(tokens)
        })
        .await;
    let revoked = || {
        XrpcError::forbidden(
            "NotAuthorized",
            format!("OAuth session for {} was revoked", session.did),
        )
    };
    match outcome {
        Ok(RefreshOutcome::Refreshed(_) | RefreshOutcome::AlreadyFresh) => Ok(state
            .store
            .session(&session.did, &session.client)?
            .ok_or_else(|| XrpcError::forbidden("NotAuthorized", "OAuth session disappeared"))?),
        Ok(RefreshOutcome::Revoked) => Err(revoked().into()),
        Ok(_) => Err(XrpcError::upstream("token refresh is backing off").into()),
        Err(error) if error.is_grant_terminal() => Err(revoked().into()),
        Err(error) => Err(XrpcError::upstream(format!("token refresh failed: {error:#}")).into()),
    }
}

fn token_rejected(response: &XrpcResponse) -> bool {
    if response.status != StatusCode::UNAUTHORIZED {
        return false;
    }
    let challenge = response
        .headers
        .get(header::WWW_AUTHENTICATE)
        .and_then(|value| value.to_str().ok())
        .unwrap_or_default()
        .to_ascii_lowercase();
    let (error, _) = response.xrpc_error_fields();
    challenge.contains("invalid_token")
        || matches!(
            error.as_str(),
            "invalid_token" | "InvalidToken" | "ExpiredToken"
        )
}

async fn send(
    state: &AppState,
    session: &Session,
    method: &Method,
    path_and_query: &str,
    body: Option<&(Vec<u8>, String)>,
    headers: &HeaderMap,
) -> anyhow::Result<XrpcResponse> {
    let auth = DPoPAuth {
        dpop_private_key_data: identify_key(&session.dpop_key)?,
        oauth_access_token: session.access_token.clone(),
    };
    let url = format!("{}{path_and_query}", session.pds.trim_end_matches('/'));
    let body = body.map(|(data, content_type)| DpopBody::Bytes {
        content_type,
        data: data.clone(),
    });
    Ok(dpop_call_with_timeout(
        &state.http,
        &auth,
        method.clone(),
        &url,
        body,
        headers,
        CallLimits::new(Duration::from_secs(60), 64 * 1024 * 1024),
    )
    .await
    .map_err(upstream_error)?)
}

pub async fn pds_call(
    state: &AppState,
    did: &str,
    client: &str,
    method: Method,
    path_and_query: &str,
    body: Option<(Vec<u8>, String)>,
    headers: &HeaderMap,
) -> anyhow::Result<Upstream> {
    if !path_and_query.starts_with('/') {
        return Err(XrpcError::invalid_request(format!("bad path: {path_and_query}")).into());
    }
    let mut session = state.store.session(did, client)?.ok_or_else(|| {
        XrpcError::forbidden(
            "NotAuthorized",
            format!("no OAuth session for {did} client {client:?}"),
        )
    })?;
    if session.expires_at <= Utc::now().timestamp() + 30 {
        session = refresh(state, session).await?;
    }
    let mut forwarded = HeaderMap::new();
    for name in FORWARDED {
        for value in headers.get_all(name) {
            forwarded.append(name, value.clone());
        }
    }
    let mut response = send(
        state,
        &session,
        &method,
        path_and_query,
        body.as_ref(),
        &forwarded,
    )
    .await?;
    if token_rejected(&response) {
        session = refresh(state, session).await?;
        response = send(
            state,
            &session,
            &method,
            path_and_query,
            body.as_ref(),
            &forwarded,
        )
        .await?;
    }
    let mut headers = response.headers;
    for name in DROPPED {
        headers.remove(name);
    }
    Ok(Upstream {
        status: response.status,
        headers,
        body: response.bytes.to_vec(),
    })
}