~ knotless Search/
↑↓ move↵ open⇧↵ all results

johnandersen777.bsky.social/atlibs

0
0
main
186 lines · 9.0 KB Raw History
// A fake PLC + handle resolver + PDS for the CLI's tests (no network): records, blobs, sessions.

use std::collections::HashMap;
use std::path::Path;
use std::sync::{Arc, Mutex};

use atlibs_host::crypto::base32;
use atlibs_host::host::{Fetch, Response};
use atlibs_host::store::Store;
use atlibs_host::BoxFut;
use serde_json::{Value, json};
use sha2::{Digest, Sha256};

use crate::Cli;
use crate::registry::{ABI, CLIENT, MODULE, Net};

pub const PDS: &str = "https://pds.test";

#[derive(Default)]
struct State {
    /// at:// uri -> (cid, value)
    records: HashMap<String, (String, Value)>,
    blobs: HashMap<String, Vec<u8>>,
    /// requests seen: "METHOD nsid"
    log: Vec<String>,
}

#[derive(Clone, Default)]
pub struct Pds(Arc<Mutex<State>>);

pub fn did_of(handle: &str) -> String {
    format!("did:plc:{}", handle.replace('.', ""))
}

pub fn raw_cid(bytes: &[u8]) -> String {
    format!("b{}", base32(&[&[0x01, 0x55, 0x12, 0x20][..], &Sha256::digest(bytes)].concat()))
}

fn record_cid(value: &Value) -> String {
    format!("b{}", base32(&[&[0x01, 0x71, 0x12, 0x20][..], &Sha256::digest(value.to_string())].concat()))
}

impl Pds {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn cli(&self, dir: &Path) -> Cli {
        let fetch: Arc<dyn Fetch> = Arc::new(self.clone());
        let state_dir = dir.join(".state");
        Cli {
            net: Net { fetch: fetch.clone(), plc: "https://plc.test".into(), resolver: PDS.into(), allow_private: false },
            fetch,
            store: Store::open(&state_dir, "atlibs-server").unwrap().into(),
            state_dir,
        }
    }

    pub fn log(&self) -> Vec<String> {
        self.0.lock().unwrap().log.clone()
    }

    pub fn get(&self, uri: &str) -> Option<Value> {
        self.0.lock().unwrap().records.get(uri).map(|(_, v)| v.clone())
    }

    pub fn put(&self, handle: &str, collection: &str, rkey: &str, value: Value) -> (String, String) {
        let uri = format!("at://{}/{collection}/{rkey}", did_of(handle));
        let cid = record_cid(&value);
        self.0.lock().unwrap().records.insert(uri.clone(), (cid.clone(), value));
        (uri, cid)
    }

    /// A module record as `atlibs publish` writes it: at `rkey` and at its version (registry::version_rkey).
    pub fn put_module(&self, handle: &str, rkey: &str, value: Value) -> (String, String) {
        let sha = atlibs_host::host::blob_sha256(value.pointer("/module/ref/$link").and_then(Value::as_str)).unwrap();
        let version = atlibs_host::registry::version_rkey(rkey, &atlibs_host::run::pin_of_sha256(&sha, &value)).unwrap();
        self.put(handle, MODULE, &version, value.clone());
        self.put(handle, MODULE, rkey, value)
    }

    /// A composition record as `atlibs publish` writes it: at `rkey` and at its version (registry::composition_version).
    pub fn put_composition(&self, handle: &str, rkey: &str, value: Value) -> (String, String) {
        let version = atlibs_host::registry::composition_version(rkey, &record_cid(&value)).unwrap();
        self.put(handle, crate::registry::COMPOSITION, &version, value.clone());
        self.put(handle, crate::registry::COMPOSITION, rkey, value)
    }

    pub fn delete(&self, uri: &str) {
        self.0.lock().unwrap().records.remove(uri);
    }

    pub fn blob(&self, bytes: &[u8]) -> Value {
        let cid = raw_cid(bytes);
        self.0.lock().unwrap().blobs.insert(cid.clone(), bytes.to_vec());
        json!({ "$type": "blob", "ref": { "$link": cid }, "mimeType": "application/wasm", "size": bytes.len() })
    }

    /// A module record (and its session's clientMetadata record when `session`); its at:// URI.
    pub fn publish_module(&self, handle: &str, rkey: &str, wasm: &[u8], session: Option<&str>) -> String {
        let mut value = json!({ "$type": MODULE, "name": rkey, "module": self.blob(wasm), "abi": ABI, "actions": [], "calls": ["com.atproto.repo.getRecord"], "createdAt": "2026-11-08T00:00:00.000Z", "id": "/etc/passwd", "owner": "someone-else", "uri": "at://did:plc:other/x/y" });
        if let Some(s) = session {
            let (uri, cid) = self.put(handle, CLIENT, s, json!({ "$type": CLIENT, "name": s, "clientMetadata": { "client_id": "https://tpl.test/c.json", "scope": "atproto repo:x.y.z?action=create" }, "createdAt": "x" }));
            value["session"] = json!({ "uri": uri, "cid": cid });
        }
        self.put_module(handle, rkey, value).0
    }

