Something went wrong. Try again.
An atproto PDS written in Go
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731package server
import ( "bytes" "context" "encoding/json" "fmt" "io" "sync" "time"
"github.com/Azure/go-autorest/autorest/to" "github.com/bluesky-social/indigo/api/atproto" "github.com/bluesky-social/indigo/atproto/atcrypto" "github.com/bluesky-social/indigo/atproto/atdata" atp "github.com/bluesky-social/indigo/atproto/repo" "github.com/bluesky-social/indigo/atproto/repo/mst" "github.com/bluesky-social/indigo/atproto/syntax" "github.com/bluesky-social/indigo/carstore" "github.com/bluesky-social/indigo/events" lexutil "github.com/bluesky-social/indigo/lex/util" "github.com/haileyok/cocoon/internal/db" "github.com/haileyok/cocoon/metrics" "github.com/haileyok/cocoon/models" "github.com/haileyok/cocoon/recording_blockstore" blocks "github.com/ipfs/go-block-format" "github.com/ipfs/go-cid" blockstore "github.com/ipfs/go-ipfs-blockstore" cbor "github.com/ipfs/go-ipld-cbor" "github.com/ipld/go-car" "github.com/multiformats/go-multihash" "gorm.io/gorm/clause")
type cachedRepo struct { mu sync.Mutex repo *atp.Repo root cid.Cid}
type RepoMan struct { db *db.DB s *Server clock *syntax.TIDClock
cacheMu sync.Mutex cache map[string]*cachedRepo}
func NewRepoMan(s *Server) *RepoMan { clock := syntax.NewTIDClock(0)
return &RepoMan{ s: s, db: s.db, clock: clock, cache: make(map[string]*cachedRepo), }}
func (rm *RepoMan) withRepo(ctx context.Context, did string, rootCid cid.Cid, fn func(r *atp.Repo) (newRoot cid.Cid, err error)) error { rm.cacheMu.Lock() cr, ok := rm.cache[did] if !ok { cr = &cachedRepo{} rm.cache[did] = cr } rm.cacheMu.Unlock()
cr.mu.Lock() defer cr.mu.Unlock()
if cr.repo == nil || cr.root != rootCid { bs := rm.s.getBlockstore(did) r, err := openRepo(ctx, bs, rootCid, did) if err != nil { return err } cr.repo = r cr.root = rootCid }
newRoot, err := fn(cr.repo) if err != nil { // invalidate on error since the tree may be partially mutated cr.repo = nil cr.root = cid.Undef return err }
cr.root = newRoot return nil}
type OpType string
var ( OpTypeCreate = OpType("com.atproto.repo.applyWrites#create") OpTypeUpdate = OpType("com.atproto.repo.applyWrites#update") OpTypeDelete = OpType("com.atproto.repo.applyWrites#delete"))
func (ot OpType) String() string { return string(ot)}
type Op struct { Type OpType `json:"$type"` Collection string `json:"collection"` Rkey *string `json:"rkey,omitempty"` Validate *bool `json:"validate,omitempty"` SwapRecord *string `json:"swapRecord,omitempty"` Record *MarshalableMap `json:"record,omitempty"`}
type MarshalableMap map[string]any
type FirehoseOp struct { Cid cid.Cid Path string Action string}
func (mm *MarshalableMap) MarshalCBOR(w io.Writer) error { data, err := atdata.MarshalCBOR(*mm) if err != nil { return err }
w.Write(data)
return nil}
type ApplyWriteResult struct { Type *string `json:"$type,omitempty"` Uri *string `json:"uri,omitempty"` Cid *string `json:"cid,omitempty"` Commit *RepoCommit `json:"commit,omitempty"` ValidationStatus *string `json:"validationStatus,omitempty"`}
type RepoCommit struct { Cid string `json:"cid"` Rev string `json:"rev"`}
func openRepo(ctx context.Context, bs blockstore.Blockstore, rootCid cid.Cid, did string) (*atp.Repo, error) { commitBlock, err := bs.Get(ctx, rootCid) if err != nil { return nil, fmt.Errorf("reading commit block: %w", err) }
var commit atp.Commit if err := commit.UnmarshalCBOR(bytes.NewReader(commitBlock.RawData())); err != nil { return nil, fmt.Errorf("parsing commit block: %w", err) }
tree, err := mst.LoadTreeFromStore(ctx, bs, commit.Data) if err != nil { return nil, fmt.Errorf("loading MST: %w", err) }
clk := syntax.ClockFromTID(syntax.TID(commit.Rev)) return &atp.Repo{ DID: syntax.DID(did), Clock: &clk, MST: *tree, RecordStore: bs, }, nil}
func commitRepo(ctx context.Context, bs blockstore.Blockstore, r *atp.Repo, signingKey []byte) (cid.Cid, string, error) { if _, err := r.MST.WriteDiffBlocks(ctx, bs); err != nil { return cid.Undef, "", fmt.Errorf("writing MST blocks: %w", err) }
commit, err := r.Commit() if err != nil { return cid.Undef, "", fmt.Errorf("creating commit: %w", err) }
privkey, err := atcrypto.ParsePrivateBytesK256(signingKey) if err != nil { return cid.Undef, "", fmt.Errorf("parsing signing key: %w", err) } if err := commit.Sign(privkey); err != nil { return cid.Undef, "", fmt.Errorf("signing commit: %w", err) }
buf := new(bytes.Buffer) if err := commit.MarshalCBOR(buf); err != nil { return cid.Undef, "", fmt.Errorf("marshaling commit: %w", err) }
pref := cid.NewPrefixV1(cid.DagCBOR, multihash.SHA2_256) commitCid, err := pref.Sum(buf.Bytes()) if err != nil { return cid.Undef, "", fmt.Errorf("computing commit CID: %w", err) }
blk, err := blocks.NewBlockWithCid(buf.Bytes(), commitCid) if err != nil { return cid.Undef, "", fmt.Errorf("creating commit block: %w", err) } if err := bs.Put(ctx, blk); err != nil { return cid.Undef, "", fmt.Errorf("writing commit block: %w", err) }
return commitCid, commit.Rev, nil}
func putRecordBlock(ctx context.Context, bs blockstore.Blockstore, rec *MarshalableMap) (cid.Cid, error) { buf := new(bytes.Buffer) if err := rec.MarshalCBOR(buf); err != nil { return cid.Undef, err }
pref := cid.NewPrefixV1(cid.DagCBOR, multihash.SHA2_256) c, err := pref.Sum(buf.Bytes()) if err != nil { return cid.Undef, err }
blk, err := blocks.NewBlockWithCid(buf.Bytes(), c) if err != nil { return cid.Undef, err } if err := bs.Put(ctx, blk); err != nil { return cid.Undef, err }
return c, nil}
// TODO make use of swap commitfunc (rm *RepoMan) applyWrites(ctx context.Context, urepo models.Repo, writes []Op, swapCommit *string) ([]ApplyWriteResult, error) { rootcid, err := cid.Cast(urepo.Root) if err != nil { return nil, err }
dbs := rm.s.getBlockstore(urepo.Did) bs := recording_blockstore.New(dbs)
var results []ApplyWriteResult var ops []*atp.Operation var entries []models.Record var newroot cid.Cid var rev string
if err := rm.withRepo(ctx, urepo.Did, rootcid, func(r *atp.Repo) (cid.Cid, error) { entries = make([]models.Record, 0, len(writes)) for i, op := range writes { // updates or deletes must supply an rkey if op.Type != OpTypeCreate && op.Rkey == nil { return cid.Undef, fmt.Errorf("invalid rkey") } else if op.Type == OpTypeCreate && op.Rkey != nil { // we should convert this op to an update if the rkey already exists path := fmt.Sprintf("%s/%s", op.Collection, *op.Rkey) existing, _ := r.MST.Get([]byte(path)) if existing != nil { op.Type = OpTypeUpdate } } else if op.Rkey == nil { // creates that don't supply an rkey will have one generated for them op.Rkey = to.StringPtr(rm.clock.Next().String()) writes[i].Rkey = op.Rkey }
path := fmt.Sprintf("%s/%s", op.Collection, *op.Rkey)
// validate the record key is actually valid _, err := syntax.ParseRecordKey(*op.Rkey) if err != nil { return cid.Undef, err }
switch op.Type { case OpTypeCreate: // HACK: this fixes some type conversions, mainly around integers b, err := json.Marshal(*op.Record) if err != nil { return cid.Undef, err } out, err := atdata.UnmarshalJSON(b) if err != nil { return cid.Undef, err } mm := MarshalableMap(out)
// HACK: if a record doesn't contain a $type, we can manually set it here based on the op's collection if mm["$type"] == "" { mm["$type"] = op.Collection }
nc, err := putRecordBlock(ctx, bs, &mm) if err != nil { return cid.Undef, err }
atpOp, err := atp.ApplyOp(&r.MST, path, &nc) if err != nil { return cid.Undef, err } ops = append(ops, atpOp)
d, err := atdata.MarshalCBOR(mm) if err != nil { return cid.Undef, err }
entries = append(entries, models.Record{ Did: urepo.Did, CreatedAt: rm.clock.Next().String(), Nsid: op.Collection, Rkey: *op.Rkey, Cid: nc.String(), Value: d, })
results = append(results, ApplyWriteResult{ Type: to.StringPtr(OpTypeCreate.String()), Uri: to.StringPtr("at://" + urepo.Did + "/" + op.Collection + "/" + *op.Rkey), Cid: to.StringPtr(nc.String()), ValidationStatus: to.StringPtr("valid"), // TODO: obviously this might not be true atm lol }) case OpTypeDelete: // try to find the old record in the database var old models.Record if err := rm.db.Raw(ctx, "SELECT value FROM records WHERE did = ? AND nsid = ? AND rkey = ?", nil, urepo.Did, op.Collection, op.Rkey).Scan(&old).Error; err != nil { return cid.Undef, err }
// TODO: this is really confusing, and looking at it i have no idea why i did this. below when we are doing deletes, we // check if `cid` here is nil to indicate if we should delete. that really doesn't make much sense and its super illogical // when reading this code. i dont feel like fixing right now though so entries = append(entries, models.Record{ Did: urepo.Did, Nsid: op.Collection, Rkey: *op.Rkey, Value: old.Value, })
atpOp, err := atp.ApplyOp(&r.MST, path, nil) if err != nil { return cid.Undef, err } ops = append(ops, atpOp)
results = append(results, ApplyWriteResult{ Type: to.StringPtr(OpTypeDelete.String()), }) case OpTypeUpdate: // HACK: same hack as above for type fixes b, err := json.Marshal(*op.Record) if err != nil { return cid.Undef, err } out, err := atdata.UnmarshalJSON(b) if err != nil { return cid.Undef, err } mm := MarshalableMap(out)
nc, err := putRecordBlock(ctx, bs, &mm) if err != nil { return cid.Undef, err }
atpOp, err := atp.ApplyOp(&r.MST, path, &nc) if err != nil { return cid.Undef, err } ops = append(ops, atpOp)
d, err := atdata.MarshalCBOR(mm) if err != nil { return cid.Undef, err }
entries = append(entries, models.Record{ Did: urepo.Did, CreatedAt: rm.clock.Next().String(), Nsid: op.Collection, Rkey: *op.Rkey, Cid: nc.String(), Value: d, })
results = append(results, ApplyWriteResult{ Type: to.StringPtr(OpTypeUpdate.String()), Uri: to.StringPtr("at://" + urepo.Did + "/" + op.Collection + "/" + *op.Rkey), Cid: to.StringPtr(nc.String()), ValidationStatus: to.StringPtr("valid"), // TODO: obviously this might not be true atm lol }) } }
// commit and get the new root var commitErr error newroot, rev, commitErr = commitRepo(ctx, bs, r, urepo.SigningKey) if commitErr != nil { return cid.Undef, commitErr }
return newroot, nil }); err != nil { return nil, err }
for _, result := range results { if result.Type != nil { metrics.RepoOperations.WithLabelValues(*result.Type).Inc() } }
// create a buffer for dumping our new cbor into buf := new(bytes.Buffer)
// first write the car header to the buffer hb, err := cbor.DumpObject(&car.CarHeader{ Roots: []cid.Cid{newroot}, Version: 1, }) if _, err := carstore.LdWrite(buf, hb); err != nil { return nil, err }
// create the repo ops for the firehose from the tracked operations repoOps := make([]*atproto.SyncSubscribeRepos_RepoOp, 0, len(ops)) for _, op := range ops { if op.IsCreate() || op.IsUpdate() { kind := "create" if op.IsUpdate() { kind = "update" }
ll := lexutil.LexLink(*op.Value) repoOps = append(repoOps, &atproto.SyncSubscribeRepos_RepoOp{ Action: kind, Path: op.Path, Cid: &ll, })
blk, err := dbs.Get(ctx, *op.Value) if err != nil { return nil, err } if _, err := carstore.LdWrite(buf, blk.Cid().Bytes(), blk.RawData()); err != nil { return nil, err } } else if op.IsDelete() { ll := lexutil.LexLink(*op.Prev) repoOps = append(repoOps, &atproto.SyncSubscribeRepos_RepoOp{ Action: "delete", Path: op.Path, Cid: nil, Prev: &ll, })
blk, err := dbs.Get(ctx, *op.Prev) if err != nil { return nil, err } if _, err := carstore.LdWrite(buf, blk.Cid().Bytes(), blk.RawData()); err != nil { return nil, err } } }
// write the writelog to the buffer for _, blk := range bs.GetWriteLog() { if _, err := carstore.LdWrite(buf, blk.Cid().Bytes(), blk.RawData()); err != nil { return nil, err } }
// blob blob blob blob blob :3 var blobs []lexutil.LexLink for _, entry := range entries { var cids []cid.Cid // whenever there is cid present, we know it's a create (dumb) if entry.Cid != "" { if err := rm.s.db.Create(ctx, &entry, []clause.Expression{clause.OnConflict{ Columns: []clause.Column{{Name: "did"}, {Name: "nsid"}, {Name: "rkey"}}, UpdateAll: true, }}).Error; err != nil { return nil, err }
// increment the given blob refs, yay cids, err = rm.incrementBlobRefs(ctx, urepo, entry.Value) if err != nil { return nil, err } } else { // as i noted above this is dumb. but we delete whenever the cid is nil. it works solely becaue the pkey // is did + collection + rkey. i still really want to separate that out, or use a different type to make // this less confusing/easy to read. alas, its 2 am and yea no if err := rm.s.db.Delete(ctx, &entry, nil).Error; err != nil { return nil, err }
// TODO: cids, err = rm.decrementBlobRefs(ctx, urepo, entry.Value) if err != nil { return nil, err } }
// add all the relevant blobs to the blobs list of blobs. blob ^.^ for _, c := range cids { blobs = append(blobs, lexutil.LexLink(c)) } }
// NOTE: using the request ctx seems a bit suss here, so using a background context. i'm not sure if this // runs sync or not rm.s.evtman.AddEvent(context.Background(), &events.XRPCStreamEvent{ RepoCommit: &atproto.SyncSubscribeRepos_Commit{ Repo: urepo.Did, Blocks: buf.Bytes(), Blobs: blobs, Rev: rev, Since: &urepo.Rev, Commit: lexutil.LexLink(newroot), Time: time.Now().Format(time.RFC3339Nano), Ops: repoOps, TooBig: false, }, })
if err := rm.s.UpdateRepo(ctx, urepo.Did, newroot, rev); err != nil { return nil, err }
for i := range results { results[i].Type = to.StringPtr(*results[i].Type + "Result") results[i].Commit = &RepoCommit{ Cid: newroot.String(), Rev: rev, } }
return results, nil}
func (rm *RepoMan) getRecordProof(ctx context.Context, urepo models.Repo, collection, rkey string) (cid.Cid, []blocks.Block, error) { commitCid, err := cid.Cast(urepo.Root) if err != nil { return cid.Undef, nil, err }
dbs := rm.s.getBlockstore(urepo.Did)
var proofBlocks []blocks.Block var recordCid *cid.Cid
if err := rm.withRepo(ctx, urepo.Did, commitCid, func(r *atp.Repo) (cid.Cid, error) { path := collection + "/" + rkey
// walk the cached in-memory tree to find the record and collect MST node CIDs on the path nodeCIDs := collectPathNodeCIDs(r.MST.Root, []byte(path))
rc, getErr := r.MST.Get([]byte(path)) if getErr != nil { return cid.Undef, getErr } if rc == nil { return cid.Undef, fmt.Errorf("record not found: %s", path) } recordCid = rc
// read the commit block commitBlk, err := dbs.Get(ctx, commitCid) if err != nil { return cid.Undef, fmt.Errorf("reading commit block for proof: %w", err) } proofBlocks = append(proofBlocks, commitBlk)
// read the MST nodes on the path for _, nc := range nodeCIDs { blk, err := dbs.Get(ctx, nc) if err != nil { return cid.Undef, fmt.Errorf("reading MST node for proof: %w", err) } proofBlocks = append(proofBlocks, blk) }
// read the record block recordBlk, err := dbs.Get(ctx, *recordCid) if err != nil { return cid.Undef, fmt.Errorf("reading record block for proof: %w", err) } proofBlocks = append(proofBlocks, recordBlk)
// read-only, return same root return commitCid, nil }); err != nil { return cid.Undef, nil, err }
return commitCid, proofBlocks, nil}
func collectPathNodeCIDs(n *mst.Node, key []byte) []cid.Cid { if n == nil { return nil }
var cids []cid.Cid if n.CID != nil { cids = append(cids, *n.CID) }
height := mst.HeightForKey(key) if height >= n.Height { // key is at or above this level, no need to descend return cids }
// find the child node that covers this key childIdx := -1 for i, e := range n.Entries { if e.IsChild() { childIdx = i continue } if e.IsValue() { if bytes.Compare(key, e.Key) <= 0 { break } childIdx = -1 } }
if childIdx >= 0 && n.Entries[childIdx].Child != nil { cids = append(cids, collectPathNodeCIDs(n.Entries[childIdx].Child, key)...) }
return cids}
func (rm *RepoMan) incrementBlobRefs(ctx context.Context, urepo models.Repo, cbor []byte) ([]cid.Cid, error) { cids, err := getBlobCidsFromCbor(cbor) if err != nil { return nil, err }
for _, c := range cids { if err := rm.db.Exec(ctx, "UPDATE blobs SET ref_count = ref_count + 1 WHERE did = ? AND cid = ?", nil, urepo.Did, c.Bytes()).Error; err != nil { return nil, err } }
return cids, nil}
func (rm *RepoMan) decrementBlobRefs(ctx context.Context, urepo models.Repo, cbor []byte) ([]cid.Cid, error) { cids, err := getBlobCidsFromCbor(cbor) if err != nil { return nil, err }
for _, c := range cids { var res struct { ID uint Count int } if err := rm.db.Raw(ctx, "UPDATE blobs SET ref_count = ref_count - 1 WHERE did = ? AND cid = ? RETURNING id, ref_count", nil, urepo.Did, c.Bytes()).Scan(&res).Error; err != nil { return nil, err }
// TODO: this does _not_ handle deletions of blobs that are on s3 storage!!!! we need to get the blob, see what // storage it is in, and clean up s3!!!! if res.Count == 0 { if err := rm.db.Exec(ctx, "DELETE FROM blobs WHERE id = ?", nil, res.ID).Error; err != nil { return nil, err } if err := rm.db.Exec(ctx, "DELETE FROM blob_parts WHERE blob_id = ?", nil, res.ID).Error; err != nil { return nil, err } } }
return cids, nil}
// to be honest, we could just store both the cbor and non-cbor in []entries above to avoid an additional// unmarshal here. this will work for now thoughfunc getBlobCidsFromCbor(cbor []byte) ([]cid.Cid, error) { var cids []cid.Cid
decoded, err := atdata.UnmarshalCBOR(cbor) if err != nil { return nil, fmt.Errorf("error unmarshaling cbor: %w", err) }
var deepiter func(any) error deepiter = func(item any) error { switch val := item.(type) { case map[string]any: if val["$type"] == "blob" { if ref, ok := val["ref"].(string); ok { c, err := cid.Parse(ref) if err != nil { return err } cids = append(cids, c) } for _, v := range val { return deepiter(v) } } case []any: for _, v := range val { deepiter(v) } }
return nil }
if err := deepiter(decoded); err != nil { return nil, err }
return cids, nil}