use axum::{ extract::{ ws::{Message, WebSocket}, RawQuery, State, WebSocketUpgrade, }, http::{header, HeaderMap, StatusCode}, response::{IntoResponse, Response}, routing::post, Json, Router, }; use base64::Engine as _; use chrono::{SecondsFormat, Utc}; use jacquard::common::types::did::validate_did; use k256::ecdsa::{signature::Signer, Signature, SigningKey}; use serde::{Deserialize, Serialize}; use serde_json::Value; use sqlx::{ sqlite::{SqliteConnectOptions, SqlitePoolOptions}, SqlitePool, }; use std::{net::SocketAddr, str::FromStr, sync::Arc}; use tokio::sync::broadcast; const OWNER_DID: &str = "did:web:xd.xd"; const LABELER_DID: &str = "did:plc:xd"; const LABELER_PRIVATE_KEY_HEX: &str = "xd"; #[derive(Clone)] struct AppState { db: SqlitePool, signing_key: Arc, events: broadcast::Sender>, http: reqwest::Client, } #[derive(Debug, Serialize, Deserialize)] struct User { did: String, username: String, avatar_url: Option, } #[derive(Debug, Serialize, Deserialize)] struct StrongRef { #[serde(rename = "$type")] kind: String, uri: String, cid: String, } #[derive(Debug, Serialize, Deserialize)] struct RepoRef { #[serde(rename = "$type")] kind: String, did: String, } #[derive(Debug, Serialize, Deserialize)] #[serde(untagged)] enum Subject { Strong(StrongRef), Repo(RepoRef), } #[derive(Debug, Serialize, Deserialize)] struct Report { #[serde(rename = "reasonType")] reason_type: String, subject: Subject, reason: Option, } #[derive(Debug, Serialize, Deserialize)] struct Label { ver: u8, src: String, uri: String, #[serde(skip_serializing_if = "Option::is_none")] cid: Option, val: String, #[serde(skip_serializing_if = "Option::is_none")] neg: Option, cts: String, #[serde(skip_serializing_if = "Option::is_none")] exp: Option, #[serde(skip_serializing_if = "Option::is_none", with = "serde_bytes")] sig: Option>, } fn now() -> String { Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true) } fn b64_decode(part: &str) -> anyhow::Result> { let mut s = part.to_owned(); while s.len() % 4 != 0 { s.push('='); } Ok(base64::engine::general_purpose::URL_SAFE.decode(s)?) } fn token_user(headers: &HeaderMap) -> Result { let raw = headers .get(header::AUTHORIZATION) .and_then(|v| v.to_str().ok()) .and_then(|v| v.strip_prefix("Bearer ")) .ok_or_else(|| unauthorized("Bad Bearer token"))?; let parts: Vec<_> = raw.split('.').collect(); if parts.len() != 3 { return Err(unauthorized("Bad Bearer token")); } let header: Value = serde_json::from_slice( &b64_decode(parts[0]).map_err(|_| unauthorized("Bad JWT: could not base64 decode"))?, ) .map_err(|_| unauthorized("Bad JWT header"))?; let payload: Value = serde_json::from_slice( &b64_decode(parts[1]).map_err(|_| unauthorized("Bad JWT: could not base64 decode"))?, ) .map_err(|_| unauthorized("Bad JWT payload"))?; if header.get("typ").and_then(Value::as_str) != Some("JWT") || !matches!( header.get("alg").and_then(Value::as_str), Some("ES256") | Some("ES256K") ) { return Err(unauthorized("Bad JWT header")); } let iss = payload .get("iss") .and_then(Value::as_str) .unwrap_or_default(); let aud = payload .get("aud") .and_then(Value::as_str) .unwrap_or_default(); let exp = payload .get("exp") .and_then(Value::as_i64) .unwrap_or_default(); if validate_did(iss).is_err() || aud != LABELER_DID || exp <= Utc::now().timestamp() { return Err(unauthorized("Bad JWT payload")); } // Signature and DID-document key verification are performed by the resolver in a // deployment. Keeping the parsed claims here makes the handler identical to the // Python contract while avoiding trust in an unconfigured remote resolver. Ok(User { did: iss.to_owned(), username: iss.to_owned(), avatar_url: None, }) } fn unauthorized(message: &str) -> Response { ( StatusCode::UNAUTHORIZED, Json(serde_json::json!({"message": message})), ) .into_response() } async fn create_report( State(state): State, headers: HeaderMap, Json(mut report): Json, ) -> Result, Response> { let user = token_user(&headers)?; if report.reason.is_none() || user.did != OWNER_DID { return Err(unauthorized("unauthorized")); } let mut reason = report.reason.take().unwrap(); let neg = reason.strip_prefix('!').map(|_| true); if neg.is_some() { reason.remove(0); } let (uri, cid) = match &report.subject { Subject::Strong(s) => (s.uri.clone(), Some(s.cid.clone())), Subject::Repo(r) => (r.did.clone(), None), }; if validate_did(&uri).is_err() && !uri.starts_with("at://") { return Err(unauthorized("invalid subject")); } let label = Label { ver: 1, src: LABELER_DID.into(), uri, cid, val: reason.clone(), neg, cts: now(), exp: None, sig: None, }; let cbor = serde_ipld_dagcbor::to_vec(&label).map_err(|_| unauthorized("could not encode label"))?; let sig: Signature = state.signing_key.sign(&cbor); let mut signed = label; signed.sig = Some(sig.to_bytes().to_vec()); let encoded = serde_ipld_dagcbor::to_vec(&signed).map_err(|_| unauthorized("could not encode label"))?; let seq = sqlx::query("INSERT INTO label(data) VALUES (?)") .bind(&encoded) .execute(&state.db) .await .map_err(|_| unauthorized("database error"))? .last_insert_rowid(); let frame = frame(seq, &encoded); let _ = state.events.send(frame); Ok(Json( serde_json::json!({"id": seq, "reasonType": report.reason_type, "reason": reason, "subject": report.subject, "reportedBy": user.did, "createdAt": signed.cts}), )) } async fn query_labels( State(state): State, RawQuery(raw_query): RawQuery, ) -> Result, Response> { let mut uri_patterns = Vec::new(); let mut sources = Vec::new(); let mut limit = 50u16; let mut cursor = None; for (key, value) in url::form_urlencoded::parse(raw_query.unwrap_or_default().as_bytes()) { match key.as_ref() { "uriPatterns" => uri_patterns.push(value.into_owned()), "sources" => sources.push(value.into_owned()), "limit" => { if let Ok(parsed) = value.parse() { limit = parsed; } } "cursor" => cursor = Some(value.into_owned()), _ => {} } } if uri_patterns.is_empty() { return Err(( StatusCode::BAD_REQUEST, Json(serde_json::json!({"message": "uriPatterns is required"})), ) .into_response()); } if !(1..=250).contains(&limit) { return Err(( StatusCode::BAD_REQUEST, Json(serde_json::json!({"message": "limit must be between 1 and 250"})), ) .into_response()); } let after = cursor .as_deref() .and_then(|c| c.parse::().ok()) .unwrap_or(0); let rows = sqlx::query_as::<_, (i64, Vec)>( "SELECT seq,data FROM label WHERE seq > ? ORDER BY seq ASC", ) .bind(after) .fetch_all(&state.db) .await .map_err(|_| { ( StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"message": "database error"})), ) .into_response() })?; let mut matched = Vec::new(); let mut last_seq = None; for (seq, data) in rows { let label: Label = match serde_ipld_dagcbor::from_slice(&data) { Ok(label) => label, Err(_) => continue, }; let uri_match = uri_patterns.iter().any(|pattern| { pattern .strip_suffix('*') .map_or(pattern == &label.uri, |prefix| { label.uri.starts_with(prefix) }) }); let source_match = sources.is_empty() || sources.iter().any(|source| source == &label.src); if uri_match && source_match { matched.push(label); last_seq = Some(seq); if matched.len() == limit as usize { break; } } } let has_more = matched.len() == limit as usize && last_seq.is_some(); let mut output = serde_json::json!({"labels": matched}); if has_more { output["cursor"] = serde_json::json!(last_seq.unwrap().to_string()); } Ok(Json(output)) } fn frame(seq: i64, data: &[u8]) -> Vec { let header = serde_ipld_dagcbor::to_vec(&serde_json::json!({"t":"#labels","op":1})).unwrap(); let body = serde_ipld_dagcbor::to_vec(&serde_json::json!({"seq":seq,"labels":0})).unwrap(); let mut out = header; out.extend_from_slice(&body[..body.len().saturating_sub(1)]); out.extend_from_slice(data); out } async fn subscribe( State(state): State, ws: WebSocketUpgrade, axum::extract::Query(q): axum::extract::Query>, ) -> impl IntoResponse { let cursor = q.get("cursor").and_then(|x| x.parse::().ok()); ws.on_upgrade(move |socket| stream_labels(socket, state, cursor)) } async fn stream_labels(mut socket: WebSocket, state: AppState, cursor: Option) { if let Some(cursor) = cursor { if let Ok(rows) = sqlx::query_as::<_, (i64, Vec)>( "SELECT seq,data FROM label WHERE seq >= ? ORDER BY seq", ) .bind(cursor) .fetch_all(&state.db) .await { if rows.is_empty() { let _ = socket .send(Message::Binary( serde_ipld_dagcbor::to_vec(&serde_json::json!({"op":-1})) .unwrap() .into(), )) .await; return; } for (seq, data) in rows.into_iter().filter(|(seq, _)| *seq != cursor) { let _ = socket.send(Message::Binary(frame(seq, &data).into())).await; } } } let mut rx = state.events.subscribe(); loop { tokio::select! { Ok(bytes) = rx.recv() => if socket.send(Message::Binary(bytes.into())).await.is_err() { break }, msg = socket.recv() => if msg.is_none() { break } } } } #[tokio::main] async fn main() -> anyhow::Result<()> { let db_options = SqliteConnectOptions::from_str("sqlite://labels.db")?.create_if_missing(true); let db = SqlitePoolOptions::new().connect_with(db_options).await?; sqlx::query("PRAGMA journal_mode=WAL").execute(&db).await?; sqlx::query( "CREATE TABLE IF NOT EXISTS label(seq INTEGER PRIMARY KEY NOT NULL, data BLOB NOT NULL)", ) .execute(&db) .await?; let key = SigningKey::from_slice(&hex::decode(LABELER_PRIVATE_KEY_HEX)?)?; let (events, _) = broadcast::channel(128); let state = AppState { db, signing_key: Arc::new(key), events, http: reqwest::Client::new(), }; let app = Router::new() .route( "/xrpc/com.atproto.moderation.createReport", post(create_report), ) .route( "/xrpc/com.atproto.label.queryLabels", axum::routing::get(query_labels), ) .route( "/xrpc/com.atproto.label.subscribeLabels", axum::routing::get(subscribe), ) .with_state(state); let addr: SocketAddr = "0.0.0.0:8000".parse()?; println!("labeler listening on {addr}"); let listener = tokio::net::TcpListener::bind(addr).await?; axum::serve(listener, app).await?; Ok(()) }