diff --git a/.env.example b/.env.example index 9df8af7..942de84 100644 --- a/.env.example +++ b/.env.example @@ -13,6 +13,10 @@ AWS_SECRET_ACCESS_KEY=minioadmin PDS_HOSTNAME=localhost:3000 PLC_URL=plc.directory +# A comma-separated list of WebSocket URLs for firehose relays to push updates to. +# e.g., RELAYS=wss://relay.bsky.social,wss://another-relay.com +RELAYS= + # Notification Service Configuration # At least one notification channel should be configured for user notifications to work. # Email notifications (via sendmail/msmtp) diff --git a/Cargo.lock b/Cargo.lock index 502fe36..6cf9e05 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -593,6 +593,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5b098575ebe77cb6d14fc7f32749631a6e44edbef6b796f89b020e99ba20d425" dependencies = [ "axum-core", + "axum-macros", + "base64 0.22.1", "bytes", "form_urlencoded", "futures-util", @@ -611,8 +613,10 @@ dependencies = [ "serde_json", "serde_path_to_error", "serde_urlencoded", + "sha1", "sync_wrapper", "tokio", + "tokio-tungstenite", "tower", "tower-layer", "tower-service", @@ -638,6 +642,17 @@ dependencies = [ "tracing", ] +[[package]] +name = "axum-macros" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "604fde5e028fea851ce1d8570bbdc034bec850d157f7569d10f347d06808c05c" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.111", +] + [[package]] name = "base-x" version = "0.2.11" @@ -862,6 +877,8 @@ dependencies = [ "cid", "ctor", "dotenvy", + "futures", + "iroh-car", "jacquard", "jacquard-axum", "jacquard-repo", @@ -872,6 +889,7 @@ dependencies = [ "rand 0.8.5", "reqwest", "serde", + "serde_bytes", "serde_ipld_dagcbor", "serde_json", "sha2", @@ -880,6 +898,7 @@ dependencies = [ "testcontainers-modules", "thiserror 2.0.17", "tokio", + "tokio-tungstenite", "tracing", "tracing-subscriber", "uuid", @@ -5558,6 +5577,20 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-tungstenite" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d25a406cddcc431a75d3d9afc6a7c0f7428d4891dd973e4d54c56b46127bf857" +dependencies = [ + "futures-util", + "log", + "native-tls", + "tokio", + "tokio-native-tls", + "tungstenite", +] + [[package]] name = "tokio-util" version = "0.7.17" @@ -5747,6 +5780,24 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "tungstenite" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8628dcc84e5a09eb3d8423d6cb682965dea9133204e8fb3efee74c2a0c259442" +dependencies = [ + "bytes", + "data-encoding", + "http 1.4.0", + "httparse", + "log", + "native-tls", + "rand 0.9.2", + "sha1", + "thiserror 2.0.17", + "utf-8", +] + [[package]] name = "typenum" version = "1.19.0" diff --git a/Cargo.toml b/Cargo.toml index acb2320..a395b66 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,13 +8,14 @@ anyhow = "1.0.100" async-trait = "0.1.89" aws-config = "1.8.11" aws-sdk-s3 = "1.116.0" -axum = "0.8.7" +axum = { version = "0.8.7", features = ["ws", "macros"] } base64 = "0.22.1" bcrypt = "0.17.1" bytes = "1.11.0" chrono = { version = "0.4.42", features = ["serde"] } cid = "0.11.1" dotenvy = "0.15.7" +futures = "0.3.30" jacquard = { version = "0.9.3", default-features = false, features = ["api", "api_bluesky", "api_full", "derive", "dns"] } jacquard-axum = "0.9.2" jacquard-repo = "0.9.2" @@ -25,6 +26,7 @@ multihash = "0.19.3" rand = "0.8.5" reqwest = { version = "0.12.24", features = ["json"] } serde = { version = "1.0.228", features = ["derive"] } +serde_bytes = "0.11.14" serde_ipld_dagcbor = "0.6.4" serde_json = "1.0.145" sha2 = "0.10.9" @@ -33,10 +35,12 @@ thiserror = "2.0.17" tokio = { version = "1.48.0", features = ["macros", "rt-multi-thread", "time", "signal", "process"] } tracing = "0.1.43" tracing-subscriber = "0.3.22" +tokio-tungstenite = { version = "0.28.0", features = ["native-tls"] } uuid = { version = "1.19.0", features = ["v4", "fast-rng"] } [dev-dependencies] ctor = "0.6.3" +iroh-car = "0.5.1" testcontainers = "0.26.0" testcontainers-modules = { version = "0.14.0", features = ["postgres"] } wiremock = "0.6.5" diff --git a/TODO.md b/TODO.md index aef40c4..c59dbac 100644 --- a/TODO.md +++ b/TODO.md @@ -43,12 +43,11 @@ Lewis' corrected big boy todofile - [x] Implement `com.atproto.server.confirmEmail`. ## Repository Operations (`com.atproto.repo`) -- [ ] Record CRUD +- [x] Record CRUD - [x] Implement `com.atproto.repo.createRecord`. - - [ ] Validate schema against Lexicon (just structure, not complex logic). - [x] Generate `rkey` (TID) if not provided. - [x] Handle MST (Merkle Search Tree) insertion. - - [ ] **Trigger Firehose Event**. + - [x] **Trigger Firehose Event**. - [x] Implement `com.atproto.repo.putRecord`. - [x] Implement `com.atproto.repo.getRecord`. - [x] Implement `com.atproto.repo.deleteRecord`. @@ -57,17 +56,17 @@ Lewis' corrected big boy todofile - [x] Implement `com.atproto.repo.applyWrites` (Batch writes). - [ ] Implement `com.atproto.repo.importRepo` (Migration). - [x] Implement `com.atproto.repo.listMissingBlobs`. -- [ ] Blob Management +- [x] Blob Management - [x] Implement `com.atproto.repo.uploadBlob`. - [x] Store blob (S3). - [x] return `blob` ref (CID + MimeType). ## Sync & Federation (`com.atproto.sync`) -- [ ] The Firehose (WebSocket) - - [ ] Implement `com.atproto.sync.subscribeRepos`. - - [ ] Broadcast real-time commit events. - - [ ] Handle cursor replay (backfill). -- [ ] Bulk Export +- [x] The Firehose (WebSocket) + - [x] Implement `com.atproto.sync.subscribeRepos`. + - [x] Broadcast real-time commit events. + - [x] Handle cursor replay (backfill). +- [x] Bulk Export - [x] Implement `com.atproto.sync.getRepo` (Return full CAR file of repo). - [x] Implement `com.atproto.sync.getBlocks` (Return specific blocks via CIDs). - [x] Implement `com.atproto.sync.getLatestCommit`. @@ -75,7 +74,7 @@ Lewis' corrected big boy todofile - [x] Implement `com.atproto.sync.getRepoStatus`. - [x] Implement `com.atproto.sync.listRepos`. - [x] Implement `com.atproto.sync.notifyOfUpdate`. -- [ ] Blob Sync +- [x] Blob Sync - [x] Implement `com.atproto.sync.getBlob`. - [x] Implement `com.atproto.sync.listBlobs`. - [x] Crawler Interaction @@ -110,17 +109,17 @@ Lewis' corrected big boy todofile - [ ] Handle this generically. ## Infrastructure & Core Components -- [ ] Sequencer (Event Log) - - [ ] Implement a `Sequencer` (backed by `repo_seq` table? Like in ref impl). - - [ ] Implement event formatting (`commit`, `handle`, `identity`, `account`). - - [ ] Implement database polling / event emission mechanism. - - [ ] Implement cursor-based event replay (`requestSeqRange`). -- [ ] Repo Storage & Consistency (in postgres) - - [ ] Implement `RepoStorage` for postgres (replaces per-user SQLite). - - [ ] Read/Write IPLD blocks to `blocks` table (global deduplication). - - [ ] Manage Repo Root in `repos` table. - - [ ] Implement Atomic Repo Transactions. - - [ ] Ensure `blocks` write, `repo_root` update, `records` index update, and `sequencer` event are committed in a single transaction. +- [x] Sequencer (Event Log) + - [x] Implement a `Sequencer` (backed by `repo_seq` table). + - [x] Implement event formatting (`commit`, `handle`, `identity`, `account`). + - [x] Implement database polling / event emission mechanism. + - [x] Implement cursor-based event replay (`requestSeqRange`). +- [x] Repo Storage & Consistency (in postgres) + - [x] Implement `RepoStorage` for postgres (replaces per-user SQLite). + - [x] Read/Write IPLD blocks to `blocks` table (global deduplication). + - [x] Manage Repo Root in `repos` table. + - [x] Implement Atomic Repo Transactions. + - [x] Ensure `blocks` write, `repo_root` update, `records` index update, and `sequencer` event are committed in a single transaction. - [ ] Implement concurrency control (row-level locking on `repos` table) to prevent concurrent writes to the same repo. - [ ] DID Cache - [ ] Implement caching layer for DID resolution (Redis or in-memory). @@ -138,9 +137,9 @@ Lewis' corrected big boy todofile - [x] Helper functions for common notification types (welcome, password reset, email verification, etc.) - [ ] Image Processing - [ ] Implement image resize/formatting pipeline (for blob uploads). -- [ ] IPLD & MST - - [ ] Implement Merkle Search Tree logic for repo signing. - - [ ] Implement CAR (Content Addressable Archive) encoding/decoding. +- [x] IPLD & MST + - [x] Implement Merkle Search Tree logic for repo signing. + - [x] Implement CAR (Content Addressable Archive) encoding/decoding. - [ ] Validation - [ ] DID PLC Operations (Sign rotation keys). - [ ] Fix any remaining TODOs in the code, everywhere, full stop. diff --git a/migrations/202512211402_repo_sequencer.sql b/migrations/202512211402_repo_sequencer.sql new file mode 100644 index 0000000..1cf3185 --- /dev/null +++ b/migrations/202512211402_repo_sequencer.sql @@ -0,0 +1,13 @@ +CREATE TABLE repo_seq ( + seq BIGSERIAL PRIMARY KEY, + did TEXT NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + event_type TEXT NOT NULL, + commit_cid TEXT, + prev_cid TEXT, + ops JSONB, + blobs TEXT[] +); + +CREATE INDEX idx_repo_seq_seq ON repo_seq(seq); +CREATE INDEX idx_repo_seq_did ON repo_seq(did); diff --git a/migrations/202512211403_add_blocks_cids_to_repo_seq.sql b/migrations/202512211403_add_blocks_cids_to_repo_seq.sql new file mode 100644 index 0000000..52b24d9 --- /dev/null +++ b/migrations/202512211403_add_blocks_cids_to_repo_seq.sql @@ -0,0 +1,2 @@ +ALTER TABLE repo_seq ADD COLUMN blocks_cids TEXT[]; + diff --git a/src/api/repo/record/batch.rs b/src/api/repo/record/batch.rs index 9b3e7dc..8fba0a5 100644 --- a/src/api/repo/record/batch.rs +++ b/src/api/repo/record/batch.rs @@ -1,17 +1,15 @@ +use crate::api::repo::record::utils::{commit_and_log, RecordOp}; +use crate::repo::tracking::TrackingBlockStore; use crate::state::AppState; use axum::{ - Json, extract::State, http::StatusCode, response::{IntoResponse, Response}, + Json, }; use chrono::Utc; use cid::Cid; -use jacquard::types::{ - did::Did, - integer::LimitedU32, - string::{Nsid, Tid}, -}; +use jacquard::types::string::Nsid; use jacquard_repo::{commit::Commit, mst::Mst, storage::BlockStore}; use serde::{Deserialize, Serialize}; use serde_json::json; @@ -98,10 +96,7 @@ pub async fn apply_writes( .unwrap_or(None); let (did, key_bytes) = match session { - Some(row) => ( - row.did, - row.key_bytes, - ), + Some(row) => (row.did, row.key_bytes), None => { return ( StatusCode::UNAUTHORIZED, @@ -143,12 +138,11 @@ pub async fn apply_writes( .into_response(); } - let user_query = sqlx::query!("SELECT id FROM users WHERE did = $1", did) + let user_id: uuid::Uuid = match sqlx::query_scalar!("SELECT id FROM users WHERE did = $1", did) .fetch_optional(&state.db) - .await; - - let user_id: uuid::Uuid = match user_query { - Ok(Some(row)) => row.id, + .await + { + Ok(Some(id)) => id, _ => { return ( StatusCode::INTERNAL_SERVER_ERROR, @@ -158,45 +152,34 @@ pub async fn apply_writes( } }; - let repo_root_query = sqlx::query!("SELECT repo_root_cid FROM repos WHERE user_id = $1", user_id) - .fetch_optional(&state.db) - .await; - - let current_root_cid = match repo_root_query { - Ok(Some(row)) => { - let cid_str: String = row.repo_root_cid; - match Cid::from_str(&cid_str) { - Ok(c) => c, - Err(_) => { - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError", "message": "Invalid repo root CID"})), - ) - .into_response(); - } + let root_cid_str: String = + match sqlx::query_scalar!("SELECT repo_root_cid FROM repos WHERE user_id = $1", user_id) + .fetch_optional(&state.db) + .await + { + Ok(Some(cid_str)) => cid_str, + _ => { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({"error": "InternalError", "message": "Repo root not found"})), + ) + .into_response(); } - } - _ => { + }; + + let current_root_cid = match Cid::from_str(&root_cid_str) { + Ok(c) => c, + Err(_) => { return ( StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError", "message": "Repo root not found"})), + Json(json!({"error": "InternalError", "message": "Invalid repo root CID"})), ) .into_response(); } }; if let Some(swap_commit) = &input.swap_commit { - let swap_cid = match Cid::from_str(swap_commit) { - Ok(c) => c, - Err(_) => { - return ( - StatusCode::BAD_REQUEST, - Json(json!({"error": "InvalidSwap", "message": "Invalid swapCommit CID"})), - ) - .into_response(); - } - }; - if swap_cid != current_root_cid { + if Cid::from_str(swap_commit).ok() != Some(current_root_cid) { return ( StatusCode::CONFLICT, Json(json!({"error": "InvalidSwap", "message": "Repo has been modified"})), @@ -205,43 +188,34 @@ pub async fn apply_writes( } } - let commit_bytes = match state.block_store.get(¤t_root_cid).await { + let tracking_store = TrackingBlockStore::new(state.block_store.clone()); + + let commit_bytes = match tracking_store.get(¤t_root_cid).await { Ok(Some(b)) => b, - Ok(None) => { + _ => { return ( StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"error": "InternalError", "message": "Commit block not found"})), ) - .into_response(); - } - Err(e) => { - error!("Failed to load commit block: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); + .into_response() } }; let commit = match Commit::from_cbor(&commit_bytes) { Ok(c) => c, - Err(e) => { - error!("Failed to parse commit: {:?}", e); + _ => { return ( StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), + Json(json!({"error": "InternalError", "message": "Failed to parse commit"})), ) - .into_response(); + .into_response() } }; - let mst_root = commit.data; - let store = Arc::new(state.block_store.clone()); - let mut mst = Mst::load(store.clone(), mst_root, None); + let mut mst = Mst::load(Arc::new(tracking_store.clone()), commit.data, None); let mut results: Vec = Vec::new(); - let mut record_ops: Vec<(String, String, Option)> = Vec::new(); + let mut ops: Vec = Vec::new(); for write in &input.writes { match write { @@ -250,248 +224,98 @@ pub async fn apply_writes( rkey, value, } => { - let collection_nsid = match collection.parse::() { - Ok(n) => n, - Err(_) => { - return ( - StatusCode::BAD_REQUEST, - Json(json!({"error": "InvalidCollection"})), - ) - .into_response(); - } - }; - let rkey = rkey .clone() .unwrap_or_else(|| Utc::now().format("%Y%m%d%H%M%S%f").to_string()); - let mut record_bytes = Vec::new(); - if let Err(e) = serde_ipld_dagcbor::to_writer(&mut record_bytes, value) { - error!("Error serializing record: {:?}", e); - return ( - StatusCode::BAD_REQUEST, - Json(json!({"error": "InvalidRecord", "message": "Failed to serialize record"})), - ) - .into_response(); - } + serde_ipld_dagcbor::to_writer(&mut record_bytes, value).unwrap(); + let record_cid = tracking_store.put(&record_bytes).await.unwrap(); - let record_cid = match state.block_store.put(&record_bytes).await { - Ok(c) => c, - Err(e) => { - error!("Failed to save record block: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } - }; - - let key = format!("{}/{}", collection_nsid, rkey); - mst = match mst.add(&key, record_cid).await { - Ok(m) => m, - Err(e) => { - error!("Failed to add to MST: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } - }; + let key = format!("{}/{}", collection.parse::().unwrap(), rkey); + mst = mst.add(&key, record_cid).await.unwrap(); let uri = format!("at://{}/{}/{}", did, collection, rkey); results.push(WriteResult::CreateResult { - uri: uri.clone(), + uri, cid: record_cid.to_string(), }); - record_ops.push((collection.clone(), rkey, Some(record_cid.to_string()))); + ops.push(RecordOp::Create { + collection: collection.clone(), + rkey, + cid: record_cid, + }); } WriteOp::Update { collection, rkey, value, } => { - let collection_nsid = match collection.parse::() { - Ok(n) => n, - Err(_) => { - return ( - StatusCode::BAD_REQUEST, - Json(json!({"error": "InvalidCollection"})), - ) - .into_response(); - } - }; - let mut record_bytes = Vec::new(); - if let Err(e) = serde_ipld_dagcbor::to_writer(&mut record_bytes, value) { - error!("Error serializing record: {:?}", e); - return ( - StatusCode::BAD_REQUEST, - Json(json!({"error": "InvalidRecord", "message": "Failed to serialize record"})), - ) - .into_response(); - } + serde_ipld_dagcbor::to_writer(&mut record_bytes, value).unwrap(); + let record_cid = tracking_store.put(&record_bytes).await.unwrap(); - let record_cid = match state.block_store.put(&record_bytes).await { - Ok(c) => c, - Err(e) => { - error!("Failed to save record block: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } - }; - - let key = format!("{}/{}", collection_nsid, rkey); - mst = match mst.update(&key, record_cid).await { - Ok(m) => m, - Err(e) => { - error!("Failed to update MST: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } - }; + let key = format!("{}/{}", collection.parse::().unwrap(), rkey); + mst = mst.update(&key, record_cid).await.unwrap(); let uri = format!("at://{}/{}/{}", did, collection, rkey); results.push(WriteResult::UpdateResult { - uri: uri.clone(), + uri, cid: record_cid.to_string(), }); - record_ops.push((collection.clone(), rkey.clone(), Some(record_cid.to_string()))); + ops.push(RecordOp::Update { + collection: collection.clone(), + rkey: rkey.clone(), + cid: record_cid, + }); } WriteOp::Delete { collection, rkey } => { - let collection_nsid = match collection.parse::() { - Ok(n) => n, - Err(_) => { - return ( - StatusCode::BAD_REQUEST, - Json(json!({"error": "InvalidCollection"})), - ) - .into_response(); - } - }; - - let key = format!("{}/{}", collection_nsid, rkey); - mst = match mst.delete(&key).await { - Ok(m) => m, - Err(e) => { - error!("Failed to delete from MST: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } - }; + let key = format!("{}/{}", collection.parse::().unwrap(), rkey); + mst = mst.delete(&key).await.unwrap(); results.push(WriteResult::DeleteResult {}); - record_ops.push((collection.clone(), rkey.clone(), None)); + ops.push(RecordOp::Delete { + collection: collection.clone(), + rkey: rkey.clone(), + }); } } } - let new_mst_root = match mst.persist().await { - Ok(c) => c, + let new_mst_root = mst.persist().await.unwrap(); + let written_cids = tracking_store.get_written_cids(); + let written_cids_str = written_cids + .iter() + .map(|c| c.to_string()) + .collect::>(); + + let commit_res = match commit_and_log( + &state, + &did, + user_id, + Some(current_root_cid), + new_mst_root, + ops, + &written_cids_str, + ) + .await + { + Ok(res) => res, Err(e) => { - error!("Failed to persist MST: {:?}", e); + error!("Commit failed: {}", e); return ( StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), + Json(json!({"error": "InternalError", "message": "Failed to commit changes"})), ) .into_response(); } }; - let did_obj = match Did::new(&did) { - Ok(d) => d, - Err(_) => { - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError", "message": "Invalid DID"})), - ) - .into_response(); - } - }; - - let rev = Tid::now(LimitedU32::MIN); - let new_commit = Commit::new_unsigned(did_obj, new_mst_root, rev.clone(), Some(current_root_cid)); - - let new_commit_bytes = match new_commit.to_cbor() { - Ok(b) => b, - Err(e) => { - error!("Failed to serialize new commit: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } - }; - - let new_root_cid = match state.block_store.put(&new_commit_bytes).await { - Ok(c) => c, - Err(e) => { - error!("Failed to save new commit: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } - }; - - let update_repo = sqlx::query!("UPDATE repos SET repo_root_cid = $1 WHERE user_id = $2", new_root_cid.to_string(), user_id) - .execute(&state.db) - .await; - - if let Err(e) = update_repo { - error!("Failed to update repo root in DB: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } - - for (collection, rkey, record_cid) in record_ops { - match record_cid { - Some(cid) => { - let _ = sqlx::query!( - "INSERT INTO records (repo_id, collection, rkey, record_cid) VALUES ($1, $2, $3, $4) - ON CONFLICT (repo_id, collection, rkey) DO UPDATE SET record_cid = $4, created_at = NOW()", - user_id, - collection, - rkey, - cid - ) - .execute(&state.db) - .await; - } - None => { - let _ = sqlx::query!( - "DELETE FROM records WHERE repo_id = $1 AND collection = $2 AND rkey = $3", - user_id, - collection, - rkey - ) - .execute(&state.db) - .await; - } - } - } - ( StatusCode::OK, Json(ApplyWritesOutput { commit: CommitInfo { - cid: new_root_cid.to_string(), - rev: rev.to_string(), + cid: commit_res.commit_cid.to_string(), + rev: commit_res.rev, }, results, }), diff --git a/src/api/repo/record/delete.rs b/src/api/repo/record/delete.rs index 299ab63..59343ab 100644 --- a/src/api/repo/record/delete.rs +++ b/src/api/repo/record/delete.rs @@ -1,16 +1,15 @@ +use crate::api::repo::record::utils::{commit_and_log, RecordOp}; +use crate::api::repo::record::write::prepare_repo_write; +use crate::repo::tracking::TrackingBlockStore; use crate::state::AppState; use axum::{ - Json, extract::State, - http::StatusCode, + http::{HeaderMap, StatusCode}, response::{IntoResponse, Response}, + Json, }; use cid::Cid; -use jacquard::types::{ - did::Did, - integer::LimitedU32, - string::{Nsid, Tid}, -}; +use jacquard::types::string::Nsid; use jacquard_repo::{commit::Commit, mst::Mst, storage::BlockStore}; use serde::Deserialize; use serde_json::json; @@ -31,122 +30,58 @@ pub struct DeleteRecordInput { pub async fn delete_record( State(state): State, - headers: axum::http::HeaderMap, + headers: HeaderMap, Json(input): Json, ) -> Response { - let auth_header = headers.get("Authorization"); - if auth_header.is_none() { - return ( - StatusCode::UNAUTHORIZED, - Json(json!({"error": "AuthenticationRequired"})), - ) - .into_response(); - } - let token = auth_header - .unwrap() - .to_str() - .unwrap_or("") - .replace("Bearer ", ""); + let (did, user_id, current_root_cid) = + match prepare_repo_write(&state, &headers, &input.repo).await { + Ok(res) => res, + Err(err_res) => return err_res, + }; - let session = sqlx::query!( - "SELECT s.did, k.key_bytes FROM sessions s JOIN users u ON s.did = u.did JOIN user_keys k ON u.id = k.user_id WHERE s.access_jwt = $1", - token - ) - .fetch_optional(&state.db) - .await - .unwrap_or(None); - - let (did, key_bytes) = match session { - Some(row) => ( - row.did, - row.key_bytes, - ), - None => { + if let Some(swap_commit) = &input.swap_commit { + if Cid::from_str(swap_commit).ok() != Some(current_root_cid) { return ( - StatusCode::UNAUTHORIZED, - Json(json!({"error": "AuthenticationFailed"})), + StatusCode::CONFLICT, + Json(json!({"error": "InvalidSwap", "message": "Repo has been modified"})), ) .into_response(); } - }; - - if let Err(_) = crate::auth::verify_token(&token, &key_bytes) { - return ( - StatusCode::UNAUTHORIZED, - Json(json!({"error": "AuthenticationFailed", "message": "Invalid token signature"})), - ) - .into_response(); } - if input.repo != did { - return (StatusCode::FORBIDDEN, Json(json!({"error": "InvalidRepo", "message": "Repo does not match authenticated user"}))).into_response(); - } + let tracking_store = TrackingBlockStore::new(state.block_store.clone()); - let user_query = sqlx::query!("SELECT id FROM users WHERE did = $1", did) - .fetch_optional(&state.db) - .await; - - let user_id: uuid::Uuid = match user_query { - Ok(Some(row)) => row.id, - _ => { - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError", "message": "User not found"})), - ) - .into_response(); - } - }; - - let repo_root_query = sqlx::query!("SELECT repo_root_cid FROM repos WHERE user_id = $1", user_id) - .fetch_optional(&state.db) - .await; - - let current_root_cid = match repo_root_query { - Ok(Some(row)) => { - let cid_str: String = row.repo_root_cid; - Cid::from_str(&cid_str).ok() - } - _ => None, - }; - - if current_root_cid.is_none() { - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError", "message": "Repo root not found"})), - ) - .into_response(); - } - let current_root_cid = current_root_cid.unwrap(); - - let commit_bytes = match state.block_store.get(¤t_root_cid).await { + let commit_bytes = match tracking_store.get(¤t_root_cid).await { Ok(Some(b)) => b, - Ok(None) => return (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"error": "InternalError", "message": "Commit block not found"}))).into_response(), - Err(e) => return (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"error": "InternalError", "message": format!("Failed to load commit block: {:?}", e)}))).into_response(), + _ => return (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"error": "InternalError", "message": "Commit block not found"}))).into_response(), }; - let commit = match Commit::from_cbor(&commit_bytes) { Ok(c) => c, - Err(e) => return (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"error": "InternalError", "message": format!("Failed to parse commit: {:?}", e)}))).into_response(), + _ => return (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"error": "InternalError", "message": "Failed to parse commit"}))).into_response(), }; - let mst_root = commit.data; - let store = Arc::new(state.block_store.clone()); - let mst = Mst::load(store.clone(), mst_root, None); - + let mst = Mst::load( + Arc::new(tracking_store.clone()), + commit.data, + None, + ); let collection_nsid = match input.collection.parse::() { Ok(n) => n, - Err(_) => { - return ( - StatusCode::BAD_REQUEST, - Json(json!({"error": "InvalidCollection"})), - ) - .into_response(); - } + Err(_) => return (StatusCode::BAD_REQUEST, Json(json!({"error": "InvalidCollection"}))).into_response(), }; - let key = format!("{}/{}", collection_nsid, input.rkey); - // TODO: Check swapRecord if provided? Skipping for brevity/robustness + if let Some(swap_record_str) = &input.swap_record { + let expected_cid = Cid::from_str(swap_record_str).ok(); + let actual_cid = mst.get(&key).await.ok().flatten(); + if expected_cid != actual_cid { + return (StatusCode::CONFLICT, Json(json!({"error": "InvalidSwap", "message": "Record has been modified or does not exist"}))).into_response(); + } + } + + if mst.get(&key).await.ok().flatten().is_none() { + return (StatusCode::OK, Json(json!({}))).into_response(); + } let new_mst = match mst.delete(&key).await { Ok(m) => m, @@ -160,73 +95,17 @@ pub async fn delete_record( Ok(c) => c, Err(e) => { error!("Failed to persist MST: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError", "message": "Failed to persist MST"})), - ) - .into_response(); + return (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"error": "InternalError", "message": "Failed to persist MST"}))).into_response(); } }; - let did_obj = match Did::new(&did) { - Ok(d) => d, - Err(_) => { - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError", "message": "Invalid DID"})), - ) - .into_response(); - } + let op = RecordOp::Delete { collection: input.collection, rkey: input.rkey }; + let written_cids = tracking_store.get_written_cids(); + let written_cids_str = written_cids.iter().map(|c| c.to_string()).collect::>(); + + if let Err(e) = commit_and_log(&state, &did, user_id, Some(current_root_cid), new_mst_root, vec![op], &written_cids_str).await { + return (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"error": "InternalError", "message": e}))).into_response(); }; - let rev = Tid::now(LimitedU32::MIN); - - let new_commit = Commit::new_unsigned(did_obj, new_mst_root, rev, Some(current_root_cid)); - - let new_commit_bytes = - match new_commit.to_cbor() { - Ok(b) => b, - Err(_e) => return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json( - json!({"error": "InternalError", "message": "Failed to serialize new commit"}), - ), - ) - .into_response(), - }; - - let new_root_cid = match state.block_store.put(&new_commit_bytes).await { - Ok(c) => c, - Err(_e) => { - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError", "message": "Failed to save new commit"})), - ) - .into_response(); - } - }; - - let update_repo = sqlx::query!("UPDATE repos SET repo_root_cid = $1 WHERE user_id = $2", new_root_cid.to_string(), user_id) - .execute(&state.db) - .await; - - if let Err(e) = update_repo { - error!("Failed to update repo root in DB: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError", "message": "Failed to update repo root in DB"})), - ) - .into_response(); - } - - let record_delete = - sqlx::query!("DELETE FROM records WHERE repo_id = $1 AND collection = $2 AND rkey = $3", user_id, input.collection, input.rkey) - .execute(&state.db) - .await; - - if let Err(e) = record_delete { - error!("Error deleting record index: {:?}", e); - } - (StatusCode::OK, Json(json!({}))).into_response() } diff --git a/src/api/repo/record/mod.rs b/src/api/repo/record/mod.rs index c983696..1eee4c2 100644 --- a/src/api/repo/record/mod.rs +++ b/src/api/repo/record/mod.rs @@ -1,12 +1,14 @@ pub mod batch; pub mod delete; pub mod read; +pub mod utils; pub mod write; pub use batch::apply_writes; pub use delete::{DeleteRecordInput, delete_record}; pub use read::{GetRecordInput, ListRecordsInput, ListRecordsOutput, get_record, list_records}; +pub use utils::*; pub use write::{ CreateRecordInput, CreateRecordOutput, PutRecordInput, PutRecordOutput, create_record, - put_record, + put_record, prepare_repo_write, }; diff --git a/src/api/repo/record/utils.rs b/src/api/repo/record/utils.rs new file mode 100644 index 0000000..565cbf9 --- /dev/null +++ b/src/api/repo/record/utils.rs @@ -0,0 +1,125 @@ +use crate::state::AppState; +use cid::Cid; +use jacquard::types::{did::Did, integer::LimitedU32, string::Tid}; +use jacquard_repo::commit::Commit; +use jacquard_repo::storage::BlockStore; +use serde_json::json; +use uuid::Uuid; + +pub enum RecordOp { + Create { collection: String, rkey: String, cid: Cid }, + Update { collection: String, rkey: String, cid: Cid }, + Delete { collection: String, rkey: String }, +} + +pub struct CommitResult { + pub commit_cid: Cid, + pub rev: String, +} + +pub async fn commit_and_log( + state: &AppState, + did: &str, + user_id: Uuid, + current_root_cid: Option, + new_mst_root: Cid, + ops: Vec, + blocks_cids: &Vec, +) -> Result { + let did_obj = Did::new(did).map_err(|e| format!("Invalid DID: {}", e))?; + let rev = Tid::now(LimitedU32::MIN); + + let new_commit = Commit::new_unsigned(did_obj, new_mst_root, rev.clone(), current_root_cid); + + let new_commit_bytes = new_commit.to_cbor().map_err(|e| format!("Failed to serialize commit: {:?}", e))?; + + let new_root_cid = state.block_store.put(&new_commit_bytes).await + .map_err(|e| format!("Failed to save commit block: {:?}", e))?; + + sqlx::query!("UPDATE repos SET repo_root_cid = $1 WHERE user_id = $2", new_root_cid.to_string(), user_id) + .execute(&state.db) + .await + .map_err(|e| format!("DB Error (repos): {}", e))?; + + for op in &ops { + match op { + RecordOp::Create { collection, rkey, cid } | RecordOp::Update { collection, rkey, cid } => { + sqlx::query!( + "INSERT INTO records (repo_id, collection, rkey, record_cid) VALUES ($1, $2, $3, $4) + ON CONFLICT (repo_id, collection, rkey) DO UPDATE SET record_cid = $4, created_at = NOW()", + user_id, + collection, + rkey, + cid.to_string() + ) + .execute(&state.db) + .await + .map_err(|e| format!("DB Error (records): {}", e))?; + } + RecordOp::Delete { collection, rkey } => { + sqlx::query!( + "DELETE FROM records WHERE repo_id = $1 AND collection = $2 AND rkey = $3", + user_id, + collection, + rkey + ) + .execute(&state.db) + .await + .map_err(|e| format!("DB Error (records): {}", e))?; + } + } + } + + let ops_json = ops.iter().map(|op| { + match op { + RecordOp::Create { collection, rkey, cid } => json!({ + "action": "create", + "path": format!("{}/{}", collection, rkey), + "cid": cid.to_string() + }), + RecordOp::Update { collection, rkey, cid } => json!({ + "action": "update", + "path": format!("{}/{}", collection, rkey), + "cid": cid.to_string() + }), + RecordOp::Delete { collection, rkey } => json!({ + "action": "delete", + "path": format!("{}/{}", collection, rkey), + "cid": null + }), + } + }).collect::>(); + + let event_type = "commit"; + let prev_cid_str = current_root_cid.map(|c| c.to_string()); + + let seq_row = sqlx::query!( + r#" + INSERT INTO repo_seq (did, event_type, commit_cid, prev_cid, ops, blobs, blocks_cids) + VALUES ($1, $2, $3, $4, $5, $6, $7) + RETURNING seq + "#, + did, + event_type, + new_root_cid.to_string(), + prev_cid_str, + json!(ops_json), + &[] as &[String], + blocks_cids, + ) + .fetch_one(&state.db) + .await + .map_err(|e| format!("DB Error (repo_seq): {}", e))?; + + sqlx::query( + &format!("NOTIFY repo_updates, '{}'", seq_row.seq) + ) + .execute(&state.db) + .await + .map_err(|e| format!("DB Error (notify): {}", e))?; + + Ok(CommitResult { + commit_cid: new_root_cid, + rev: rev.to_string(), + }) +} diff --git a/src/api/repo/record/write.rs b/src/api/repo/record/write.rs index 936e83c..eef5a93 100644 --- a/src/api/repo/record/write.rs +++ b/src/api/repo/record/write.rs @@ -1,23 +1,115 @@ +use crate::api::repo::record::utils::{commit_and_log, RecordOp}; +use crate::repo::tracking::TrackingBlockStore; use crate::state::AppState; use axum::{ - Json, extract::State, - http::StatusCode, + http::{HeaderMap, StatusCode}, response::{IntoResponse, Response}, + Json, }; use chrono::Utc; use cid::Cid; -use jacquard::types::{ - did::Did, - integer::LimitedU32, - string::{Nsid, Tid}, -}; +use jacquard::types::string::Nsid; use jacquard_repo::{commit::Commit, mst::Mst, storage::BlockStore}; use serde::{Deserialize, Serialize}; use serde_json::json; use std::str::FromStr; use std::sync::Arc; use tracing::error; +use uuid::Uuid; + +pub async fn prepare_repo_write( + state: &AppState, + headers: &HeaderMap, + repo_did: &str, +) -> Result<(String, Uuid, Cid), Response> { + let auth_header = headers.get("Authorization").ok_or_else(|| { + ( + StatusCode::UNAUTHORIZED, + Json(json!({"error": "AuthenticationRequired"})), + ) + .into_response() + })?; + let token = auth_header + .to_str() + .unwrap_or("") + .replace("Bearer ", ""); + + let session = sqlx::query!( + "SELECT s.did, k.key_bytes FROM sessions s JOIN users u ON s.did = u.did JOIN user_keys k ON u.id = k.user_id WHERE s.access_jwt = $1", + token + ) + .fetch_optional(&state.db) + .await + .map_err(|e| { + error!("DB error fetching session: {}", e); + (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"error": "InternalError"}))).into_response() + })? + .ok_or_else(|| { + ( + StatusCode::UNAUTHORIZED, + Json(json!({"error": "AuthenticationFailed"})), + ) + .into_response() + })?; + + crate::auth::verify_token(&token, &session.key_bytes).map_err(|_| { + ( + StatusCode::UNAUTHORIZED, + Json(json!({"error": "AuthenticationFailed", "message": "Invalid token signature"})), + ) + .into_response() + })?; + + if repo_did != session.did { + return Err(( + StatusCode::FORBIDDEN, + Json(json!({"error": "InvalidRepo", "message": "Repo does not match authenticated user"})), + ) + .into_response()); + } + + let user_id = sqlx::query_scalar!("SELECT id FROM users WHERE did = $1", session.did) + .fetch_optional(&state.db) + .await + .map_err(|e| { + error!("DB error fetching user: {}", e); + (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"error": "InternalError"}))).into_response() + })? + .ok_or_else(|| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({"error": "InternalError", "message": "User not found"})), + ) + .into_response() + })?; + + let root_cid_str: String = + sqlx::query_scalar!("SELECT repo_root_cid FROM repos WHERE user_id = $1", user_id) + .fetch_optional(&state.db) + .await + .map_err(|e| { + error!("DB error fetching repo root: {}", e); + (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"error": "InternalError"}))).into_response() + })? + .ok_or_else(|| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({"error": "InternalError", "message": "Repo root not found"})), + ) + .into_response() + })?; + + let current_root_cid = Cid::from_str(&root_cid_str).map_err(|_| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({"error": "InternalError", "message": "Invalid repo root CID"})), + ) + .into_response() + })?; + + Ok((session.did, user_id, current_root_cid)) +} #[derive(Deserialize)] #[allow(dead_code)] @@ -40,145 +132,47 @@ pub struct CreateRecordOutput { pub async fn create_record( State(state): State, - headers: axum::http::HeaderMap, + headers: HeaderMap, Json(input): Json, ) -> Response { - let auth_header = headers.get("Authorization"); - if auth_header.is_none() { - return ( - StatusCode::UNAUTHORIZED, - Json(json!({"error": "AuthenticationRequired"})), - ) - .into_response(); - } - let token = auth_header - .unwrap() - .to_str() - .unwrap_or("") - .replace("Bearer ", ""); + let (did, user_id, current_root_cid) = + match prepare_repo_write(&state, &headers, &input.repo).await { + Ok(res) => res, + Err(err_res) => return err_res, + }; - let session = sqlx::query!( - "SELECT s.did, k.key_bytes FROM sessions s JOIN users u ON s.did = u.did JOIN user_keys k ON u.id = k.user_id WHERE s.access_jwt = $1", - token - ) - .fetch_optional(&state.db) - .await - .unwrap_or(None); - - let (did, key_bytes) = match session { - Some(row) => ( - row.did, - row.key_bytes, - ), - None => { + if let Some(swap_commit) = &input.swap_commit { + if Cid::from_str(swap_commit).ok() != Some(current_root_cid) { return ( - StatusCode::UNAUTHORIZED, - Json(json!({"error": "AuthenticationFailed"})), + StatusCode::CONFLICT, + Json(json!({"error": "InvalidSwap", "message": "Repo has been modified"})), ) .into_response(); } - }; - - if let Err(_) = crate::auth::verify_token(&token, &key_bytes) { - return ( - StatusCode::UNAUTHORIZED, - Json(json!({"error": "AuthenticationFailed", "message": "Invalid token signature"})), - ) - .into_response(); } - if input.repo != did { - return (StatusCode::FORBIDDEN, Json(json!({"error": "InvalidRepo", "message": "Repo does not match authenticated user"}))).into_response(); - } + let tracking_store = TrackingBlockStore::new(state.block_store.clone()); - let user_query = sqlx::query!("SELECT id FROM users WHERE did = $1", did) - .fetch_optional(&state.db) - .await; - - let user_id: uuid::Uuid = match user_query { - Ok(Some(row)) => row.id, - _ => { - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError", "message": "User not found"})), - ) - .into_response(); - } - }; - - let repo_root_query = sqlx::query!("SELECT repo_root_cid FROM repos WHERE user_id = $1", user_id) - .fetch_optional(&state.db) - .await; - - let current_root_cid = match repo_root_query { - Ok(Some(row)) => { - let cid_str: String = row.repo_root_cid; - Cid::from_str(&cid_str).ok() - } - _ => None, - }; - - if current_root_cid.is_none() { - error!("Repo root not found for user {}", did); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError", "message": "Repo root not found"})), - ) - .into_response(); - } - let current_root_cid = current_root_cid.unwrap(); - - let commit_bytes = match state.block_store.get(¤t_root_cid).await { + let commit_bytes = match tracking_store.get(¤t_root_cid).await { Ok(Some(b)) => b, - Ok(None) => { - error!("Commit block not found: {}", current_root_cid); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } - Err(e) => { - error!("Failed to load commit block: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } + _ => return (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"error": "InternalError", "message": "Commit block not found"}))).into_response(), }; - let commit = match Commit::from_cbor(&commit_bytes) { Ok(c) => c, - Err(e) => { - error!("Failed to parse commit: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } + _ => return (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"error": "InternalError", "message": "Failed to parse commit"}))).into_response(), }; - let mst_root = commit.data; - let store = Arc::new(state.block_store.clone()); - let mst = Mst::load(store.clone(), mst_root, None); + let mst = Mst::load( + Arc::new(tracking_store.clone()), + commit.data, + None, + ); let collection_nsid = match input.collection.parse::() { Ok(n) => n, - Err(_) => { - return ( - StatusCode::BAD_REQUEST, - Json(json!({"error": "InvalidCollection"})), - ) - .into_response(); - } + Err(_) => return (StatusCode::BAD_REQUEST, Json(json!({"error": "InvalidCollection"}))).into_response(), }; - let rkey = input - .rkey - .unwrap_or_else(|| Utc::now().format("%Y%m%d%H%M%S%f").to_string()); - if input.validate.unwrap_or(true) { if input.collection == "app.bsky.feed.post" { if input.record.get("text").is_none() || input.record.get("createdAt").is_none() { @@ -191,130 +185,39 @@ pub async fn create_record( } } - let mut record_bytes = Vec::new(); - if let Err(e) = serde_ipld_dagcbor::to_writer(&mut record_bytes, &input.record) { - error!("Error serializing record: {:?}", e); - return ( - StatusCode::BAD_REQUEST, - Json(json!({"error": "InvalidRecord", "message": "Failed to serialize record"})), - ) - .into_response(); - } + let rkey = input.rkey.unwrap_or_else(|| Utc::now().format("%Y%m%d%H%M%S%f").to_string()); - let record_cid = match state.block_store.put(&record_bytes).await { + let mut record_bytes = Vec::new(); + if serde_ipld_dagcbor::to_writer(&mut record_bytes, &input.record).is_err() { + return (StatusCode::BAD_REQUEST, Json(json!({"error": "InvalidRecord", "message": "Failed to serialize record"}))).into_response(); + } + let record_cid = match tracking_store.put(&record_bytes).await { Ok(c) => c, - Err(e) => { - error!("Failed to save record block: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } + _ => return (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"error": "InternalError", "message": "Failed to save record block"}))).into_response(), }; let key = format!("{}/{}", collection_nsid, rkey); let new_mst = match mst.add(&key, record_cid).await { Ok(m) => m, - Err(e) => { - error!("Failed to add to MST: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } + _ => return (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"error": "InternalError", "message": "Failed to add to MST"}))).into_response(), }; - let new_mst_root = match new_mst.persist().await { Ok(c) => c, - Err(e) => { - error!("Failed to persist MST: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } + _ => return (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"error": "InternalError", "message": "Failed to persist MST"}))).into_response(), }; - let did_obj = match Did::new(&did) { - Ok(d) => d, - Err(_) => { - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError", "message": "Invalid DID"})), - ) - .into_response(); - } + let op = RecordOp::Create { collection: input.collection.clone(), rkey: rkey.clone(), cid: record_cid }; + let written_cids = tracking_store.get_written_cids(); + let written_cids_str = written_cids.iter().map(|c| c.to_string()).collect::>(); + + if let Err(e) = commit_and_log(&state, &did, user_id, Some(current_root_cid), new_mst_root, vec![op], &written_cids_str).await { + return (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"error": "InternalError", "message": e}))).into_response(); }; - let rev = Tid::now(LimitedU32::MIN); - - let new_commit = Commit::new_unsigned(did_obj, new_mst_root, rev, Some(current_root_cid)); - - let new_commit_bytes = match new_commit.to_cbor() { - Ok(b) => b, - Err(e) => { - error!("Failed to serialize new commit: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } - }; - - let new_root_cid = match state.block_store.put(&new_commit_bytes).await { - Ok(c) => c, - Err(e) => { - error!("Failed to save new commit: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } - }; - - let update_repo = sqlx::query!("UPDATE repos SET repo_root_cid = $1 WHERE user_id = $2", new_root_cid.to_string(), user_id) - .execute(&state.db) - .await; - - if let Err(e) = update_repo { - error!("Failed to update repo root in DB: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } - - let record_insert = sqlx::query!( - "INSERT INTO records (repo_id, collection, rkey, record_cid) VALUES ($1, $2, $3, $4) - ON CONFLICT (repo_id, collection, rkey) DO UPDATE SET record_cid = $4, created_at = NOW()", - user_id, - input.collection, - rkey, - record_cid.to_string() - ) - .execute(&state.db) - .await; - - if let Err(e) = record_insert { - error!("Error inserting record index: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError", "message": "Failed to index record"})), - ) - .into_response(); - } - - let output = CreateRecordOutput { - uri: format!("at://{}/{}/{}", input.repo, input.collection, rkey), + (StatusCode::OK, Json(CreateRecordOutput { + uri: format!("at://{}/{}/{}", did, input.collection, rkey), cid: record_cid.to_string(), - }; - (StatusCode::OK, Json(output)).into_response() + })).into_response() } #[derive(Deserialize)] @@ -340,142 +243,42 @@ pub struct PutRecordOutput { pub async fn put_record( State(state): State, - headers: axum::http::HeaderMap, + headers: HeaderMap, Json(input): Json, ) -> Response { - let auth_header = headers.get("Authorization"); - if auth_header.is_none() { - return ( - StatusCode::UNAUTHORIZED, - Json(json!({"error": "AuthenticationRequired"})), - ) - .into_response(); - } - let token = auth_header - .unwrap() - .to_str() - .unwrap_or("") - .replace("Bearer ", ""); + let (did, user_id, current_root_cid) = + match prepare_repo_write(&state, &headers, &input.repo).await { + Ok(res) => res, + Err(err_res) => return err_res, + }; - let session = sqlx::query!( - "SELECT s.did, k.key_bytes FROM sessions s JOIN users u ON s.did = u.did JOIN user_keys k ON u.id = k.user_id WHERE s.access_jwt = $1", - token - ) - .fetch_optional(&state.db) - .await - .unwrap_or(None); - - let (did, key_bytes) = match session { - Some(row) => ( - row.did, - row.key_bytes, - ), - None => { - return ( - StatusCode::UNAUTHORIZED, - Json(json!({"error": "AuthenticationFailed"})), - ) - .into_response(); + if let Some(swap_commit) = &input.swap_commit { + if Cid::from_str(swap_commit).ok() != Some(current_root_cid) { + return (StatusCode::CONFLICT, Json(json!({"error": "InvalidSwap", "message": "Repo has been modified"}))).into_response(); } - }; - - if let Err(_) = crate::auth::verify_token(&token, &key_bytes) { - return ( - StatusCode::UNAUTHORIZED, - Json(json!({"error": "AuthenticationFailed", "message": "Invalid token signature"})), - ) - .into_response(); } - if input.repo != did { - return (StatusCode::FORBIDDEN, Json(json!({"error": "InvalidRepo", "message": "Repo does not match authenticated user"}))).into_response(); - } + let tracking_store = TrackingBlockStore::new(state.block_store.clone()); - let user_query = sqlx::query!("SELECT id FROM users WHERE did = $1", did) - .fetch_optional(&state.db) - .await; - - let user_id: uuid::Uuid = match user_query { - Ok(Some(row)) => row.id, - _ => { - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError", "message": "User not found"})), - ) - .into_response(); - } - }; - - let repo_root_query = sqlx::query!("SELECT repo_root_cid FROM repos WHERE user_id = $1", user_id) - .fetch_optional(&state.db) - .await; - - let current_root_cid = match repo_root_query { - Ok(Some(row)) => { - let cid_str: String = row.repo_root_cid; - Cid::from_str(&cid_str).ok() - } - _ => None, - }; - - if current_root_cid.is_none() { - error!("Repo root not found for user {}", did); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError", "message": "Repo root not found"})), - ) - .into_response(); - } - let current_root_cid = current_root_cid.unwrap(); - - let commit_bytes = match state.block_store.get(¤t_root_cid).await { + let commit_bytes = match tracking_store.get(¤t_root_cid).await { Ok(Some(b)) => b, - Ok(None) => { - error!("Commit block not found: {}", current_root_cid); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError", "message": "Commit block not found"})), - ) - .into_response(); - } - Err(e) => { - error!("Failed to load commit block: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError", "message": "Failed to load commit block"})), - ) - .into_response(); - } + _ => return (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"error": "InternalError", "message": "Commit block not found"}))).into_response(), }; - let commit = match Commit::from_cbor(&commit_bytes) { Ok(c) => c, - Err(e) => { - error!("Failed to parse commit: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError", "message": "Failed to parse commit"})), - ) - .into_response(); - } + _ => return (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"error": "InternalError", "message": "Failed to parse commit"}))).into_response(), }; - let mst_root = commit.data; - let store = Arc::new(state.block_store.clone()); - let mst = Mst::load(store.clone(), mst_root, None); - + let mst = Mst::load( + Arc::new(tracking_store.clone()), + commit.data, + None, + ); let collection_nsid = match input.collection.parse::() { Ok(n) => n, - Err(_) => { - return ( - StatusCode::BAD_REQUEST, - Json(json!({"error": "InvalidCollection"})), - ) - .into_response(); - } + Err(_) => return (StatusCode::BAD_REQUEST, Json(json!({"error": "InvalidCollection"}))).into_response(), }; - - let rkey = input.rkey.clone(); + let key = format!("{}/{}", collection_nsid, input.rkey); if input.validate.unwrap_or(true) { if input.collection == "app.bsky.feed.post" { @@ -489,183 +292,47 @@ pub async fn put_record( } } - let mut record_bytes = Vec::new(); - if let Err(e) = serde_ipld_dagcbor::to_writer(&mut record_bytes, &input.record) { - error!("Error serializing record: {:?}", e); - return ( - StatusCode::BAD_REQUEST, - Json(json!({"error": "InvalidRecord", "message": "Failed to serialize record"})), - ) - .into_response(); - } - - let record_cid = match state.block_store.put(&record_bytes).await { - Ok(c) => c, - Err(e) => { - error!("Failed to save record block: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError", "message": "Failed to save record block"})), - ) - .into_response(); - } - }; - - let key = format!("{}/{}", collection_nsid, rkey); - - let existing = match mst.get(&key).await { - Ok(v) => v, - Err(e) => { - error!("Failed to check MST key: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json( - json!({"error": "InternalError", "message": "Failed to check existing record"}), - ), - ) - .into_response(); - } - }; - if let Some(swap_record_str) = &input.swap_record { - let swap_record_cid = match Cid::from_str(swap_record_str) { - Ok(c) => c, - Err(_) => { - return ( - StatusCode::BAD_REQUEST, - Json( - json!({"error": "InvalidSwapRecord", "message": "Invalid swapRecord CID"}), - ), - ) - .into_response(); - } - }; - match &existing { - Some(current_cid) if *current_cid != swap_record_cid => { - return ( - StatusCode::CONFLICT, - Json(json!({"error": "InvalidSwap", "message": "Record has been modified"})), - ) - .into_response(); - } - None => { - return ( - StatusCode::CONFLICT, - Json(json!({"error": "InvalidSwap", "message": "Record does not exist"})), - ) - .into_response(); - } - _ => {} + let expected_cid = Cid::from_str(swap_record_str).ok(); + let actual_cid = mst.get(&key).await.ok().flatten(); + if expected_cid != actual_cid { + return (StatusCode::CONFLICT, Json(json!({"error": "InvalidSwap", "message": "Record has been modified or does not exist"}))).into_response(); } } - let new_mst = if existing.is_some() { - match mst.update(&key, record_cid).await { - Ok(m) => m, - Err(e) => { - error!("Failed to update MST: {:?}", e); - return (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"error": "InternalError", "message": format!("Failed to update MST: {:?}", e)}))).into_response(); - } - } + let existing_cid = mst.get(&key).await.ok().flatten(); + + let mut record_bytes = Vec::new(); + if serde_ipld_dagcbor::to_writer(&mut record_bytes, &input.record).is_err() { + return (StatusCode::BAD_REQUEST, Json(json!({"error": "InvalidRecord", "message": "Failed to serialize record"}))).into_response(); + } + let record_cid = match tracking_store.put(&record_bytes).await { + Ok(c) => c, + _ => return (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"error": "InternalError", "message": "Failed to save record block"}))).into_response(), + }; + + let new_mst = if existing_cid.is_some() { + mst.update(&key, record_cid).await.unwrap() } else { - match mst.add(&key, record_cid).await { - Ok(m) => m, - Err(e) => { - error!("Failed to add to MST: {:?}", e); - return (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"error": "InternalError", "message": format!("Failed to add to MST: {:?}", e)}))).into_response(); - } - } + mst.add(&key, record_cid).await.unwrap() + }; + let new_mst_root = new_mst.persist().await.unwrap(); + + let op = if existing_cid.is_some() { + RecordOp::Update { collection: input.collection.clone(), rkey: input.rkey.clone(), cid: record_cid } + } else { + RecordOp::Create { collection: input.collection.clone(), rkey: input.rkey.clone(), cid: record_cid } }; - let new_mst_root = match new_mst.persist().await { - Ok(c) => c, - Err(e) => { - error!("Failed to persist MST: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError", "message": "Failed to persist MST"})), - ) - .into_response(); - } + let written_cids = tracking_store.get_written_cids(); + let written_cids_str = written_cids.iter().map(|c| c.to_string()).collect::>(); + + if let Err(e) = commit_and_log(&state, &did, user_id, Some(current_root_cid), new_mst_root, vec![op], &written_cids_str).await { + return (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"error": "InternalError", "message": e}))).into_response(); }; - let did_obj = match Did::new(&did) { - Ok(d) => d, - Err(_) => { - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError", "message": "Invalid DID"})), - ) - .into_response(); - } - }; - - let rev = Tid::now(LimitedU32::MIN); - - let new_commit = Commit::new_unsigned(did_obj, new_mst_root, rev, Some(current_root_cid)); - - let new_commit_bytes = match new_commit.to_cbor() { - Ok(b) => b, - Err(e) => { - error!("Failed to serialize new commit: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json( - json!({"error": "InternalError", "message": "Failed to serialize new commit"}), - ), - ) - .into_response(); - } - }; - - let new_root_cid = match state.block_store.put(&new_commit_bytes).await { - Ok(c) => c, - Err(e) => { - error!("Failed to save new commit: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError", "message": "Failed to save new commit"})), - ) - .into_response(); - } - }; - - let update_repo = sqlx::query!("UPDATE repos SET repo_root_cid = $1 WHERE user_id = $2", new_root_cid.to_string(), user_id) - .execute(&state.db) - .await; - - if let Err(e) = update_repo { - error!("Failed to update repo root in DB: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError", "message": "Failed to update repo root in DB"})), - ) - .into_response(); - } - - let record_insert = sqlx::query!( - "INSERT INTO records (repo_id, collection, rkey, record_cid) VALUES ($1, $2, $3, $4) - ON CONFLICT (repo_id, collection, rkey) DO UPDATE SET record_cid = $4, created_at = NOW()", - user_id, - input.collection, - rkey, - record_cid.to_string() - ) - .execute(&state.db) - .await; - - if let Err(e) = record_insert { - error!("Error inserting record index: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError", "message": "Failed to index record"})), - ) - .into_response(); - } - - let output = PutRecordOutput { - uri: format!("at://{}/{}/{}", input.repo, input.collection, rkey), + (StatusCode::OK, Json(PutRecordOutput { + uri: format!("at://{}/{}/{}", did, input.collection, input.rkey), cid: record_cid.to_string(), - }; - (StatusCode::OK, Json(output)).into_response() + })).into_response() } diff --git a/src/lib.rs b/src/lib.rs index 8fc47de..715e6b5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -131,6 +131,10 @@ pub fn app(state: AppState) -> Router { "/xrpc/com.atproto.sync.getRecord", get(sync::get_record), ) + .route( + "/xrpc/com.atproto.sync.subscribeRepos", + get(sync::subscribe_repos), + ) .route( "/xrpc/com.atproto.moderation.createReport", post(api::moderation::create_report), diff --git a/src/main.rs b/src/main.rs index aecb98a..3e8539b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -24,6 +24,15 @@ async fn main() { let state = AppState::new(pool.clone()).await; + bspds::sync::listener::start_sequencer_listener(state.clone()).await; + let relays = std::env::var("RELAYS") + .unwrap_or_default() + .split(',') + .filter(|s| !s.is_empty()) + .map(|s| s.to_string()) + .collect(); + bspds::sync::relay_client::start_relay_clients(state.clone(), relays, None).await; + let (shutdown_tx, shutdown_rx) = watch::channel(false); let mut notification_service = NotificationService::new(pool); diff --git a/src/repo/mod.rs b/src/repo/mod.rs index 8097b71..37b669f0 100644 --- a/src/repo/mod.rs +++ b/src/repo/mod.rs @@ -7,6 +7,8 @@ use multihash::Multihash; use sha2::{Digest, Sha256}; use sqlx::PgPool; +pub mod tracking; + #[derive(Clone)] pub struct PostgresBlockStore { pool: PgPool, diff --git a/src/repo/tracking.rs b/src/repo/tracking.rs new file mode 100644 index 0000000..5170d45 --- /dev/null +++ b/src/repo/tracking.rs @@ -0,0 +1,62 @@ +use crate::repo::PostgresBlockStore; +use bytes::Bytes; +use cid::Cid; +use jacquard_repo::error::RepoError; +use jacquard_repo::repo::CommitData; +use jacquard_repo::storage::BlockStore; +use std::sync::{Arc, Mutex}; + +#[derive(Clone)] +pub struct TrackingBlockStore { + inner: PostgresBlockStore, + written_cids: Arc>>, +} + +impl TrackingBlockStore { + pub fn new(store: PostgresBlockStore) -> Self { + Self { + inner: store, + written_cids: Arc::new(Mutex::new(Vec::new())), + } + } + + pub fn get_written_cids(&self) -> Vec { + self.written_cids.lock().unwrap().clone() + } +} + +impl BlockStore for TrackingBlockStore { + async fn get(&self, cid: &Cid) -> Result, RepoError> { + self.inner.get(cid).await + } + + async fn put(&self, data: &[u8]) -> Result { + let cid = self.inner.put(data).await?; + self.written_cids.lock().unwrap().push(cid.clone()); + Ok(cid) + } + + async fn has(&self, cid: &Cid) -> Result { + self.inner.has(cid).await + } + + async fn put_many( + &self, + blocks: impl IntoIterator + Send, + ) -> Result<(), RepoError> { + let blocks: Vec<_> = blocks.into_iter().collect(); + let cids: Vec = blocks.iter().map(|(cid, _)| cid.clone()).collect(); + self.inner.put_many(blocks).await?; + self.written_cids.lock().unwrap().extend(cids); + Ok(()) + } + + async fn get_many(&self, cids: &[Cid]) -> Result>, RepoError> { + self.inner.get_many(cids).await + } + + async fn apply_commit(&self, commit: CommitData) -> Result<(), RepoError> { + self.put_many(commit.blocks).await?; + Ok(()) + } +} diff --git a/src/state.rs b/src/state.rs index 3bc2a46..a7d9221 100644 --- a/src/state.rs +++ b/src/state.rs @@ -1,23 +1,28 @@ use crate::repo::PostgresBlockStore; use crate::storage::{BlobStorage, S3BlobStorage}; +use crate::sync::firehose::SequencedEvent; use sqlx::PgPool; use std::sync::Arc; +use tokio::sync::broadcast; #[derive(Clone)] pub struct AppState { pub db: PgPool, pub block_store: PostgresBlockStore, pub blob_store: Arc, + pub firehose_tx: broadcast::Sender, } impl AppState { pub async fn new(db: PgPool) -> Self { let block_store = PostgresBlockStore::new(db.clone()); let blob_store = S3BlobStorage::new().await; + let (firehose_tx, _) = broadcast::channel(1000); Self { db, block_store, blob_store: Arc::new(blob_store), + firehose_tx, } } } diff --git a/src/sync/firehose.rs b/src/sync/firehose.rs new file mode 100644 index 0000000..74ee2e4 --- /dev/null +++ b/src/sync/firehose.rs @@ -0,0 +1,16 @@ +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use chrono::{DateTime, Utc}; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SequencedEvent { + pub seq: i64, + pub did: String, + pub created_at: DateTime, + pub event_type: String, + pub commit_cid: Option, + pub prev_cid: Option, + pub ops: Option, + pub blobs: Option>, + pub blocks_cids: Option>, +} diff --git a/src/sync/frame.rs b/src/sync/frame.rs new file mode 100644 index 0000000..160f9a4 --- /dev/null +++ b/src/sync/frame.rs @@ -0,0 +1,59 @@ +use serde::{Deserialize, Serialize}; +use crate::sync::firehose::SequencedEvent; + +#[derive(Debug, Serialize, Deserialize)] +pub struct Frame { + #[serde(rename = "op")] + pub op: i64, + #[serde(rename = "d")] + pub data: FrameData, +} + +#[derive(Debug, Serialize, Deserialize)] +#[serde(untagged)] +pub enum FrameData { + Commit(Box), +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct CommitFrame { + pub seq: i64, + pub rebase: bool, + #[serde(rename = "tooBig")] + pub too_big: bool, + pub repo: String, + pub commit: String, + pub prev: Option, + #[serde(with = "serde_bytes")] + pub blocks: Vec, + pub ops: Vec, + pub blobs: Vec, + pub time: String, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct RepoOp { + pub action: String, + pub path: String, + pub cid: Option, +} + +impl From for CommitFrame { + fn from(event: SequencedEvent) -> Self { + let ops = serde_json::from_value::>(event.ops.unwrap_or_default()) + .unwrap_or_else(|_| vec![]); + + CommitFrame { + seq: event.seq, + rebase: false, + too_big: false, + repo: event.did, + commit: event.commit_cid.unwrap_or_default(), + prev: event.prev_cid, + blocks: Vec::new(), + ops, + blobs: event.blobs.unwrap_or_default(), + time: event.created_at.to_rfc3339(), + } + } +} diff --git a/src/sync/listener.rs b/src/sync/listener.rs new file mode 100644 index 0000000..05636be --- /dev/null +++ b/src/sync/listener.rs @@ -0,0 +1,53 @@ +use crate::state::AppState; +use crate::sync::firehose::SequencedEvent; +use sqlx::postgres::PgListener; +use tracing::{error, info, warn}; + +pub async fn start_sequencer_listener(state: AppState) { + tokio::spawn(async move { + info!("Starting sequencer listener background task"); + loop { + if let Err(e) = listen_loop(state.clone()).await { + error!("Sequencer listener failed: {}. Restarting in 5s...", e); + tokio::time::sleep(tokio::time::Duration::from_secs(5)).await; + } + } + }); +} + +async fn listen_loop(state: AppState) -> anyhow::Result<()> { + let mut listener = PgListener::connect_with(&state.db).await?; + listener.listen("repo_updates").await?; + info!("Connected to Postgres and listening for 'repo_updates'"); + + loop { + let notification = listener.recv().await?; + let payload = notification.payload(); + + let seq_id: i64 = match payload.parse() { + Ok(id) => id, + Err(e) => { + warn!("Received invalid payload in repo_updates: '{}'. Error: {}", payload, e); + continue; + } + }; + + let event = sqlx::query_as!( + SequencedEvent, + r#" + SELECT seq, did, created_at, event_type, commit_cid, prev_cid, ops, blobs, blocks_cids + FROM repo_seq + WHERE seq = $1 + "#, + seq_id + ) + .fetch_optional(&state.db) + .await?; + + if let Some(event) = event { + let _ = state.firehose_tx.send(event); + } else { + warn!("Received notification for seq {} but could not find row in repo_seq", seq_id); + } + } +} diff --git a/src/sync/mod.rs b/src/sync/mod.rs index c363bf2..2f3965c 100644 --- a/src/sync/mod.rs +++ b/src/sync/mod.rs @@ -2,9 +2,16 @@ pub mod blob; pub mod car; pub mod commit; pub mod crawl; +pub mod firehose; +pub mod frame; +pub mod listener; +pub mod relay_client; pub mod repo; +pub mod subscribe_repos; +pub mod util; pub use blob::{get_blob, list_blobs}; pub use commit::{get_latest_commit, get_repo_status, list_repos}; pub use crawl::{notify_of_update, request_crawl}; -pub use repo::{get_blocks, get_record, get_repo}; +pub use repo::{get_blocks, get_repo, get_record}; +pub use subscribe_repos::subscribe_repos; diff --git a/src/sync/relay_client.rs b/src/sync/relay_client.rs new file mode 100644 index 0000000..86e59e1 --- /dev/null +++ b/src/sync/relay_client.rs @@ -0,0 +1,84 @@ +use crate::state::AppState; +use crate::sync::util::format_event_for_sending; +use futures::{sink::SinkExt, stream::StreamExt}; +use std::time::Duration; +use tokio::sync::mpsc; +use tokio_tungstenite::{connect_async, tungstenite::Message}; +use tracing::{error, info, warn}; + +async fn run_relay_client(state: AppState, url: String, ready_tx: Option>) { + info!("Starting firehose client for relay: {}", url); + loop { + match connect_async(&url).await { + Ok((mut ws_stream, _)) => { + info!("Connected to firehose relay: {}", url); + if let Some(tx) = ready_tx.as_ref() { + tx.send(()).await.ok(); + } + + let mut rx = state.firehose_tx.subscribe(); + + loop { + tokio::select! { + Ok(event) = rx.recv() => { + match format_event_for_sending(&state, event).await { + Ok(bytes) => { + if let Err(e) = ws_stream.send(Message::Binary(bytes.into())).await { + warn!("Failed to send event to {}: {}. Disconnecting.", url, e); + break; + } + } + Err(e) => { + error!("Failed to format event for relay {}: {}", url, e); + } + } + } + Some(msg) = ws_stream.next() => { + if let Ok(Message::Close(_)) = msg { + warn!("Relay {} closed connection.", url); + break; + } + } + else => break, + } + } + } + Err(e) => { + error!("Failed to connect to firehose relay {}: {}", url, e); + } + } + warn!( + "Disconnected from {}. Reconnecting in 5 seconds...", + url + ); + tokio::time::sleep(Duration::from_secs(5)).await; + } +} + +pub async fn start_relay_clients( + state: AppState, + relays: Vec, + mut ready_rx: Option>, +) { + if relays.is_empty() { + return; + } + + let (ready_tx, mut internal_ready_rx) = mpsc::channel(1); + + for url in relays { + let ready_tx = if ready_rx.is_some() { + Some(ready_tx.clone()) + } else { + None + }; + tokio::spawn(run_relay_client(state.clone(), url, ready_tx)); + } + + if let Some(mut rx) = ready_rx.take() { + tokio::spawn(async move { + internal_ready_rx.recv().await; + rx.close(); + }); + } +} diff --git a/src/sync/repo.rs b/src/sync/repo.rs index b9d27c6..e5a9d3a 100644 --- a/src/sync/repo.rs +++ b/src/sync/repo.rs @@ -1,342 +1,212 @@ use crate::state::AppState; -use crate::sync::car::{encode_car_header, ld_write}; +use crate::sync::car::encode_car_header; use axum::{ - Json, - body::Body, extract::{Query, State}, http::StatusCode, - http::header, response::{IntoResponse, Response}, + Json, }; -use bytes::Bytes; use cid::Cid; -use jacquard_repo::{commit::Commit, storage::BlockStore}; +use jacquard_repo::storage::BlockStore; use serde::Deserialize; use serde_json::json; -use std::collections::HashSet; +use std::io::Write; +use std::str::FromStr; use tracing::error; #[derive(Deserialize)] -pub struct GetBlocksParams { +pub struct GetBlocksQuery { pub did: String, pub cids: String, } pub async fn get_blocks( State(state): State, - Query(params): Query, + Query(query): Query, ) -> Response { - let did = params.did.trim(); - - if did.is_empty() { - return ( - StatusCode::BAD_REQUEST, - Json(json!({"error": "InvalidRequest", "message": "did is required"})), - ) - .into_response(); - } - - let cid_strings: Vec<&str> = params.cids.split(',').map(|s| s.trim()).filter(|s| !s.is_empty()).collect(); - - if cid_strings.is_empty() { - return ( - StatusCode::BAD_REQUEST, - Json(json!({"error": "InvalidRequest", "message": "cids is required"})), - ) - .into_response(); - } - - let repo_result = sqlx::query!( - r#" - SELECT r.repo_root_cid - FROM repos r - JOIN users u ON r.user_id = u.id - WHERE u.did = $1 - "#, - did - ) - .fetch_optional(&state.db) - .await; - - let repo_root_cid_str = match repo_result { - Ok(Some(row)) => row.repo_root_cid, - Ok(None) => { - return ( - StatusCode::NOT_FOUND, - Json(json!({"error": "RepoNotFound", "message": "Could not find repo for DID"})), - ) - .into_response(); - } - Err(e) => { - error!("DB error in get_blocks: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } - }; - - let root_cid = match repo_root_cid_str.parse::() { - Ok(c) => c, - Err(e) => { - error!("Failed to parse root CID: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } - }; - - let mut requested_cids: Vec = Vec::new(); - for cid_str in &cid_strings { - match cid_str.parse::() { - Ok(c) => requested_cids.push(c), - Err(e) => { - error!("Failed to parse CID '{}': {:?}", cid_str, e); - return ( - StatusCode::BAD_REQUEST, - Json(json!({"error": "InvalidRequest", "message": format!("Invalid CID: {}", cid_str)})), - ) - .into_response(); - } - } - } - - let mut buf = Vec::new(); - let car_header = encode_car_header(&root_cid); - if let Err(e) = ld_write(&mut buf, &car_header) { - error!("Failed to write CAR header: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } - - for cid in &requested_cids { - let cid_bytes = cid.to_bytes(); - let block_result = sqlx::query!( - "SELECT data FROM blocks WHERE cid = $1", - &cid_bytes - ) + let user_exists = sqlx::query!("SELECT id FROM users WHERE did = $1", query.did) .fetch_optional(&state.db) - .await; + .await + .unwrap_or(None); - match block_result { - Ok(Some(row)) => { - let mut block_data = Vec::new(); - block_data.extend_from_slice(&cid_bytes); - block_data.extend_from_slice(&row.data); - if let Err(e) = ld_write(&mut buf, &block_data) { - error!("Failed to write block: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } - } - Ok(None) => { - return ( - StatusCode::NOT_FOUND, - Json(json!({"error": "BlockNotFound", "message": format!("Block not found: {}", cid)})), - ) - .into_response(); - } - Err(e) => { - error!("DB error fetching block: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } + if user_exists.is_none() { + return (StatusCode::NOT_FOUND, "Repo not found").into_response(); + } + + let cids_str: Vec<&str> = query.cids.split(',').collect(); + let mut cids = Vec::new(); + for s in cids_str { + match Cid::from_str(s) { + Ok(cid) => cids.push(cid), + Err(_) => return (StatusCode::BAD_REQUEST, "Invalid CID").into_response(), } } - Response::builder() - .status(StatusCode::OK) - .header(header::CONTENT_TYPE, "application/vnd.ipld.car") - .body(Body::from(buf)) - .unwrap() + let blocks_res = state.block_store.get_many(&cids).await; + let blocks = match blocks_res { + Ok(blocks) => blocks, + Err(e) => { + error!("Failed to get blocks: {}", e); + return (StatusCode::INTERNAL_SERVER_ERROR, "Failed to get blocks").into_response(); + } + }; + + let root_cid = cids.first().cloned().unwrap_or_default(); + + if cids.is_empty() { + return (StatusCode::BAD_REQUEST, "No CIDs provided").into_response(); + } + + let header = encode_car_header(&root_cid); + + let mut car_bytes = header; + + for (i, block_opt) in blocks.into_iter().enumerate() { + if let Some(block) = block_opt { + let cid = cids[i]; + let cid_bytes = cid.to_bytes(); + let total_len = cid_bytes.len() + block.len(); + + let mut writer = Vec::new(); + crate::sync::car::write_varint(&mut writer, total_len as u64).unwrap(); + writer.write_all(&cid_bytes).unwrap(); + writer.write_all(&block).unwrap(); + + car_bytes.extend_from_slice(&writer); + } + } + + ( + StatusCode::OK, + [(axum::http::header::CONTENT_TYPE, "application/vnd.ipld.car")], + car_bytes, + ) + .into_response() } #[derive(Deserialize)] -pub struct GetRepoParams { +pub struct GetRepoQuery { pub did: String, pub since: Option, } pub async fn get_repo( State(state): State, - Query(params): Query, + Query(query): Query, ) -> Response { - let did = params.did.trim(); + let repo_row = sqlx::query!( + r#" + SELECT r.repo_root_cid + FROM repos r + JOIN users u ON u.id = r.user_id + WHERE u.did = $1 + "#, + query.did + ) + .fetch_optional(&state.db) + .await + .unwrap_or(None); - if did.is_empty() { - return ( - StatusCode::BAD_REQUEST, - Json(json!({"error": "InvalidRequest", "message": "did is required"})), - ) - .into_response(); - } + let head_str = match repo_row { + Some(r) => r.repo_root_cid, + None => { + let user_exists = sqlx::query!("SELECT id FROM users WHERE did = $1", query.did) + .fetch_optional(&state.db) + .await + .unwrap_or(None); - let user_result = sqlx::query!("SELECT id FROM users WHERE did = $1", did) - .fetch_optional(&state.db) - .await; - - let user_id = match user_result { - Ok(Some(row)) => row.id, - Ok(None) => { - return ( - StatusCode::NOT_FOUND, - Json(json!({"error": "RepoNotFound", "message": "Could not find repo for DID"})), - ) - .into_response(); - } - Err(e) => { - error!("DB error in get_repo: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); + if user_exists.is_none() { + return ( + StatusCode::NOT_FOUND, + Json(json!({"error": "RepoNotFound", "message": "Repo not found"})), + ) + .into_response(); + } else { + return ( + StatusCode::NOT_FOUND, + Json(json!({"error": "RepoNotFound", "message": "Repo not initialized"})), + ) + .into_response(); + } } }; - let repo_result = sqlx::query!("SELECT repo_root_cid FROM repos WHERE user_id = $1", user_id) - .fetch_optional(&state.db) - .await; - - let repo_root_cid_str = match repo_result { - Ok(Some(row)) => row.repo_root_cid, - Ok(None) => { - return ( - StatusCode::NOT_FOUND, - Json(json!({"error": "RepoNotFound", "message": "Repository not initialized"})), - ) - .into_response(); - } - Err(e) => { - error!("DB error in get_repo: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } - }; - - let root_cid = match repo_root_cid_str.parse::() { + let head_cid = match Cid::from_str(&head_str) { Ok(c) => c, - Err(e) => { - error!("Failed to parse root CID: {:?}", e); + Err(_) => { return ( StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), + Json(json!({"error": "InternalError", "message": "Invalid head CID"})), ) .into_response(); } }; - let commit_bytes = match state.block_store.get(&root_cid).await { - Ok(Some(b)) => b, - Ok(None) => { - error!("Commit block not found: {}", root_cid); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); + let mut car_bytes = encode_car_header(&head_cid); + + let mut stack = vec![head_cid]; + let mut visited = std::collections::HashSet::new(); + let mut limit = 20000; + + while let Some(cid) = stack.pop() { + if visited.contains(&cid) { + continue; } - Err(e) => { - error!("Failed to load commit block: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } - }; + visited.insert(cid); + if limit == 0 { break; } + limit -= 1; - let commit = match Commit::from_cbor(&commit_bytes) { - Ok(c) => c, - Err(e) => { - error!("Failed to parse commit: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } - }; + if let Ok(Some(block)) = state.block_store.get(&cid).await { + let cid_bytes = cid.to_bytes(); + let total_len = cid_bytes.len() + block.len(); + let mut writer = Vec::new(); + crate::sync::car::write_varint(&mut writer, total_len as u64).unwrap(); + writer.write_all(&cid_bytes).unwrap(); + writer.write_all(&block).unwrap(); + car_bytes.extend_from_slice(&writer); - let mut collected_blocks: Vec<(Cid, Bytes)> = Vec::new(); - let mut visited: HashSet> = HashSet::new(); - - collected_blocks.push((root_cid, commit_bytes.clone())); - visited.insert(root_cid.to_bytes()); - - let mst_root_cid = commit.data; - if !visited.contains(&mst_root_cid.to_bytes()) { - visited.insert(mst_root_cid.to_bytes()); - if let Ok(Some(data)) = state.block_store.get(&mst_root_cid).await { - collected_blocks.push((mst_root_cid, data)); - } - } - - let records = sqlx::query!("SELECT record_cid FROM records WHERE repo_id = $1", user_id) - .fetch_all(&state.db) - .await - .unwrap_or_default(); - - for record in records { - if let Ok(cid) = record.record_cid.parse::() { - if !visited.contains(&cid.to_bytes()) { - visited.insert(cid.to_bytes()); - if let Ok(Some(data)) = state.block_store.get(&cid).await { - collected_blocks.push((cid, data)); - } + if let Ok(value) = serde_ipld_dagcbor::from_slice::(&block) { + extract_links_json(&value, &mut stack); } } } - let mut buf = Vec::new(); - let car_header = encode_car_header(&root_cid); - if let Err(e) = ld_write(&mut buf, &car_header) { - error!("Failed to write CAR header: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } + ( + StatusCode::OK, + [(axum::http::header::CONTENT_TYPE, "application/vnd.ipld.car")], + car_bytes, + ) + .into_response() +} - for (cid, data) in &collected_blocks { - let mut block_data = Vec::new(); - block_data.extend_from_slice(&cid.to_bytes()); - block_data.extend_from_slice(data); - if let Err(e) = ld_write(&mut buf, &block_data) { - error!("Failed to write block: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); +fn extract_links_json(value: &serde_json::Value, stack: &mut Vec) { + match value { + serde_json::Value::Object(map) => { + if let Some(serde_json::Value::String(s)) = map.get("/") { + if let Ok(cid) = Cid::from_str(s) { + stack.push(cid); + } + } else if let Some(serde_json::Value::String(s)) = map.get("$link") { + if let Ok(cid) = Cid::from_str(s) { + stack.push(cid); + } + } else { + for v in map.values() { + extract_links_json(v, stack); + } + } } + serde_json::Value::Array(arr) => { + for v in arr { + extract_links_json(v, stack); + } + } + _ => {} } - - Response::builder() - .status(StatusCode::OK) - .header(header::CONTENT_TYPE, "application/vnd.ipld.car") - .body(Body::from(buf)) - .unwrap() } #[derive(Deserialize)] -pub struct GetRecordParams { +pub struct GetRecordQuery { pub did: String, pub collection: String, pub rkey: String, @@ -344,213 +214,71 @@ pub struct GetRecordParams { pub async fn get_record( State(state): State, - Query(params): Query, + Query(query): Query, ) -> Response { - let did = params.did.trim(); - let collection = params.collection.trim(); - let rkey = params.rkey.trim(); - - if did.is_empty() { - return ( - StatusCode::BAD_REQUEST, - Json(json!({"error": "InvalidRequest", "message": "did is required"})), - ) - .into_response(); - } - - if collection.is_empty() { - return ( - StatusCode::BAD_REQUEST, - Json(json!({"error": "InvalidRequest", "message": "collection is required"})), - ) - .into_response(); - } - - if rkey.is_empty() { - return ( - StatusCode::BAD_REQUEST, - Json(json!({"error": "InvalidRequest", "message": "rkey is required"})), - ) - .into_response(); - } - - let user_result = sqlx::query!("SELECT id FROM users WHERE did = $1", did) + let user = sqlx::query!("SELECT id FROM users WHERE did = $1", query.did) .fetch_optional(&state.db) - .await; + .await + .unwrap_or(None); - let user_id = match user_result { - Ok(Some(row)) => row.id, - Ok(None) => { - return ( + let user_id = match user { + Some(u) => u.id, + None => { + return ( StatusCode::NOT_FOUND, - Json(json!({"error": "RepoNotFound", "message": "Could not find repo for DID"})), - ) - .into_response(); - } - Err(e) => { - error!("DB error in sync get_record: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), + Json(json!({"error": "RepoNotFound", "message": "Repo not found"})), ) .into_response(); } }; - let record_result = sqlx::query!( + let record = sqlx::query!( "SELECT record_cid FROM records WHERE repo_id = $1 AND collection = $2 AND rkey = $3", user_id, - collection, - rkey + query.collection, + query.rkey ) .fetch_optional(&state.db) - .await; + .await + .unwrap_or(None); - let record_cid_str = match record_result { - Ok(Some(row)) => row.record_cid, - Ok(None) => { - return ( + let record_cid_str = match record { + Some(r) => r.record_cid, + None => { + return ( StatusCode::NOT_FOUND, Json(json!({"error": "RecordNotFound", "message": "Record not found"})), ) .into_response(); } - Err(e) => { - error!("DB error in sync get_record: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } }; - let record_cid = match record_cid_str.parse::() { + let cid = match Cid::from_str(&record_cid_str) { Ok(c) => c, - Err(e) => { - error!("Failed to parse record CID: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } + Err(_) => return (StatusCode::INTERNAL_SERVER_ERROR, "Invalid CID").into_response(), }; - let repo_result = sqlx::query!("SELECT repo_root_cid FROM repos WHERE user_id = $1", user_id) - .fetch_optional(&state.db) - .await; - - let repo_root_cid_str = match repo_result { - Ok(Some(row)) => row.repo_root_cid, - Ok(None) => { - return ( - StatusCode::NOT_FOUND, - Json(json!({"error": "RepoNotFound", "message": "Repository not initialized"})), - ) - .into_response(); - } - Err(e) => { - error!("DB error in sync get_record: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } - }; - - let root_cid = match repo_root_cid_str.parse::() { - Ok(c) => c, - Err(e) => { - error!("Failed to parse root CID: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } - }; - - let mut collected_blocks: Vec<(Cid, Bytes)> = Vec::new(); - - let commit_bytes = match state.block_store.get(&root_cid).await { + let block_res = state.block_store.get(&cid).await; + let block = match block_res { Ok(Some(b)) => b, - Ok(None) => { - error!("Commit block not found: {}", root_cid); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } - Err(e) => { - error!("Failed to load commit block: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } + _ => return (StatusCode::NOT_FOUND, "Block not found").into_response(), }; - collected_blocks.push((root_cid, commit_bytes.clone())); + let header = encode_car_header(&cid); + let mut car_bytes = header; - let commit = match Commit::from_cbor(&commit_bytes) { - Ok(c) => c, - Err(e) => { - error!("Failed to parse commit: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } - }; + let cid_bytes = cid.to_bytes(); + let total_len = cid_bytes.len() + block.len(); + let mut writer = Vec::new(); + crate::sync::car::write_varint(&mut writer, total_len as u64).unwrap(); + writer.write_all(&cid_bytes).unwrap(); + writer.write_all(&block).unwrap(); + car_bytes.extend_from_slice(&writer); - let mst_root_cid = commit.data; - if let Ok(Some(data)) = state.block_store.get(&mst_root_cid).await { - collected_blocks.push((mst_root_cid, data)); - } - - if let Ok(Some(data)) = state.block_store.get(&record_cid).await { - collected_blocks.push((record_cid, data)); - } else { - return ( - StatusCode::NOT_FOUND, - Json(json!({"error": "RecordNotFound", "message": "Record block not found"})), - ) - .into_response(); - } - - let mut buf = Vec::new(); - let car_header = encode_car_header(&root_cid); - if let Err(e) = ld_write(&mut buf, &car_header) { - error!("Failed to write CAR header: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } - - for (cid, data) in &collected_blocks { - let mut block_data = Vec::new(); - block_data.extend_from_slice(&cid.to_bytes()); - block_data.extend_from_slice(data); - if let Err(e) = ld_write(&mut buf, &block_data) { - error!("Failed to write block: {:?}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "InternalError"})), - ) - .into_response(); - } - } - - Response::builder() - .status(StatusCode::OK) - .header(header::CONTENT_TYPE, "application/vnd.ipld.car") - .body(Body::from(buf)) - .unwrap() + ( + StatusCode::OK, + [(axum::http::header::CONTENT_TYPE, "application/vnd.ipld.car")], + car_bytes, + ) + .into_response() } diff --git a/src/sync/subscribe_repos.rs b/src/sync/subscribe_repos.rs new file mode 100644 index 0000000..983f841 --- /dev/null +++ b/src/sync/subscribe_repos.rs @@ -0,0 +1,91 @@ +use crate::state::AppState; +use crate::sync::firehose::SequencedEvent; +use crate::sync::util::format_event_for_sending; +use axum::{ + extract::{ws::Message, ws::WebSocket, ws::WebSocketUpgrade, Query, State}, + response::Response, +}; +use futures::{sink::SinkExt, stream::StreamExt}; +use serde::Deserialize; +use tracing::{error, info, warn}; + +#[derive(Deserialize)] +pub struct SubscribeReposParams { + pub cursor: Option, +} + +#[axum::debug_handler] +pub async fn subscribe_repos( + ws: WebSocketUpgrade, + State(state): State, + Query(params): Query, +) -> Response { + ws.on_upgrade(move |socket| handle_socket(socket, state, params)) +} + +async fn send_event( + socket: &mut WebSocket, + state: &AppState, + event: SequencedEvent, +) -> Result<(), anyhow::Error> { + let bytes = format_event_for_sending(state, event).await?; + socket.send(Message::Binary(bytes.into())).await?; + Ok(()) +} + +async fn handle_socket(mut socket: WebSocket, state: AppState, params: SubscribeReposParams) { + info!(cursor = ?params.cursor, "New firehose subscriber"); + + if let Some(cursor) = params.cursor { + let events = sqlx::query_as!( + SequencedEvent, + r#" + SELECT seq, did, created_at, event_type, commit_cid, prev_cid, ops, blobs, blocks_cids + FROM repo_seq + WHERE seq > $1 + ORDER BY seq ASC + "#, + cursor + ) + .fetch_all(&state.db) + .await; + + match events { + Ok(events) => { + for event in events { + if let Err(e) = send_event(&mut socket, &state, event).await { + warn!("Failed to send backfill event: {}", e); + return; + } + } + } + Err(e) => { + error!("Failed to fetch backfill events: {}", e); + socket.close().await.ok(); + return; + } + } + } + + let mut rx = state.firehose_tx.subscribe(); + + loop { + tokio::select! { + Ok(event) = rx.recv() => { + if let Err(e) = send_event(&mut socket, &state, event).await { + warn!("Failed to send event: {}", e); + break; + } + } + Some(Ok(msg)) = socket.next() => { + if let Message::Close(_) = msg { + info!("Client closed connection"); + break; + } + } + else => { + break; + } + } + } +} diff --git a/src/sync/util.rs b/src/sync/util.rs new file mode 100644 index 0000000..c2be79d --- /dev/null +++ b/src/sync/util.rs @@ -0,0 +1,52 @@ +use crate::state::AppState; +use crate::sync::firehose::SequencedEvent; +use crate::sync::frame::{CommitFrame, Frame, FrameData}; +use cid::Cid; +use jacquard_repo::car::write_car; +use jacquard_repo::storage::BlockStore; +use std::fs; +use std::str::FromStr; +use tokio::fs::File; +use tokio::io::AsyncReadExt; +use uuid::Uuid; + +pub async fn format_event_for_sending( + state: &AppState, + event: SequencedEvent, +) -> Result, anyhow::Error> { + let block_cids_str = event.blocks_cids.clone().unwrap_or_default(); + let mut frame: CommitFrame = event.into(); + + let mut car_bytes = Vec::new(); + if !block_cids_str.is_empty() { + let temp_path = format!("/tmp/{}.car", Uuid::new_v4()); + let mut blocks = std::collections::BTreeMap::new(); + + for cid_str in block_cids_str { + let cid = Cid::from_str(&cid_str)?; + let data = state + .block_store + .get(&cid) + .await? + .ok_or_else(|| anyhow::anyhow!("Block not found: {}", cid))?; + blocks.insert(cid, data); + } + + let root = Cid::from_str(&frame.commit)?; + write_car(&temp_path, vec![root], blocks).await?; + + let mut file = File::open(&temp_path).await?; + file.read_to_end(&mut car_bytes).await?; + fs::remove_file(&temp_path)?; + } + frame.blocks = car_bytes; + + let frame = Frame { + op: 1, + data: FrameData::Commit(Box::new(frame)), + }; + + let mut bytes = Vec::new(); + serde_ipld_dagcbor::to_writer(&mut bytes, &frame)?; + Ok(bytes) +} diff --git a/tests/common/mod.rs b/tests/common/mod.rs index 7e95ecb..bc12156 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -19,6 +19,7 @@ use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; static SERVER_URL: OnceLock = OnceLock::new(); +static APP_PORT: OnceLock = OnceLock::new(); static DB_CONTAINER: OnceLock> = OnceLock::new(); static S3_CONTAINER: OnceLock> = OnceLock::new(); static MOCK_APPVIEW: OnceLock = OnceLock::new(); @@ -53,6 +54,11 @@ pub fn client() -> Client { Client::new() } +#[allow(dead_code)] +pub fn app_port() -> u16 { + *APP_PORT.get().expect("APP_PORT not initialized") +} + pub async fn base_url() -> &'static str { SERVER_URL.get_or_init(|| { let (tx, rx) = std::sync::mpsc::channel(); @@ -153,7 +159,7 @@ pub async fn base_url() -> &'static str { .await .expect("Failed to start Postgres"); let connection_string = format!( - "postgres://postgres:postgres@127.0.0.1:{}/postgres", + "postgres://postgres:postgres@127.0.0.1:{}", container .get_host_port_ipv4(5432) .await @@ -186,12 +192,16 @@ async fn spawn_app(database_url: String) -> String { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); + APP_PORT.set(addr.port()).ok(); unsafe { std::env::set_var("PDS_HOSTNAME", addr.to_string()); } let state = AppState::new(pool).await; + + bspds::sync::listener::start_sequencer_listener(state.clone()).await; + let app = bspds::app(state); tokio::spawn(async move { diff --git a/tests/firehose.rs b/tests/firehose.rs new file mode 100644 index 0000000..d05f2d7 --- /dev/null +++ b/tests/firehose.rs @@ -0,0 +1,81 @@ +mod common; +use common::*; + +use bspds::sync::frame::{Frame, FrameData}; +use cid::Cid; +use futures::{stream::StreamExt, SinkExt}; +use iroh_car::CarReader; +use reqwest::StatusCode; +use serde_json::{json, Value}; +use std::io::Cursor; +use std::str::FromStr; +use tokio_tungstenite::{connect_async, tungstenite}; + +#[tokio::test] +async fn test_firehose_subscription() { + let client = client(); + let (token, did) = create_account_and_login(&client).await; + + let url = format!( + "ws://127.0.0.1:{}/xrpc/com.atproto.sync.subscribeRepos", + app_port() + ); + let (mut ws_stream, _) = connect_async(&url).await.expect("Failed to connect"); + + let post_text = "Hello from the firehose test!"; + let post_payload = json!({ + "repo": did, + "collection": "app.bsky.feed.post", + "record": { + "$type": "app.bsky.feed.post", + "text": post_text, + "createdAt": chrono::Utc::now().to_rfc3339(), + } + }); + let res = client + .post(format!( + "{}/xrpc/com.atproto.repo.createRecord", + base_url().await + )) + .bearer_auth(token) + .json(&post_payload) + .send() + .await + .expect("Failed to create post"); + assert_eq!(res.status(), StatusCode::OK); + + let msg = ws_stream.next().await.unwrap().unwrap(); + + let frame: Frame = match msg { + tungstenite::Message::Binary(bin) => { + serde_ipld_dagcbor::from_slice(&bin).expect("Failed to deserialize frame") + } + _ => panic!("Expected binary message"), + }; + + let FrameData::Commit(commit) = frame.data; + assert_eq!(commit.repo, did); + assert_eq!(commit.ops.len(), 1); + assert!(!commit.blocks.is_empty()); + + let op = &commit.ops[0]; + let record_cid = Cid::from_str(&op.cid.clone().unwrap()).unwrap(); + + let mut car_reader = CarReader::new(Cursor::new(&commit.blocks)).await.unwrap(); + let mut record_block: Option> = None; + while let Ok(Some((cid, block))) = car_reader.next_block().await { + if cid == record_cid { + record_block = Some(block); + break; + } + } + let record_block = record_block.expect("Record block not found in CAR"); + + let record: Value = serde_ipld_dagcbor::from_slice(&record_block).unwrap(); + assert_eq!(record["text"], post_text); + + ws_stream + .send(tungstenite::Message::Close(None)) + .await + .ok(); +} \ No newline at end of file diff --git a/tests/relay_client.rs b/tests/relay_client.rs new file mode 100644 index 0000000..23dc655 --- /dev/null +++ b/tests/relay_client.rs @@ -0,0 +1,68 @@ +mod common; +use common::*; + +use axum::{extract::ws::Message, routing::get, Router}; +use bspds::{ + state::AppState, + sync::{firehose::SequencedEvent, relay_client::start_relay_clients}, +}; +use chrono::Utc; +use tokio::net::TcpListener; +use tokio::sync::mpsc; + +async fn mock_relay_server( + listener: TcpListener, + event_tx: mpsc::Sender>, + ready_tx: mpsc::Sender<()>, +) { + let handler = |ws: axum::extract::ws::WebSocketUpgrade| async { + ws.on_upgrade(move |mut socket| async move { + ready_tx.send(()).await.unwrap(); + if let Some(Ok(Message::Binary(bytes))) = socket.recv().await { + event_tx.send(bytes.to_vec()).await.unwrap(); + } + }) + }; + let app = Router::new().route("/", get(handler)); + + axum::serve(listener, app.into_make_service()) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_outbound_relay_client() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (event_tx, mut event_rx) = mpsc::channel(1); + let (ready_tx, ready_rx) = mpsc::channel(1); + tokio::spawn(mock_relay_server(listener, event_tx, ready_tx)); + let relay_url = format!("ws://{}", addr); + + let db_url = get_db_connection_string().await; + let pool = sqlx::postgres::PgPoolOptions::new() + .connect(&db_url) + .await + .unwrap(); + let state = AppState::new(pool).await; + + start_relay_clients(state.clone(), vec![relay_url], Some(ready_rx)).await; + + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + + let dummy_event = SequencedEvent { + seq: 1, + did: "did:plc:test".to_string(), + created_at: Utc::now(), + event_type: "commit".to_string(), + commit_cid: None, + prev_cid: None, + ops: None, + blobs: None, + blocks_cids: None, + }; + state.firehose_tx.send(dummy_event).unwrap(); + + let received_bytes = event_rx.recv().await.expect("Did not receive event"); + assert!(!received_bytes.is_empty()); +}