johnandersen777.bsky.social/zkwid

Switch repository

Loading…

mod proxy;

use anyhow::{Context, bail, ensure};
use atproto_identity::key::{KeyData, KeyType, generate_key};
use atproto_oauth::resources::{AuthorizationServer, pds_resources};
use atproto_oauth::workflow::{OAuthClient, OAuthRequest, OAuthRequestState, TokenResponse, generate_state, oauth_complete, oauth_init};
use atproto_testkit::FakePlcDirectory;
use rand::distr::{Alphanumeric, SampleString};
use reqwest::{StatusCode, Url};
use serde_json::{Value, json};
use std::net::SocketAddr;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::sync::Mutex;
use std::time::{Duration, Instant};
use tokio::process::{Child, Command};
use zkwid::lexicon::{Blob, StrongRef};

pub const N_HOST: &str = "zkwid.test";
pub const H_HOST: &str = "h.zkwid.test";
pub const PDS_HOST: &str = "pds.zkwid.test";
pub const PLC_HOST: &str = "plc.zkwid.test";
pub const USER_HANDLE: &str = "alice.pds.zkwid.test";

const PDS_WILDCARD: &str = "*.pds.zkwid.test";

#[derive(Clone, Debug)]
pub struct Account {
    pub did: String,
    pub handle: String,
    pub password: String,
    pub access_jwt: String,
}

pub struct Login {
    pub redirect: String,
    pub tokens: TokenResponse,
    pub dpop_key: KeyData,
}

pub struct Devnet {
    pub dir: PathBuf,
    pub ca_pem: PathBuf,
    pub proxy_url: String,
    pub pds_url: String,
    pub plc_url: String,
    pub user: Account,
    pub pds_local: String,
    pub pds_log: PathBuf,
    routes: proxy::Routes,
    tasks: Mutex<Vec<proxy::Task>>,
    _plc: FakePlcDirectory,
    _pds: Child,
}

impl Devnet {
    pub async fn start() -> anyhow::Result<Devnet> {
        let dir = std::env::temp_dir().join(format!("zkwid-devnet-{}-{}", std::process::id(), random(8).to_lowercase()));
        std::fs::create_dir_all(&dir)?;
        let certs = proxy::Certs::generate(&[N_HOST, H_HOST, PDS_HOST, PDS_WILDCARD, PLC_HOST])?;
        let ca_pem = dir.join("ca.pem");
        std::fs::write(&ca_pem, &certs.ca_pem)?;
        let routes = proxy::Routes::default();
        let (proxy_addr, proxy_task) = proxy::start(routes.clone(), certs.acceptor()?).await?;
        let proxy_url = format!("http://{proxy_addr}");
        unsafe {
            for name in ["NO_PROXY", "no_proxy", "ALL_PROXY", "all_proxy", "HTTP_PROXY", "http_proxy"] {
                std::env::remove_var(name);
            }
            std::env::set_var("HTTPS_PROXY", &proxy_url);
            std::env::set_var("https_proxy", &proxy_url);
            std::env::set_var("SSL_CERT_FILE", &ca_pem);
        }

        let plc = FakePlcDirectory::start().await;
        plc.auto_publish();
        routes.set(PLC_HOST, plc.origin().trim_start_matches("http://").parse()?);

        let (pds, pds_local, pds_log) = spawn_pds(&dir, &proxy_url, &ca_pem, plc.origin()).await?;
        let pds_addr: SocketAddr = pds_local.trim_start_matches("http://").parse()?;
        routes.set(PDS_HOST, pds_addr);
        routes.set(PDS_WILDCARD, pds_addr);

        let mut devnet = Devnet {
            dir,
            ca_pem,
            proxy_url,
            pds_url: format!("https://{PDS_HOST}"),
            plc_url: format!("https://{PLC_HOST}"),
            user: Account { did: String::new(), handle: USER_HANDLE.to_string(), password: random(24), access_jwt: String::new() },
            pds_local,
            pds_log,
            routes,
            tasks: Mutex::new(vec![proxy_task]),
            _plc: plc,
            _pds: pds,
        };
        devnet.create_user().await?;
        devnet.verify_identity().await?;
        Ok(devnet)
    }

    pub fn http(&self) -> reqwest::Client {
        reqwest::Client::builder()
            .redirect(reqwest::redirect::Policy::none())
            .build()
            .expect("reqwest client")
    }

    pub async fn route(&self, host: &str, backend: SocketAddr) -> anyhow::Result<()> {
        self.routes.set(host, backend);
        Ok(())
    }

