@atproto/repo 0.0.1

package/src/util.ts

@@ -0,0 +1,88 @@
+import { CID } from 'multiformats/cid'
+import * as auth from '@atproto/auth'
+import Repo from './repo'
+import { DataDiff, MST } from './mst'
+import { IpldStore } from './blockstore'
+import { DataStore, RecordWriteOp, def } from './types'
+
+export const ucanForOperation = async (
+  prevData: DataStore,
+  newData: DataStore,
+  rootDid: string,
+  authStore: auth.AuthStore,
+): Promise<string> => {
+  const diff = await prevData.diff(newData)
+  const neededCaps = diff.neededCapabilities(rootDid)
+  const ucanForOp = await authStore.createUcanForCaps(rootDid, neededCaps, 30)
+  return auth.encodeUcan(ucanForOp)
+}
+
+export const getCommitPath = async (
+  blockstore: IpldStore,
+  earliest: CID | null,
+  latest: CID,
+): Promise<CID[] | null> => {
+  let curr: CID | null = latest
+  const path: CID[] = []
+  while (curr !== null) {
+    path.push(curr)
+    const commit = await blockstore.get(curr, def.commit)
+    if (earliest && curr.equals(earliest)) {
+      return path.reverse()
+    }
+    const root = await blockstore.get(commit.root, def.repoRoot)
+    if (!earliest && root.prev === null) {
+      return path.reverse()
+    }
+    curr = root.prev
+  }
+  return null
+}
+
+export const getWriteOpLog = async (
+  blockstore: IpldStore,
+  earliest: CID | null,
+  latest: CID,
+): Promise<RecordWriteOp[][]> => {
+  const commits = await getCommitPath(blockstore, earliest, latest)
+  if (!commits) throw new Error('Could not find shared history')
+  const heads = await Promise.all(commits.map((c) => Repo.load(blockstore, c)))
+  // Turn commit path into list of diffs
+  let prev: DataStore = await MST.create(blockstore) // Empty
+  const msts = heads.map((h) => h.data)
+  const diffs: DataDiff[] = []
+  for (const mst of msts) {
+    diffs.push(await prev.diff(mst))
+    prev = mst
+  }
+  // Map MST diffs to write ops
+  return Promise.all(diffs.map((diff) => diffToWriteOps(blockstore, diff)))
+}
+
+export const diffToWriteOps = (
+  blockstore: IpldStore,
+  diff: DataDiff,
+): Promise<RecordWriteOp[]> => {
+  return Promise.all([
+    ...diff.addList().map(async (add) => {
+      const { collection, rkey } = parseRecordKey(add.key)
+      const value = await blockstore.getUnchecked(add.cid)
+      return { action: 'create' as const, collection, rkey, value }
+    }),
+    ...diff.updateList().map(async (upd) => {
+      const { collection, rkey } = parseRecordKey(upd.key)
+      const value = await blockstore.getUnchecked(upd.cid)
+      return { action: 'update' as const, collection, rkey, value }
+    }),
+    ...diff.deleteList().map((del) => {
+      const { collection, rkey } = parseRecordKey(del.key)
+      return { action: 'delete' as const, collection, rkey }
+    }),
+  ])
+}
+
+export const parseRecordKey = (key: string) => {
+  const parts = key.split('/')
+  if (parts.length !== 2) throw new Error(`Invalid record key: ${key}`)
+  return { collection: parts[0], rkey: parts[1] }
+}

package/src/verify.ts

