Something went wrong. Try again.
A lexicon-driven AppView for ATProto. forked from gamesgamesgamesgames.games/happyview
Something went wrong. Try again.
Rust
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360use mlua::{Lua, Result as LuaResult};use std::sync::Arc;
use crate::AppState;use crate::db::{adapt_sql, now_rfc3339};use crate::profile;
/// Register the `atproto` table with AT Protocol utility functions.pub fn register_atproto_api(lua: &Lua, state: Arc<AppState>) -> LuaResult<()> { let atproto_table = lua.create_table()?;
let state_clone = state.clone(); let resolve_fn = lua.create_async_function(move |_lua, did: String| { let state = state_clone.clone(); async move { let result = profile::resolve_pds_endpoint(&state.http, &state.config.plc_url, &did).await;
match result { Ok(endpoint) => Ok(Some(endpoint)), Err(_) => Ok(None), } } })?;
atproto_table.set("resolve_service_endpoint", resolve_fn)?;
// get_labels(uri) -> array of { src, uri, val, cts } let state_clone = state.clone(); let get_labels_fn = lua.create_async_function(move |lua, uri: String| { let state = state_clone.clone(); async move { let backend = state.db_backend; let now = now_rfc3339(); let sql = adapt_sql( "SELECT src, uri, val, cts FROM labels WHERE uri = ? AND (exp IS NULL OR exp > ?)", backend, ); let rows: Vec<(String, String, String, String)> = sqlx::query_as(&sql) .bind(&uri) .bind(&now) .fetch_all(&state.db) .await .map_err(|e| mlua::Error::runtime(format!("label query failed: {e}")))?;
let result = lua.create_table()?; let mut idx = 1;
for (src, label_uri, val, cts) in &rows { let label = lua.create_table()?; label.set("src", src.as_str())?; label.set("uri", label_uri.as_str())?; label.set("val", val.as_str())?; label.set("cts", cts.as_str())?; result.set(idx, label)?; idx += 1; }
// Check for self-labels in the record itself. let record_sql = adapt_sql("SELECT did, record FROM records WHERE uri = ?", backend); let record: Option<(String, String)> = sqlx::query_as(&record_sql) .bind(&uri) .fetch_optional(&state.db) .await .map_err(|e| mlua::Error::runtime(format!("record query failed: {e}")))?;
if let Some((did, record_str)) = record { let record_val: serde_json::Value = serde_json::from_str(&record_str).unwrap_or(serde_json::json!({})); if let Some(labels) = record_val.get("labels") && let Some(values) = labels.get("values") && let Some(arr) = values.as_array() { for item in arr { if let Some(val) = item.get("val").and_then(|v| v.as_str()) { let label = lua.create_table()?; label.set("src", did.as_str())?; label.set("uri", uri.as_str())?; label.set("val", val)?; label.set("cts", "")?; result.set(idx, label)?; idx += 1; } } } }
Ok(mlua::Value::Table(result)) } })?; atproto_table.set("get_labels", get_labels_fn)?;
// get_labels_batch(uris) -> table keyed by URI let state_clone = state.clone(); let get_labels_batch_fn = lua.create_async_function(move |lua, uris: mlua::Table| { let state = state_clone.clone(); async move { let backend = state.db_backend; // Collect URIs from the Lua table. let uri_list: Vec<String> = uris .sequence_values::<String>() .collect::<Result<Vec<_>, _>>()?;
let now = now_rfc3339();
// Query labels for all URIs (one query per URI since AnyPool doesn't support array binding). let label_sql = adapt_sql( "SELECT src, uri, val, cts FROM labels WHERE uri = ? AND (exp IS NULL OR exp > ?)", backend, ); let mut rows: Vec<(String, String, String, String)> = Vec::new(); for uri in &uri_list { let mut uri_rows: Vec<(String, String, String, String)> = sqlx::query_as(&label_sql) .bind(uri) .bind(&now) .fetch_all(&state.db) .await .map_err(|e| { mlua::Error::runtime(format!("label batch query failed: {e}")) })?; rows.append(&mut uri_rows); }
// Query records for self-labels. let record_sql = adapt_sql( "SELECT uri, did, record FROM records WHERE uri = ?", backend, ); let mut records: Vec<(String, String, String)> = Vec::new(); for uri in &uri_list { let mut uri_records: Vec<(String, String, String)> = sqlx::query_as(&record_sql) .bind(uri) .fetch_all(&state.db) .await .map_err(|e| mlua::Error::runtime(format!("record batch query failed: {e}")))?; records.append(&mut uri_records); }
// Build result table keyed by URI. let result = lua.create_table()?;
// Initialize empty arrays for each URI. let mut counters: std::collections::HashMap<String, i32> = std::collections::HashMap::new(); for uri in &uri_list { result.set(uri.as_str(), lua.create_table()?)?; counters.insert(uri.clone(), 1); }
// Add external labels. for (src, uri, val, cts) in &rows { let label = lua.create_table()?; label.set("src", src.as_str())?; label.set("uri", uri.as_str())?; label.set("val", val.as_str())?; label.set("cts", cts.as_str())?;
let uri_table: mlua::Table = result.get(uri.as_str())?; let idx = counters.get(uri).copied().unwrap_or(1); uri_table.set(idx, label)?; counters.insert(uri.clone(), idx + 1); }
// Add self-labels from records. for (uri, did, record_str) in &records { let record_val: serde_json::Value = serde_json::from_str(record_str).unwrap_or(serde_json::json!({})); if let Some(labels) = record_val.get("labels") && let Some(values) = labels.get("values") && let Some(arr) = values.as_array() { for item in arr { if let Some(val) = item.get("val").and_then(|v| v.as_str()) { let label = lua.create_table()?; label.set("src", did.as_str())?; label.set("uri", uri.as_str())?; label.set("val", val)?; label.set("cts", "")?;
let uri_table: mlua::Table = result.get(uri.as_str())?; let idx = counters.get(uri).copied().unwrap_or(1); uri_table.set(idx, label)?; counters.insert(uri.clone(), idx + 1); } } } }
Ok(mlua::Value::Table(result)) } })?; atproto_table.set("get_labels_batch", get_labels_batch_fn)?;
lua.globals().set("atproto", atproto_table)?; Ok(())}
#[cfg(test)]mod tests { use super::*; use crate::config::Config; use crate::db::DatabaseBackend; use crate::lexicon::LexiconRegistry; use tokio::sync::watch;
fn test_state_with_plc(plc_url: &str) -> AppState { let config = Config { host: "127.0.0.1".into(), port: 3000, database_url: String::new(), database_backend: crate::db::DatabaseBackend::Sqlite, public_url: String::new(), session_secret: "test-secret".into(), tap_url: String::new(), tap_admin_password: None, relay_url: String::new(), plc_url: plc_url.to_string(), static_dir: String::new(), event_log_retention_days: 30, app_name: None, logo_uri: None, tos_uri: None, policy_uri: None, token_encryption_key: None, }; let (tx, _) = watch::channel(vec![]); let (labeler_tx, _) = watch::channel(()); sqlx::any::install_default_drivers(); let test_db = sqlx::AnyPool::connect_lazy("sqlite::memory:").unwrap(); let atrium_http = std::sync::Arc::new(atrium_oauth::DefaultHttpClient::default()); let did_resolver = atrium_identity::did::CommonDidResolver::new( atrium_identity::did::CommonDidResolverConfig { plc_directory_url: "https://plc.directory".into(), http_client: std::sync::Arc::clone(&atrium_http), }, ); let handle_resolver = atrium_identity::handle::AtprotoHandleResolver::new( atrium_identity::handle::AtprotoHandleResolverConfig { dns_txt_resolver: crate::dns::NativeDnsResolver::new(), http_client: atrium_http, }, ); let oauth = atrium_oauth::OAuthClient::new(atrium_oauth::OAuthClientConfig { client_metadata: atrium_oauth::AtprotoLocalhostClientMetadata { redirect_uris: Some(vec!["http://127.0.0.1:0/auth/callback".into()]), scopes: Some(vec![atrium_oauth::Scope::Known( atrium_oauth::KnownScope::Atproto, )]), }, keys: None, state_store: crate::auth::oauth_store::DbStateStore::new( test_db.clone(), crate::db::DatabaseBackend::Sqlite, ), session_store: crate::auth::oauth_store::DbSessionStore::new( test_db.clone(), crate::db::DatabaseBackend::Sqlite, ), resolver: atrium_oauth::OAuthResolverConfig { did_resolver, handle_resolver, authorization_server_metadata: Default::default(), protected_resource_metadata: Default::default(), }, }) .expect("Failed to create test OAuth client"); AppState { config, http: reqwest::Client::new(), db: test_db, db_backend: DatabaseBackend::Sqlite, lexicons: LexiconRegistry::new(), collections_tx: tx, labeler_subscriptions_tx: labeler_tx, rate_limiter: crate::rate_limit::RateLimiter::new( false, crate::rate_limit::RateLimitConfig { capacity: 100, refill_rate: 2.0, default_query_cost: 1, default_procedure_cost: 1, default_proxy_cost: 1, }, vec![], ), oauth: std::sync::Arc::new(oauth), cookie_key: axum_extra::extract::cookie::Key::derive_from( b"test-secret-for-tests-only-not-production", ), plugin_registry: std::sync::Arc::new(crate::plugin::PluginRegistry::new()), wasm_runtime: std::sync::Arc::new( crate::plugin::WasmRuntime::new().expect("wasm runtime"), ), attestation_signer: None, } }
#[tokio::test] async fn resolve_service_endpoint_returns_endpoint() { let mock = wiremock::MockServer::start().await;
let did_doc = serde_json::json!({ "id": "did:plc:test123", "alsoKnownAs": ["at://test.example.com"], "service": [{ "id": "#atproto_pds", "type": "AtprotoPersonalDataServer", "serviceEndpoint": "https://pds.example.com" }] });
wiremock::Mock::given(wiremock::matchers::method("GET")) .and(wiremock::matchers::path("/did:plc:test123")) .respond_with(wiremock::ResponseTemplate::new(200).set_body_json(&did_doc)) .mount(&mock) .await;
let state = test_state_with_plc(&mock.uri()); let lua = mlua::Lua::new(); register_atproto_api(&lua, Arc::new(state)).unwrap();
let chunk = r#"return atproto.resolve_service_endpoint("did:plc:test123")"#; let result: String = lua.load(chunk).eval_async().await.unwrap(); assert_eq!(result, "https://pds.example.com"); }
#[tokio::test] async fn resolve_service_endpoint_returns_nil_on_failure() { let mock = wiremock::MockServer::start().await;
wiremock::Mock::given(wiremock::matchers::method("GET")) .and(wiremock::matchers::path("/did:plc:unknown")) .respond_with(wiremock::ResponseTemplate::new(404)) .mount(&mock) .await;
let state = test_state_with_plc(&mock.uri()); let lua = mlua::Lua::new(); register_atproto_api(&lua, Arc::new(state)).unwrap();
let chunk = r#"return atproto.resolve_service_endpoint("did:plc:unknown")"#; let result: mlua::Value = lua.load(chunk).eval_async().await.unwrap(); assert!(matches!(result, mlua::Value::Nil)); }
#[tokio::test] async fn resolve_did_web() { let mock = wiremock::MockServer::start().await;
let state = test_state_with_plc(&mock.uri()); let lua = mlua::Lua::new(); register_atproto_api(&lua, Arc::new(state)).unwrap();
let chunk = r#"return type(atproto.resolve_service_endpoint)"#; let result: String = lua.load(chunk).eval_async().await.unwrap(); assert_eq!(result, "function"); }}