    pub async fn serve(&self, host: &str, app: axum::Router) -> anyhow::Result<SocketAddr> {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
        let addr = listener.local_addr()?;
        let task = tokio::spawn(async move {
            axum::serve(listener, app).await.ok();
        });
        self.tasks.lock().unwrap().push(proxy::Task(task));
        self.route(host, addr).await?;
        Ok(addr)
    }

    pub async fn upload_blob(&self, bytes: &[u8], mime: &str) -> anyhow::Result<Blob> {
        let response = self
            .http()
            .post(self.xrpc("com.atproto.repo.uploadBlob"))
            .bearer_auth(&self.user.access_jwt)
            .header(reqwest::header::CONTENT_TYPE, mime)
            .body(bytes.to_vec())
            .send()
            .await?;
        let body = json_ok(response, "uploadBlob").await?;
        Ok(serde_json::from_value(body["blob"].clone())?)
    }

    pub async fn put_record(&self, collection: &str, rkey: &str, value: &Value) -> anyhow::Result<StrongRef> {
        let body = json!({ "repo": self.user.did, "collection": collection, "rkey": rkey, "record": value, "validate": false });
        let response = self.http().post(self.xrpc("com.atproto.repo.putRecord")).bearer_auth(&self.user.access_jwt).json(&body).send().await?;
        strong_ref(&json_ok(response, "putRecord").await?)
    }

    pub async fn get_record(&self, collection: &str, rkey: &str) -> anyhow::Result<Option<(Value, String)>> {
        let response = self
            .http()
            .get(self.xrpc("com.atproto.repo.getRecord"))
            .query(&[("repo", self.user.did.as_str()), ("collection", collection), ("rkey", rkey)])
            .send()
            .await?;
        if response.status() == StatusCode::BAD_REQUEST || response.status() == StatusCode::NOT_FOUND {
            let body: Value = response.json().await.unwrap_or(Value::Null);
            if body["error"] == "RecordNotFound" {
                return Ok(None);
            }
            bail!("getRecord failed: {body}");
        }
        let body = json_ok(response, "getRecord").await?;
        Ok(Some((body["value"].clone(), text(&body, "cid")?)))
    }

    pub async fn list_records(&self, collection: &str) -> anyhow::Result<Vec<(String, Value)>> {
        let mut records = Vec::new();
        let mut cursor: Option<String> = None;
        loop {
            let mut query = vec![("repo", self.user.did.clone()), ("collection", collection.to_string()), ("limit", "100".to_string())];
            if let Some(cursor) = &cursor {
                query.push(("cursor", cursor.clone()));
            }
            let response = self.http().get(self.xrpc("com.atproto.repo.listRecords")).query(&query).send().await?;
            let body = json_ok(response, "listRecords").await?;
            let page = body["records"].as_array().cloned().unwrap_or_default();
            for record in &page {
                records.push((text(record, "uri")?, record["value"].clone()));
            }
            cursor = body["cursor"].as_str().map(str::to_string);
            if page.is_empty() || cursor.is_none() {
                return Ok(records);
            }
        }
    }

    pub async fn delete_record(&self, collection: &str, rkey: &str) -> anyhow::Result<()> {
        let body = json!({ "repo": self.user.did, "collection": collection, "rkey": rkey });
        let response = self.http().post(self.xrpc("com.atproto.repo.deleteRecord")).bearer_auth(&self.user.access_jwt).json(&body).send().await?;
        json_ok(response, "deleteRecord").await?;
        Ok(())
    }

    pub async fn consent(&self, authorize_url: &str) -> anyhow::Result<String> {
        let mut endpoint = Url::parse(authorize_url)?;
        let request_uri = endpoint
            .query_pairs()
            .find(|(key, _)| key == "request_uri")
            .map(|(_, value)| value.into_owned())
            .context("authorize URL has no request_uri")?;
        endpoint.set_query(None);
        let body = json!({ "request_uri": request_uri, "identifier": self.user.handle, "password": self.user.password, "approve": true });
        let response = self.http().post(endpoint).json(&body).send().await?;
        let granted = json_ok(response, "consent").await?;
        let mut redirect = Url::parse(&text(&granted, "redirect_uri")?)?;
        let params = [("code", text(&granted, "code")?), ("state", text(&granted, "state")?), ("iss", text(&granted, "iss")?)];
        if granted["response_mode"] == "fragment" {
            let fragment = url::form_urlencoded::Serializer::new(String::new()).extend_pairs(params).finish();
            redirect.set_fragment(Some(&fragment));
        } else {
            redirect.query_pairs_mut().extend_pairs(params);
        }
        Ok(redirect.into())
    }