@@ -0,0 +1,62 @@
+import { CID } from 'multiformats/cid'
+import * as auth from '@atproto/auth'
+import { IpldStore } from './blockstore'
+import Repo from './repo'
+import { DataDiff } from './mst'
+import * as util from './util'
+import { def } from './types'
+
+export const verifyUpdates = async (
+  blockstore: IpldStore,
+  earliest: CID | null,
+  latest: CID,
+  verifier: auth.Verifier,
+): Promise<DataDiff> => {
+  const commitPath = await util.getCommitPath(blockstore, earliest, latest)
+  if (commitPath === null) {
+    throw new Error('Could not find shared history')
+  }
+  const fullDiff = new DataDiff()
+  if (commitPath.length === 0) return fullDiff
+  let prevRepo = await Repo.load(blockstore, commitPath[0])
+  for (const commit of commitPath.slice(1)) {
+    const nextRepo = await Repo.load(blockstore, commit)
+    const diff = await prevRepo.data.diff(nextRepo.data)
+
+    if (!nextRepo.root.meta.equals(prevRepo.root.meta)) {
+      throw new Error('Not supported: repo metadata updated')
+    }
+
+    let didForSignature: string
+    if (nextRepo.root.auth_token) {
+      // verify auth token covers all necessary writes
+      const encodedToken = await blockstore.get(
+        nextRepo.root.auth_token,
+        def.string,
+      )
+      const token = await verifier.validateUcan(encodedToken)
+      const neededCaps = diff.neededCapabilities(prevRepo.did)
+      for (const cap of neededCaps) {
+        await verifier.verifyAtpUcan(token, prevRepo.did, cap)
+      }
+      didForSignature = token.payload.iss
+    } else {
+      didForSignature = prevRepo.did
+    }
+
+    // verify signature matches repo root + auth token
+    // const commit = await toRepo.getCommit()
+    const validSig = await verifier.verifySignature(
+      didForSignature,
+      nextRepo.commit.root.bytes,
+      nextRepo.commit.sig,
+    )
+    if (!validSig) {
+      throw new Error(`Invalid signature on commit: ${nextRepo.cid.toString()}`)
+    }
+
+    fullDiff.addDiff(diff)
+    prevRepo = nextRepo
+  }
+  return fullDiff
+}

package/tests/_util.ts

