chadtmiller.com/pds.js

Switch repository

Loading…

// dev.pdsjs.repo.rebuildMst rebuilds the tree from the records under it.
// buildMst over listAllRecords is the oracle: the rebuilt root must equal
// what a full rebuild produces, or the repair has not repaired anything.

import { PersonalDataServer } from '@pdsjs/core';
import { buildMst } from '@pdsjs/core/mst';
import {
  cborDecode,
  cborEncodeDagCbor,
  cidToString,
  createCid,
  parseCommitBlock,
} from '@pdsjs/core/repo';
import { beforeEach, describe, expect, test } from 'vitest';
import {
  createMemorySharedStorage,
  createMemoryStorage,
} from './helpers/memory-storage.js';

const DID = 'did:plc:rebuildtestaaaaaaaaaaaaa';
const PASSWORD = 'test-password';

/** Root CID a full rebuild would produce for everything in storage. */
async function oracleRoot(/** @type {any} */ storage) {
  return buildMst(await storage.listAllRecords(), async () => {});
}

/** The head commit, failing the test rather than the typechecker if absent. */
async function head(/** @type {any} */ storage) {
  const commit = await storage.getLatestCommit();
  expect(commit, 'repo has no commits').toBeTruthy();
  return /** @type {{seq: number, cid: string, rev: string}} */ (commit);
}

/** The MST root the repo's head commit points at. */
async function headRoot(/** @type {any} */ storage) {
  const head = await storage.getLatestCommit();
  if (!head) return null;
  const bytes = await storage.getBlock(head.cid);
  return bytes ? parseCommitBlock(bytes).data : null;
}