    pub fn corrupt_blob(&self, module_uri: &str) {
        let cid = self.get(module_uri).unwrap()["module"]["ref"]["$link"].as_str().unwrap().to_string();
        self.0.lock().unwrap().blobs.insert(cid, b"tampered".to_vec());
    }

    fn answer(&self, method: &str, url: &str, body: Option<Vec<u8>>) -> (u16, Vec<u8>) {
        let parsed = url::Url::parse(url).unwrap();
        let q: HashMap<String, String> = parsed.query_pairs().into_owned().collect();
        let json_body = || body.as_deref().and_then(|b| serde_json::from_slice::<Value>(b).ok()).unwrap_or(Value::Null);
        let ok = |v: Value| (200, v.to_string().into_bytes());
        let not_found = |e: &str| (400, json!({ "error": e, "message": "not found" }).to_string().into_bytes());
        if parsed.host_str() == Some("plc.test") {
            let did = parsed.path().trim_start_matches('/');
            let handle = did.trim_start_matches("did:plc:").replace("test", ".test");
            // mallory.test resolves to alice's DID, whose document does not claim it
            return ok(json!({ "id": did, "alsoKnownAs": [format!("at://{handle}")], "service": [{ "id": "#atproto_pds", "type": "AtprotoPersonalDataServer", "serviceEndpoint": PDS }] }));
        }
        let nsid = parsed.path().trim_start_matches("/xrpc/").to_string();
        self.0.lock().unwrap().log.push(format!("{method} {nsid}"));
        let mut state = self.0.lock().unwrap();
        match nsid.as_str() {
            "com.atproto.identity.resolveHandle" => {
                let handle = q["handle"].as_str();
                ok(json!({ "did": did_of(if handle == "mallory.test" { "alice.test" } else { handle }) }))
            }
            "com.atproto.repo.getRecord" => {
                let uri = format!("at://{}/{}/{}", q["repo"], q["collection"], q["rkey"]);
                match state.records.get(&uri) {
                    Some((cid, value)) => ok(json!({ "uri": uri, "cid": cid, "value": value })),
                    None => not_found("RecordNotFound"),
                }
            }
            "com.atproto.repo.listRecords" => {
                let prefix = format!("at://{}/{}/", q["repo"], q["collection"]);
                let records: Vec<Value> = state.records.iter().filter(|(u, _)| u.starts_with(&prefix)).map(|(u, (c, v))| json!({ "uri": u, "cid": c, "value": v })).collect();
                ok(json!({ "records": records }))
            }
            "com.atproto.sync.getBlob" => match state.blobs.get(&q["cid"]) {
                Some(b) => (200, b.clone()),
                None => not_found("BlobNotFound"),
            },
            "com.atproto.server.createSession" => {
                let handle = json_body()["identifier"].as_str().unwrap_or_default().to_string();
                let did = did_of(&handle);
                ok(json!({ "did": did, "handle": handle, "accessJwt": "access", "refreshJwt": "refresh", "didDoc": { "id": did, "service": [{ "id": "#atproto_pds", "serviceEndpoint": PDS }] } }))
            }
            "com.atproto.repo.uploadBlob" => {
                drop(state);
                ok(json!({ "blob": self.blob(&body.unwrap_or_default()) }))
            }
            "com.atproto.repo.deleteRecord" => {
                let input = json_body();
                let uri = format!("at://{}/{}/{}", input["repo"].as_str().unwrap_or_default(), input["collection"].as_str().unwrap_or_default(), input["rkey"].as_str().unwrap_or_default());
                state.records.remove(&uri);
                ok(json!({}))
            }
            "com.atproto.repo.createRecord" | "com.atproto.repo.putRecord" => {
                let input = json_body();
                let rkey = input["rkey"].as_str().unwrap_or("3l2tid").to_string();
                let uri = format!("at://{}/{}/{rkey}", input["repo"].as_str().unwrap_or_default(), input["collection"].as_str().unwrap_or_default());
                if nsid.ends_with("createRecord") && state.records.contains_key(&uri) {
                    return (400, json!({ "error": "InvalidRequest", "message": "exists" }).to_string().into_bytes());
                }
                let cid = record_cid(&input["record"]);
                state.records.insert(uri.clone(), (cid.clone(), input["record"].clone()));
                ok(json!({ "uri": uri, "cid": cid }))
            }
            _ => (501, json!({ "error": "MethodNotImplemented" }).to_string().into_bytes()),
        }
    }
}

impl Fetch for Pds {
    fn fetch(&self, method: &str, url: &str, _headers: Vec<(String, String)>, body: Option<Vec<u8>>) -> BoxFut<'_, Result<Response, String>> {
        let (status, body) = self.answer(method, url, body);
        Box::pin(async move { Ok(Response { status, headers: vec![], body }) })
    }
}