@@ -0,0 +1,254 @@
+import { CID } from 'multiformats'
+import { TID } from '@atproto/common'
+import * as auth from '@atproto/auth'
+import IpldStore from '../src/blockstore/ipld-store'
+import { Repo } from '../src/repo'
+import { MemoryBlockstore } from '../src/blockstore'
+import { DataDiff, MST } from '../src/mst'
+import fs from 'fs'
+import { RecordWriteOp } from '../src'
+
+type IdMapping = Record<string, CID>
+
+const fakeStore = new MemoryBlockstore()
+
+export const randomCid = async (store: IpldStore = fakeStore): Promise<CID> => {
+  const str = randomStr(50)
+  return store.stage({ test: str })
+}
+
+export const generateBulkTids = (count: number): TID[] => {
+  const ids: TID[] = []
+  for (let i = 0; i < count; i++) {
+    ids.push(TID.next())
+  }
+  return ids
+}
+
+export const generateBulkTidMapping = async (
+  count: number,
+  blockstore: IpldStore = fakeStore,
+): Promise<IdMapping> => {
+  const ids = generateBulkTids(count)
+  const obj: IdMapping = {}
+  for (const id of ids) {
+    obj[id.toString()] = await randomCid(blockstore)
+  }
+  return obj
+}
+
+export const keysFromMapping = (mapping: IdMapping): TID[] => {
+  return Object.keys(mapping).map((id) => TID.fromStr(id))
+}
+
+export const keysFromMappings = (mappings: IdMapping[]): TID[] => {
+  return mappings.map(keysFromMapping).flat()
+}
+
+export const randomStr = (len: number): string => {
+  let result = ''
+  const CHARS = '123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz'
+  for (let i = 0; i < len; i++) {
+    result += CHARS.charAt(Math.floor(Math.random() * CHARS.length))
+  }
+  return result
+}
+
+export const shuffle = <T>(arr: T[]): T[] => {
+  const toShuffle = [...arr]
+  const shuffled: T[] = []
+  while (toShuffle.length > 0) {
+    const index = Math.floor(Math.random() * toShuffle.length)
+    shuffled.push(toShuffle[index])
+    toShuffle.splice(index, 1)
+  }
+  return shuffled
+}
+
+export const generateObject = (): Record<string, string> => {
+  return {
+    name: randomStr(100),
+  }
+}
+
+// Mass repo mutations & checking
+// -------------------------------
+
+export const testCollections = ['com.example.posts', 'com.example.likes']
+
+export type CollectionData = Record<string, unknown>
+export type RepoData = Record<string, CollectionData>
+
+export const fillRepo = async (
+  repo: Repo,
+  authStore: auth.AuthStore,
+  itemsPerCollection: number,
+): Promise<{ repo: Repo; data: RepoData }> => {
+  const repoData: RepoData = {}
+  const writes: RecordWriteOp[] = []
+  for (const collName of testCollections) {
+    const collData: CollectionData = {}
+    for (let i = 0; i < itemsPerCollection; i++) {
+      const object = generateObject()
+      const rkey = TID.nextStr()
+      collData[rkey] = object
+      writes.push({
+        action: 'create',
+        collection: collName,
+        rkey,
+        value: object,
+      })
+    }
+    repoData[collName] = collData
+  }
+  const updated = await repo.stageUpdate(writes).createCommit(authStore)
+  return {
+    repo: updated,
+    data: repoData,
+  }
+}
+
+export const editRepo = async (
+  repo: Repo,
+  prevData: RepoData,
+  authStore: auth.AuthStore,
+  params: {
+    adds?: number
+    updates?: number
+    deletes?: number
+  },
+): Promise<{ repo: Repo; data: RepoData }> => {
+  const { adds = 0, updates = 0, deletes = 0 } = params
+  const repoData: RepoData = {}
+  const writes: RecordWriteOp[] = []
+  for (const collName of testCollections) {
+    const collData = prevData[collName]
+    const shuffled = shuffle(Object.entries(collData))
+
+    for (let i = 0; i < adds; i++) {
+      const object = generateObject()
+      const rkey = TID.nextStr()
+      collData[rkey] = object
+      writes.push({
+        action: 'create',
+        collection: collName,
+        rkey,
+        value: object,
+      })
+    }
+
+    const toUpdate = shuffled.slice(0, updates)
+    for (let i = 0; i < toUpdate.length; i++) {
+      const object = generateObject()
+      const rkey = toUpdate[i][0]
+      writes.push({
+        action: 'update',
+        collection: collName,
+        rkey,
+        value: object,
+      })
+      collData[rkey] = object
+    }
+
+    const toDelete = shuffled.slice(updates, deletes)
+    for (let i = 0; i < toDelete.length; i++) {
+      const rkey = toDelete[i][0]
+      writes.push({
+        action: 'delete',
+        collection: collName,
+        rkey,
+      })
+      delete collData[rkey]
+    }
+    repoData[collName] = collData
+  }
+  const updated = await repo.stageUpdate(writes).createCommit(authStore)
+  return {
+    repo: updated,
+    data: repoData,
+  }
+}
+
+export const checkRepo = async (repo: Repo, data: RepoData): Promise<void> => {
+  for (const collName of Object.keys(data)) {
+    const collData = data[collName]
+    for (const rkey of Object.keys(collData)) {
+      const record = await repo.getRecord(collName, rkey)
+      expect(record).toEqual(collData[rkey])
+    }
+  }
+}
+
+export const checkRepoDiff = async (
+  diff: DataDiff,
+  before: RepoData,
+  after: RepoData,
+): Promise<void> => {
+  const getObjectCid = async (
+    key: string,
+    data: RepoData,
+  ): Promise<CID | undefined> => {
+    const parts = key.split('/')
+    const collection = parts[0]
+    const obj = (data[collection] || {})[parts[1]]
+    return obj === undefined ? undefined : fakeStore.stage(obj)
+  }
+
+  for (const add of diff.addList()) {
+    const beforeCid = await getObjectCid(add.key, before)
+    const afterCid = await getObjectCid(add.key, after)
+
+    expect(beforeCid).toBeUndefined()
+    expect(afterCid).toEqual(add.cid)
+  }
+
+  for (const update of diff.updateList()) {
+    const beforeCid = await getObjectCid(update.key, before)
+    const afterCid = await getObjectCid(update.key, after)
+
+    expect(beforeCid).toEqual(update.prev)
+    expect(afterCid).toEqual(update.cid)
+  }
+
+  for (const del of diff.deleteList()) {
+    const beforeCid = await getObjectCid(del.key, before)
+    const afterCid = await getObjectCid(del.key, after)
+
+    expect(beforeCid).toEqual(del.cid)
+    expect(afterCid).toBeUndefined()
+  }
+}
+
+// Logging
+// ----------------
+
+export const writeMstLog = async (filename: string, tree: MST) => {
+  let log = ''
+  for await (const entry of tree.walk()) {
+    if (entry.isLeaf()) continue
+    const layer = await entry.getLayer()
+    log += `Layer ${layer}: ${entry.pointer}\n`
+    log += '--------------\n'
+    const entries = await entry.getEntries()
+    for (const e of entries) {
+      if (e.isLeaf()) {
+        log += `Key: ${e.key} (${e.value})\n`
+      } else {
+        log += `Subtree: ${e.pointer}\n`
+      }
+    }
+    log += '\n\n'
+  }
+  fs.writeFileSync(filename, log)
+}
+
+export const saveMstEntries = (filename: string, entries: [string, CID][]) => {
+  const writable = entries.map(([key, val]) => [key, val.toString()])
+  fs.writeFileSync(filename, JSON.stringify(writable))
+}
+
+export const loadMstEntries = (filename: string): [string, CID][] => {
+  const contents = fs.readFileSync(filename)
+  const parsed = JSON.parse(contents.toString())
+  return parsed.map(([key, value]) => [key, CID.parse(value)])
+}