    pub async fn authorization_server(&self) -> anyhow::Result<AuthorizationServer> {
        let (_, server) = pds_resources(&self.http(), &self.pds_url).await?;
        Ok(server)
    }

    pub async fn oauth_login(&self, client: &OAuthClient, scope: &str) -> anyhow::Result<Login> {
        let http = self.http();
        let server = self.authorization_server().await?;
        let dpop_key = generate_key(KeyType::P256Private)?;
        let (pkce_verifier, code_challenge) = atproto_oauth::pkce::generate();
        let state = OAuthRequestState { state: generate_state(), code_challenge, scope: scope.to_string() };
        let par = oauth_init(&http, client, &dpop_key, Some(&self.user.handle), &server, &state).await?;
        let mut authorize = Url::parse(&server.authorization_endpoint)?;
        authorize.query_pairs_mut().append_pair("client_id", &client.client_id).append_pair("request_uri", &par.request_uri);
        let redirect = self.consent(authorize.as_str()).await?;
        let callback = Url::parse(&redirect)?;
        let param = |name: &str| callback.query_pairs().find(|(key, _)| key == name).map(|(_, value)| value.into_owned());
        ensure!(param("state").as_deref() == Some(state.state.as_str()), "callback state mismatch: {redirect}");
        let code = param("code").with_context(|| format!("callback has no code: {redirect}"))?;
        let now = chrono::Utc::now();
        let request = OAuthRequest {
            oauth_state: state.state.clone(),
            subject: self.user.did.clone(),
            issuer: server.issuer.clone(),
            pkce_verifier,
            signing_public_key: String::new(),
            dpop_private_key: String::new(),
            created_at: now,
            expires_at: now + chrono::Duration::minutes(10),
        };
        let tokens = oauth_complete(&http, client, &dpop_key, &code, &self.user.did, param("iss").as_deref(), &request, &server).await?;
        Ok(Login { redirect, tokens, dpop_key })
    }

    fn xrpc(&self, nsid: &str) -> String {
        format!("{}/xrpc/{nsid}", self.pds_url)
    }

    async fn create_user(&mut self) -> anyhow::Result<()> {
        let body = json!({ "handle": self.user.handle, "password": self.user.password, "email": "alice@zkwid.test" });
        let response = self.http().post(self.xrpc("com.atproto.server.createAccount")).json(&body).send().await?;
        let created = json_ok(response, "createAccount").await?;
        self.user.did = text(&created, "did")?;
        self.user.access_jwt = text(&created, "accessJwt")?;
        ensure!(self.user.did.starts_with("did:plc:"), "account DID is not did:plc: {}", self.user.did);
        Ok(())
    }

    async fn verify_identity(&self) -> anyhow::Result<()> {
        let http = self.http();
        let well_known = format!("https://{}/.well-known/atproto-did", self.user.handle);
        let response = http.get(&well_known).send().await.with_context(|| format!("GET {well_known}"))?;
        ensure!(response.status().is_success(), "GET {well_known}: {}", response.status());
        let did = response.text().await?;
        ensure!(did.trim() == self.user.did, "{well_known} answered {did}, expected {}", self.user.did);
        let document_url = format!("{}/{}", self.plc_url, self.user.did);
        let document = json_ok(http.get(&document_url).send().await?, "PLC document").await?;
        ensure!(document["id"] == self.user.did.as_str(), "PLC document for {} has id {}", self.user.did, document["id"]);
        let endpoint = document["service"]
            .as_array()
            .and_then(|services| services.iter().find(|service| service["id"].as_str().is_some_and(|id| id.ends_with("#atproto_pds"))))
            .and_then(|service| service["serviceEndpoint"].as_str());
        ensure!(endpoint == Some(self.pds_url.as_str()), "PLC document for {} names PDS {endpoint:?}", self.user.did);
        Ok(())
    }
}

impl Drop for Devnet {
    fn drop(&mut self) {
        self._pds.start_kill().ok();
        self.tasks.lock().unwrap().clear();
        std::fs::remove_dir_all(&self.dir).ok();
    }
}

fn random(len: usize) -> String {
    Alphanumeric.sample_string(&mut rand::rng(), len)
}

fn text(value: &Value, field: &str) -> anyhow::Result<String> {
    value[field].as_str().map(str::to_string).with_context(|| format!("missing {field} in {value}"))
}

