Something went wrong. Try again.
Chris’s PDS implementation
Something went wrong. Try again.
14 kB · 394 lines
Rust
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395//! Stage `ATrium` blocks, sign the final root, and atomically publish the index, head, and outbox.use super::{MAX_REPO_RETRIES, RepositoryService, car::map_block_error};use crate::error::ErrorContext;use crate::{ database::{PendingRepoCommit, RepoBlockStore, RepoHead}, domain::{ account::Account, event::RepositoryEvent, record::{Record, RecordMutation}, }, error::{Error, Result},};use atrium_api::com::atproto::repo::defs::CommitMetaData;use atrium_api::types::string::{Cid as AtriumCid, Did, Tid};use atrium_crypto::verify::verify_signature;use atrium_repo::{ Repository, blockstore::{AsyncBlockStoreWrite as _, DAG_CBOR as ATRIUM_DAG_CBOR, SHA2_256}, repo::Error as AtriumRepoError,};use cid::Cid;use cpds_protocol::atproto_record;use serde::Serialize;
#[derive(Serialize)]struct UnsignedCommit<'a> { did: Did, version: i64, data: Cid, rev: &'a Tid, prev: Option<Cid>,}
#[derive(Serialize)]struct SignedCommit<'a> { #[serde(flatten)] content: UnsignedCommit<'a>, #[serde(rename = "sig", with = "serde_bytes")] signature: Vec<u8>,}
fn record_path(record: &Record) -> String { format!("{}/{}", record.collection, record.rkey)}
async fn stage_record_blocks( blocks: &mut RepoBlockStore, mutations: &[RecordMutation],) -> Result<()> { for mutation in mutations { let record = match mutation { RecordMutation::Create(record) | RecordMutation::Update { replacement: record, .. } => record, RecordMutation::Delete(_) => continue, }; let data = atproto_record(&record.value)?; let bytes = serde_ipld_dagcbor::to_vec(&data) .context(Error::Internal, "validated record cannot be encoded")?; let cid = blocks .write_block(ATRIUM_DAG_CBOR, SHA2_256, &bytes) .await .map_err(map_block_error)?; if cid != record.cid { return Err(Error::Internal( "ATrium record CID differs from validated record CID".to_owned(), )); } } Ok(())}
async fn apply_repo_mutations( repo: &mut Repository<&mut RepoBlockStore>, mutations: &[RecordMutation],) -> Result<Cid> { let mut tree = repo.tree(); for mutation in mutations { match mutation { RecordMutation::Create(record) => tree.add(&record_path(record), record.cid).await, RecordMutation::Update { replacement, .. } => { tree.update(&record_path(replacement), replacement.cid) .await } RecordMutation::Delete(record) => tree.delete(&record_path(record)).await, } .map_err(|error| map_repo_error(error.into()))?; } Ok(tree.root())}
async fn event_block_cids(blocks: &mut RepoBlockStore, root: Cid) -> Result<Vec<Cid>> { let reachable = { let mut repo = Repository::open(&mut *blocks, root) .await .map_err(map_repo_error)?; repo.export() .await .map_err(map_repo_error)? .collect::<Vec<_>>() }; Ok(blocks .reachable_writes(reachable)? .map(|(cid, _)| *cid) .collect())}
pub(super) fn invalid_swap(message: &str) -> Error { Error::xrpc("InvalidSwap", message)}
pub(super) fn map_repo_error(error: AtriumRepoError) -> Error { match error { AtriumRepoError::BlockStore(error) => map_block_error(error), AtriumRepoError::Mst(atrium_repo::mst::Error::KeyAlreadyExists) => { Error::Conflict("record already exists".to_owned()) } AtriumRepoError::Mst(atrium_repo::mst::Error::KeyNotFound) => Error::NotFound, error => Error::Internal(format!("ATrium repository operation failed: {error}")), }}
impl RepositoryService { fn record_event(mutation: &RecordMutation) -> RepositoryEvent { match mutation { RecordMutation::Create(record) => RepositoryEvent::RecordCreated { did: record.did.clone(), collection: record.collection.clone(), rkey: record.rkey.clone(), cid: record.cid, }, RecordMutation::Update { current, replacement, } => RepositoryEvent::RecordUpdated { did: replacement.did.clone(), collection: replacement.collection.clone(), rkey: replacement.rkey.clone(), cid: replacement.cid, previous_cid: current.cid, }, RecordMutation::Delete(record) => RepositoryEvent::RecordDeleted { did: record.did.clone(), collection: record.collection.clone(), rkey: record.rkey.clone(), previous_cid: record.cid, }, } }
pub(super) async fn commit_record_mutations( &self, actor: &Account, expected_commit: Option<&Cid>, mutations: &[RecordMutation], ) -> Result<CommitMetaData> { for _ in 0..MAX_REPO_RETRIES { let previous = self.database.repo_head(&actor.did)?; if expected_commit .is_some_and(|expected| previous.as_ref().is_none_or(|head| expected != &head.root)) { return Err(invalid_swap("repository commit does not match")); } let rev = previous.as_ref().map_or_else( || Ok(self.tids.next()), |head| { self.tids.next_after(&head.rev).ok_or_else(|| { Error::Internal("stored repository revision is invalid".to_owned()) }) }, )?; let tid = Tid::new(rev.clone()) .context(Error::Internal, "generated repository revision is invalid")?; let mut blocks = RepoBlockStore::new(self.database.clone()); stage_record_blocks(&mut blocks, mutations).await?; let mut repo = self .open_or_create_repo(&mut blocks, actor, previous.as_ref()) .await?; let previous_data = previous.as_ref().map(|_| repo.commit().data()); let data = apply_repo_mutations(&mut repo, mutations).await?; drop(repo); let root = self .write_signed_commit(&mut blocks, actor, data, &tid) .await?; let event_blocks = event_block_cids(&mut blocks, root).await?; let staged = blocks.writes(); let commit = PendingRepoCommit { did: actor.did.clone(), expected_head: previous.as_ref().map(|head| head.root), previous_rev: previous.as_ref().map(|head| head.rev.clone()), previous_data, head: RepoHead { root, rev: rev.clone(), }, blocks: staged, event_blocks, mutations: mutations.to_vec(), }; match self .database .commit_repo(&commit, self.repo_event_retention) { Ok(()) => { self.publish_events(mutations.iter().map(Self::record_event)) .await; return Ok(CommitMetaData { cid: AtriumCid::new(root), rev: tid, }); } Err(Error::Conflict(message)) if message == "repository head changed" => {} Err(error) => return Err(error), } } Err(Error::Conflict( "repository changed too frequently; retry the request".to_owned(), )) }
async fn write_signed_commit( &self, blocks: &mut RepoBlockStore, actor: &Account, data: Cid, rev: &Tid, ) -> Result<Cid> { let did = Did::new(actor.did.clone()) .context(Error::Internal, "hosted account DID is invalid")?; let commit = UnsignedCommit { did, version: 3, data, rev, prev: None, }; let signature = self.sign( &serde_ipld_dagcbor::to_vec(&commit) .context(Error::Internal, "repository commit cannot be encoded")?, )?; let signed = SignedCommit { content: commit, signature, }; let bytes = serde_ipld_dagcbor::to_vec(&signed).context( Error::Internal, "signed repository commit cannot be encoded", )?; blocks .write_block(ATRIUM_DAG_CBOR, SHA2_256, &bytes) .await .map_err(map_block_error) }
async fn open_or_create_repo<'a>( &self, blocks: &'a mut RepoBlockStore, actor: &Account, head: Option<&RepoHead>, ) -> Result<Repository<&'a mut RepoBlockStore>> { if let Some(head) = head { return self.open_repo(blocks, head.root, &head.rev).await; }
let did = Did::new(actor.did.clone()) .context(Error::Internal, "hosted account DID is invalid")?; let builder = Repository::create(&mut *blocks, did) .await .map_err(map_repo_error)?; let signature = self.sign(&builder.bytes())?; let mut repo = builder.finalize(signature).await.map_err(map_repo_error)?;
for record in self.database.all_records(&actor.did)? { let path = format!("{}/{}", record.collection, record.rkey); let data = atproto_record(&record.value).context( Error::Internal, format_args!("stored record {record} is invalid"), )?; let (builder, cid) = repo.add_raw(&path, &data).await.map_err(map_repo_error)?; if cid != record.cid { return Err(Error::Internal(format!( "stored record {record} has noncanonical CID" ))); } let signature = self.sign(&builder.bytes())?; builder.finalize(signature).await.map_err(map_repo_error)?; } Ok(repo) }
pub(super) async fn open_repo<'a>( &self, blocks: &'a mut RepoBlockStore, root: Cid, expected_rev: &str, ) -> Result<Repository<&'a mut RepoBlockStore>> { let repo = Repository::open(&mut *blocks, root) .await .map_err(map_repo_error)?; let commit = repo.commit(); if commit.rev().to_string() != expected_rev { return Err(Error::Internal( "stored repository head revision does not match its commit".to_owned(), )); } verify_signature(&self.signing_did, &commit.bytes(), commit.sig()) .context(Error::Internal, "stored repository signature is invalid")?; Ok(repo) }
fn sign(&self, bytes: &[u8]) -> Result<Vec<u8>> { self.signer .sign(bytes) .context(Error::Internal, "repository commit signing failed") }}
#[cfg(test)]mod tests { use super::*;
#[derive(Serialize)] struct PreviousUnsignedCommit { did: Did, version: i64, data: Cid, rev: Tid, prev: Option<Cid>, }
#[derive(Serialize)] struct PreviousSignedCommit { #[serde(flatten)] content: PreviousUnsignedCommit, #[serde(rename = "sig", with = "serde_bytes")] signature: Vec<u8>, }
#[test] fn borrowed_revision_keeps_unsigned_and_signed_commit_bytes() { let data = cpds_protocol::dag_cbor_cid(b"data"); let previous = cpds_protocol::dag_cbor_cid(b"previous"); for rev in ["2222222222222", "3jqfcqzm3fo2j", "jzzzzzzzzzzzz"] { let rev = Tid::new(rev.to_owned()).unwrap(); for prev in [None, Some(previous)] { for signature in [Vec::new(), vec![0, 255], (0..=255).collect()] { let owned = PreviousUnsignedCommit { did: Did::new("did:web:alice.test".to_owned()).unwrap(), version: 3, data, rev: rev.clone(), prev, }; let borrowed = UnsignedCommit { did: owned.did.clone(), version: owned.version, data, rev: &rev, prev, }; assert_eq!( serde_ipld_dagcbor::to_vec(&borrowed).unwrap(), serde_ipld_dagcbor::to_vec(&owned).unwrap() ); assert_eq!( serde_json::to_vec(&borrowed).unwrap(), serde_json::to_vec(&owned).unwrap() ); let owned = PreviousSignedCommit { content: owned, signature: signature.clone(), }; let borrowed = SignedCommit { content: borrowed, signature, }; let expected = serde_ipld_dagcbor::to_vec(&owned).unwrap(); let actual = serde_ipld_dagcbor::to_vec(&borrowed).unwrap(); assert_eq!(actual, expected); assert_eq!( cpds_protocol::dag_cbor_cid(&actual), cpds_protocol::dag_cbor_cid(&expected) ); assert_eq!( serde_json::to_vec(&borrowed).unwrap(), serde_json::to_vec(&owned).unwrap() ); } } } }}