package/tests/mst.test.ts

@@ -0,0 +1,280 @@
+import { MST, DataAdd, DataUpdate, DataDelete } from '../src/mst'
+import { countPrefixLen } from '../src/mst/util'
+
+import { MemoryBlockstore } from '../src/blockstore'
+import * as util from './_util'
+
+import { CID } from 'multiformats'
+
+describe('Merkle Search Tree', () => {
+  let blockstore: MemoryBlockstore
+  let mst: MST
+  let mapping: Record<string, CID>
+  let shuffled: [string, CID][]
+
+  beforeAll(async () => {
+    blockstore = new MemoryBlockstore()
+    mst = await MST.create(blockstore)
+    mapping = await util.generateBulkTidMapping(1000, blockstore)
+    shuffled = util.shuffle(Object.entries(mapping))
+  })
+
+  it('adds records', async () => {
+    for (const entry of shuffled) {
+      mst = await mst.add(entry[0], entry[1])
+    }
+    for (const entry of shuffled) {
+      const got = await mst.get(entry[0])
+      expect(entry[1].equals(got)).toBeTruthy()
+    }
+
+    const totalSize = await mst.leafCount()
+    expect(totalSize).toBe(1000)
+  })
+
+  it('edits records', async () => {
+    let editedMst = mst
+    const toEdit = shuffled.slice(0, 100)
+
+    const edited: [string, CID][] = []
+    for (const entry of toEdit) {
+      const newCid = await util.randomCid()
+      editedMst = await editedMst.update(entry[0], newCid)
+      edited.push([entry[0], newCid])
+    }
+
+    for (const entry of edited) {
+      const got = await editedMst.get(entry[0])
+      expect(entry[1].equals(got)).toBeTruthy()
+    }
+
+    const totalSize = await editedMst.leafCount()
+    expect(totalSize).toBe(1000)
+  })
+
+  it('deletes records', async () => {
+    let deletedMst = mst
+    const toDelete = shuffled.slice(0, 100)
+    const theRest = shuffled.slice(100)
+    for (const entry of toDelete) {
+      deletedMst = await deletedMst.delete(entry[0])
+    }
+
+    const totalSize = await deletedMst.leafCount()
+    expect(totalSize).toBe(900)
+
+    for (const entry of toDelete) {
+      const got = await deletedMst.get(entry[0])
+      expect(got).toBe(null)
+    }
+    for (const entry of theRest) {
+      const got = await deletedMst.get(entry[0])
+      expect(entry[1].equals(got)).toBeTruthy()
+    }
+  })
+
+  it('is order independent', async () => {
+    const allNodes = await mst.allNodes()
+
+    let recreated = await MST.create(blockstore)
+    const reshuffled = util.shuffle(Object.entries(mapping))
+    for (const entry of reshuffled) {
+      recreated = await recreated.add(entry[0], entry[1])
+    }
+    const allReshuffled = await recreated.allNodes()
+
+    expect(allNodes.length).toBe(allReshuffled.length)
+    for (let i = 0; i < allNodes.length; i++) {
+      expect(await allNodes[i].equals(allReshuffled[i])).toBeTruthy()
+    }
+  })
+
+  it('saves and loads from blockstore', async () => {
+    const cid = await mst.stage()
+    const loaded = await MST.load(blockstore, cid)
+    const origNodes = await mst.allNodes()
+    const loadedNodes = await loaded.allNodes()
+    expect(origNodes.length).toBe(loadedNodes.length)
+    for (let i = 0; i < origNodes.length; i++) {
+      expect(await origNodes[i].equals(loadedNodes[i])).toBeTruthy()
+    }
+  })
+
+  it('diffs', async () => {
+    let toDiff = mst
+
+    const toAdd = Object.entries(
+      await util.generateBulkTidMapping(100, blockstore),
+    )
+    const toEdit = shuffled.slice(500, 600)
+    const toDel = shuffled.slice(400, 500)
+
+    const expectedAdds: Record<string, DataAdd> = {}
+    const expectedUpdates: Record<string, DataUpdate> = {}
+    const expectedDels: Record<string, DataDelete> = {}
+
+    for (const entry of toAdd) {
+      toDiff = await toDiff.add(entry[0], entry[1])
+      expectedAdds[entry[0]] = { key: entry[0], cid: entry[1] }
+    }
+    for (const entry of toEdit) {
+      const updated = await util.randomCid()
+      toDiff = await toDiff.update(entry[0], updated)
+      expectedUpdates[entry[0]] = {
+        key: entry[0],
+        prev: entry[1],
+        cid: updated,
+      }
+    }
+    for (const entry of toDel) {
+      toDiff = await toDiff.delete(entry[0])
+      expectedDels[entry[0]] = { key: entry[0], cid: entry[1] }
+    }
+
+    const diff = await mst.diff(toDiff)
+
+    expect(diff.addList().length).toBe(100)
+    expect(diff.updateList().length).toBe(100)
+    expect(diff.deleteList().length).toBe(100)
+
+    expect(diff.adds).toEqual(expectedAdds)
+    expect(diff.updates).toEqual(expectedUpdates)
+    expect(diff.deletes).toEqual(expectedDels)
+
+    // ensure we correctly report all added CIDs
+    for await (const entry of toDiff.walk()) {
+      let cid: CID
+      if (entry.isTree()) {
+        cid = await entry.getPointer()
+      } else {
+        cid = entry.value
+      }
+      const found = (await blockstore.has(cid)) || diff.newCids.has(cid)
+      expect(found).toBeTruthy()
+    }
+  })
+
+  // Special Cases (these are made for fanout 32)
+  // ------------
+
+  it('trims the top of an MST on stage', async () => {
+    const layer0 = [
+      '3j6hnk65jis2t',
+      '3j6hnk65jit2t',
+      '3j6hnk65jiu2t',
+      '3j6hnk65jne2t',
+      '3j6hnk65jnm2t',
+    ]
+    const layer1 = '3j6hnk65jju2t'
+    mst = await MST.create(blockstore, [], { fanout: 32 })
+    const cid = await util.randomCid()
+    const tids = [...layer0, layer1]
+    for (const tid of tids) {
+      mst = await mst.add(tid, cid)
+    }
+    const layer = await mst.getLayer()
+    expect(layer).toBe(1)
+    mst = await mst.delete(layer1)
+    const root = await mst.stage()
+    const loaded = MST.load(blockstore, root)
+    const loadedLayer = await loaded.getLayer()
+    expect(loadedLayer).toBe(0)
+  })
+
+  // These are some tricky things that can come up that may not be included in a randomized tree
+
+  /**
+   *   `f` gets added & it does two node splits (e is no longer grouped with g/h)
+   *
+   *                *                                  *
+   *       _________|________                      ____|_____
+   *       |   |    |    |   |                    |    |     |
+   *       *   d    *    i   *       ->           *    f     *
+   *     __|__    __|__    __|__                __|__      __|___
+   *    |  |  |  |  |  |  |  |  |              |  |  |    |  |   |
+   *    a  b  c  e  g  h  j  k  l              *  d  *    *  i   *
+   *                                         __|__   |   _|_   __|__
+   *                                        |  |  |  |  |   | |  |  |
+   *                                        a  b  c  e  g   h j  k  l
+   *
+   */
+  it('handles splits that must go 2 deep', async () => {
+    const layer0 = [
+      '3j6hnk65jis2t',
+      '3j6hnk65jit2t',
+      '3j6hnk65jiu2t',
+      '3j6hnk65jne2t',
+      '3j6hnk65jnm2t',
+      '3j6hnk65jnn2t',
+      '3j6hnk65kvx2t',
+      '3j6hnk65kvy2t',
+      '3j6hnk65kvz2t',
+    ]
+    const layer1 = ['3j6hnk65jju2t', '3j6hnk65kve2t']
+    const layer2 = '3j6hnk65jng2t'
+    mst = await MST.create(blockstore, [], { fanout: 32 })
+    const cid = await util.randomCid()
+    for (const tid of layer0) {
+      mst = await mst.add(tid, cid)
+    }
+    for (const tid of layer1) {
+      mst = await mst.add(tid, cid)
+    }
+    mst = await mst.add(layer2, cid)
+    const layer = await mst.getLayer()
+    expect(layer).toBe(2)
+
+    const root = await mst.stage()
+    mst = MST.load(blockstore, root, { fanout: 32 })
+
+    const allTids = [...layer0, ...layer1, layer2]
+    for (const tid of allTids) {
+      const got = await mst.get(tid)
+      expect(cid.equals(got)).toBeTruthy()
+    }
+  })
+  /**
+   *   `b` gets added & it hashes to 2 levels above any existing leaves
+   *
+   *          *        ->            *
+   *        __|__                  __|__
+   *       |     |                |  |  |
+   *       a     c                *  b  *
+   *                              |     |
+   *                              *     *
+   *                              |     |
+   *                              a     c
+   *
+   */
+  it('handles new layers that are two higher than existing', async () => {
+    const layer0 = ['3j6hnk65jis2t', '3j6hnk65kvz2t']
+    const layer2 = '3j6hnk65jng2t'
+    mst = await MST.create(blockstore, [], { fanout: 32 })
+    const cid = await util.randomCid()
+    for (const tid of layer0) {
+      mst = await mst.add(tid, cid)
+    }
+    mst = await mst.add(layer2, cid)
+
+    const root = await mst.stage()
+    mst = MST.load(blockstore, root, { fanout: 32 })
+
+    const layer = await mst.getLayer()
+    expect(layer).toBe(2)
+    const allTids = [...layer0, layer2]
+    for (const tid of allTids) {
+      const got = await mst.get(tid)
+      expect(cid.equals(got)).toBeTruthy()
+    }
+  })
+})
+
+describe('utils', () => {
+  it('counts prefix length', () => {
+    expect(countPrefixLen('abc', 'abc')).toBe(3)
+    expect(countPrefixLen('', 'abc')).toBe(0)
+    expect(countPrefixLen('abc', '')).toBe(0)
+    expect(countPrefixLen('ab', 'abc')).toBe(2)
+    expect(countPrefixLen('abc', 'ab')).toBe(2)
+  })
+})