describe('rebuildMst', () => {
  /** @type {ReturnType<typeof createMemoryStorage>} */
  let storage;
  /** @type {PersonalDataServer} */
  let pds;
  /** @type {string} */
  let token;

  beforeEach(async () => {
    storage = createMemoryStorage();
    pds = new PersonalDataServer({
      actorStorage: /** @type {any} */ (storage),
      sharedStorage: /** @type {any} */ (createMemorySharedStorage()),
      blobs: {
        async get() {
          return null;
        },
        async put() {},
        async delete() {},
      },
      jwtSecret: 'test-secret',
      hostname: 'pds.example.com',
      password: PASSWORD,
    });

    await pds.fetch(
      new Request(`https://pds.example.com/init?did=${DID}`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
          did: DID,
          privateKey: Buffer.from(
            crypto.getRandomValues(new Uint8Array(32)),
          ).toString('hex'),
          handle: 'rebuild.test',
          password: PASSWORD,
        }),
      }),
    );

    const sessionRes = await pds.fetch(
      new Request(
        'https://pds.example.com/xrpc/com.atproto.server.createSession',
        {
          method: 'POST',
          headers: { 'Content-Type': 'application/json' },
          body: JSON.stringify({ identifier: DID, password: PASSWORD }),
        },
      ),
    );
    const sessionBody = await sessionRes.json();
    expect(sessionBody.accessJwt, 'session setup failed').toBeTruthy();
    token = sessionBody.accessJwt;
  });

  /** @param {number} i */
  async function createRecord(i) {
    const res = await pds.fetch(
      new Request(
        'https://pds.example.com/xrpc/com.atproto.repo.createRecord',
        {
          method: 'POST',
          headers: {
            'Content-Type': 'application/json',
            Authorization: `Bearer ${token}`,
          },
          body: JSON.stringify({
            repo: DID,
            collection: 'app.bsky.feed.post',
            record: { text: `post ${i}`, createdAt: '2026-01-01T00:00:00Z' },
          }),
        },
      ),
    );
    expect(res.status).toBe(200);
  }

  /** @param {string|null} [auth] Pass null to send no Authorization header. */
  function rebuild(auth = token) {
    return pds.fetch(
      new Request('https://pds.example.com/xrpc/dev.pdsjs.repo.rebuildMst', {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          ...(auth ? { Authorization: `Bearer ${auth}` } : {}),
        },
        body: '{}',
      }),
    );
  }

  /**
   * Put a record into storage without committing, so the tree no longer
   * covers every record. Stands in for any drift between the two: a
   * duplicate key, a node lost to a partial write.
   */
  async function orphanRecord(/** @type {string} */ rkey) {
    const uri = `at://${DID}/app.bsky.feed.post/${rkey}`;
    const encoded = cborEncodeDagCbor({
      $type: 'app.bsky.feed.post',
      text: 'orphan',
      createdAt: '2026-01-01T00:00:00Z',
    });
    const cid = cidToString(await createCid(encoded));
    await storage.putBlock(cid, encoded);
    await storage.putRecord(uri, cid, 'app.bsky.feed.post', rkey, encoded);
    return cid;
  }

  test('rebuilds a tree that has drifted from its records', async () => {
    for (let i = 0; i < 30; i++) await createRecord(i);
    await orphanRecord('3korphanaaaaa');

    // The head does not cover the orphan, so it disagrees with the records.
    expect(await headRoot(storage)).not.toBe(await oracleRoot(storage));

    const res = await rebuild();
    expect(res.status).toBe(200);
    const body = await res.json();
    expect(body.changed).toBe(true);

    expect(await headRoot(storage)).toBe(await oracleRoot(storage));
    const latest = await head(storage);
    expect(latest.rev).toBe(body.rev);
    expect(latest.cid).toBe(body.cid);
  });

  test('stores every node of the rebuilt tree', async () => {
    for (let i = 0; i < 40; i++) await createRecord(i);
    await orphanRecord('3korphanbbbbb');
    await rebuild();

    // A root nothing can be read beneath is no better than the broken tree.
    const seen = new Set();
    /** @param {string|null} cid */
    const walk = async (cid) => {
      if (!cid || seen.has(cid)) return;
      seen.add(cid);
      const data = await storage.getBlock(cid);
      expect(data, `node ${cid} missing from storage`).toBeTruthy();
      const node = cborDecode(/** @type {Uint8Array} */ (data));
      if (node.l) await walk(cidToString(node.l));
      for (const entry of node.e || []) {
        if (entry.t) await walk(cidToString(entry.t));
      }
    };
    await walk(await headRoot(storage));
    expect(seen.size).toBeGreaterThan(1);
  });

  test('announces the new root with #sync, carrying no ops', async () => {
    for (let i = 0; i < 10; i++) await createRecord(i);
    await orphanRecord('3korphanccccc');

    const before = storage._events.length;
    await rebuild();

    const emitted = storage._events.slice(before);
    expect(emitted.length).toBe(1);
    const evt = cborDecode(emitted[0].evt);
    expect(evt.type).toBe('sync');
    // A commit would have to explain the move with ops it does not have.
    expect(evt.ops).toBeUndefined();
    expect(evt.rev).toBe((await head(storage)).rev);
  });

  test('is a no-op when the tree already matches the records', async () => {
    for (let i = 0; i < 10; i++) await createRecord(i);

    const before = await head(storage);
    const events = storage._events.length;

    const res = await rebuild();
    expect(res.status).toBe(200);
    const body = await res.json();
    expect(body.changed).toBe(false);
    expect(body.rev).toBe(before.rev);

    // Nothing announced, and the head did not move.
    expect(storage._events.length).toBe(events);
    expect((await head(storage)).cid).toBe(before.cid);
  });

  test('rebuilding twice is stable', async () => {
    for (let i = 0; i < 20; i++) await createRecord(i);
    await orphanRecord('3korphanddddd');

    const first = await (await rebuild()).json();
    expect(first.changed).toBe(true);
    const second = await (await rebuild()).json();
    expect(second.changed).toBe(false);
    expect(second.cid).toBe(first.cid);
  });

  test('refuses a caller that is not the account', async () => {
    for (let i = 0; i < 5; i++) await createRecord(i);
    const before = await head(storage);

    expect((await rebuild(null)).status).toBe(401);
    expect((await rebuild('not-a-token')).status).toBe(401);

    // Refused means nothing happened, not merely nothing returned.
    expect((await head(storage)).cid).toBe(before.cid);
  });
});