MST walk, and repo tweak (#984)

Adds a `Walk` method to `Tree` for visiting all key/value in the MST.

Also updates `GetRecordBytes` to return the CID, along with actual
bytes. It needs to be looked up from the MST anyways, and this results
in one less call/lookup in some cases. Code which doesn't need the CID
can just ignore that return type.

Author: bnewbold

atproto/repo/mst/node.go

@@ -291,6 +291,25 @@ func (n *Node) writeToMap(m map[string]cid.Cid) error {
 	return nil
 }
 
+func (n *Node) walk(f func(key []byte, val cid.Cid) error) error {
+	if n == nil {
+		return fmt.Errorf("nil tree pointer")
+	}
+	for _, e := range n.Entries {
+		if e.IsValue() {
+			if err := f(e.Key, *e.Value); err != nil {
+				return err
+			}
+		}
+		if e.Child != nil {
+			if err := e.Child.walk(f); err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
 // Reads the value (CID) corresponding to the key. If key is not in the tree, returns (nil, nil).
 //
 // n: Node at top of sub-tree to operate on. Must not be nil.

atproto/repo/mst/tree.go

@@ -79,6 +79,11 @@ func (t *Tree) Get(key []byte) (*cid.Cid, error) {
 	return t.Root.getCID(key, -1)
 }
 
+// Walks the Tree, invoking the callback function on each key/value pair.
+func (t *Tree) Walk(f func(key []byte, val cid.Cid) error) error {
+	return t.Root.walk(f)
+}
+
 // Creates a new Tree by loading key/value pairs from a map.
 func LoadTreeFromMap(m map[string]cid.Cid) (*Tree, error) {
 	if m == nil {

atproto/repo/repo.go

@@ -48,17 +48,17 @@ func (repo *Repo) GetRecordCID(ctx context.Context, collection syntax.NSID, rkey
 	return c, nil
 }
 
-func (repo *Repo) GetRecordBytes(ctx context.Context, collection syntax.NSID, rkey syntax.RecordKey) ([]byte, error) {
+func (repo *Repo) GetRecordBytes(ctx context.Context, collection syntax.NSID, rkey syntax.RecordKey) ([]byte, *cid.Cid, error) {
 	c, err := repo.GetRecordCID(ctx, collection, rkey)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	blk, err := repo.RecordStore.Get(ctx, *c)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	// TODO: not verifying CID
-	return blk.RawData(), nil
+	return blk.RawData(), c, nil
 }
 
 // Snapshots the current state of the repository, resulting in a new (unsigned) `Commit` struct.

atproto/repo/sync.go

@@ -61,17 +61,14 @@ func VerifyCommitMessage(ctx context.Context, msg *comatproto.SyncSubscribeRepos
 			if err != nil {
 				return nil, fmt.Errorf("invalid repo path in ops list: %w", err)
 			}
-			val, err := repo.GetRecordCID(ctx, nsid, rkey)
+			// don't use the returned bytes, but do actually read them out of store (not just CID)
+			_, val, err := repo.GetRecordBytes(ctx, nsid, rkey)
 			if err != nil {
 				return nil, err
 			}
 			if *c != *val {
 				return nil, fmt.Errorf("record op doesn't match MST tree value")
 			}
-			_, err = repo.GetRecordBytes(ctx, nsid, rkey)
-			if err != nil {
-				return nil, err
-			}
 		}
 	}