fn strong_ref(value: &Value) -> anyhow::Result<StrongRef> {
    Ok(StrongRef { uri: text(value, "uri")?, cid: text(value, "cid")? })
}

async fn json_ok(response: reqwest::Response, what: &str) -> anyhow::Result<Value> {
    let status = response.status();
    let body = response.text().await?;
    ensure!(status.is_success(), "{what} failed ({status}): {body}");
    serde_json::from_str(&body).with_context(|| format!("{what} returned non-JSON: {body}"))
}

fn pds_binary() -> anyhow::Result<PathBuf> {
    let path = std::env::var_os("ZKWID_PDS_BIN")
        .map(PathBuf::from)
        .unwrap_or_else(|| Path::new(env!("CARGO_MANIFEST_DIR")).join("../../../atproto-crates/target/debug/pds"));
    ensure!(
        path.is_file(),
        "PDS binary not found at {}. Build it with `cargo build -p atproto-pds --bin pds --features clap` in ../atproto-crates, or set ZKWID_PDS_BIN",
        path.display()
    );
    Ok(path)
}

async fn spawn_pds(dir: &Path, proxy_url: &str, ca_pem: &Path, plc_origin: &str) -> anyhow::Result<(Child, String, PathBuf)> {
    let binary = pds_binary()?;
    let data = dir.join("pds");
    std::fs::create_dir_all(&data)?;
    let port = std::net::TcpListener::bind("127.0.0.1:0")?.local_addr()?.port();
    let log_path = dir.join("pds.log");
    let log = std::fs::File::create(&log_path)?;
    let mut child = Command::new(&binary)
        .env_clear()
        .env("PATH", std::env::var_os("PATH").unwrap_or_default())
        .env("HOME", std::env::var_os("HOME").unwrap_or_else(|| dir.into()))
        .env("RUST_LOG", std::env::var("ZKWID_PDS_LOG").unwrap_or_else(|_| "warn,atproto_pds=info".into()))
        .env("PDS_BIND", "127.0.0.1")
        .env("PDS_PORT", port.to_string())
        .env("PDS_DATA_DIRECTORY", &data)
        .env("PDS_SERVICE_DID", format!("did:web:{PDS_HOST}"))
        .env("PDS_HOSTNAME", PDS_HOST)
        .env("PDS_SERVICE_HANDLE_DOMAINS", format!(".{PDS_HOST}"))
        .env("PDS_DID_PLC_URL", plc_origin)
        .env("PDS_JWT_SECRET", random(48))
        .env("PDS_ADMIN_PASSWORD", random(32))
        .env("PDS_ALLOW_DEV_DEFAULTS", "true")
        .env("PDS_PRODUCTION", "false")
        .env("PDS_INVITE_REQUIRED", "false")
        .env("PDS_DPOP_NONCE_REQUIRED", "true")
        .env("PDS_TRUSTED_PROXY_HOPS", "1")
        .env("PDS_RATE_LIMIT", "1000000")
        .env("PDS_RATE_LIMIT_AUTH", "1000000")
        .env("HTTPS_PROXY", proxy_url)
        .env("https_proxy", proxy_url)
        .env("SSL_CERT_FILE", ca_pem)
        .stdin(Stdio::null())
        .stdout(log.try_clone()?)
        .stderr(log)
        .kill_on_drop(true)
        .spawn()
        .with_context(|| format!("spawning PDS {}", binary.display()))?;
    let local = format!("http://127.0.0.1:{port}");
    let health = format!("{local}/xrpc/_health");
    let http = reqwest::Client::builder().no_proxy().timeout(Duration::from_secs(2)).build()?;
    let deadline = Instant::now() + Duration::from_secs(30);
    loop {
        if let Some(status) = child.try_wait()? {
            bail!("PDS {} exited ({status}) before becoming ready; log {}:\n{}", binary.display(), log_path.display(), tail(&log_path));
        }
        if let Ok(response) = http.get(&health).send().await
            && response.status().is_success()
        {
            return Ok((child, local, log_path));
        }
        if Instant::now() > deadline {
            bail!("PDS {} not healthy at {health} within 30s; log {}:\n{}", binary.display(), log_path.display(), tail(&log_path));
        }
        tokio::time::sleep(Duration::from_millis(100)).await;
    }
}

fn tail(path: &Path) -> String {
    let log = std::fs::read_to_string(path).unwrap_or_default();
    let lines: Vec<&str> = log.lines().collect();
    lines[lines.len().saturating_sub(40)..].join("\n")
}